← Back to C-Kernel-Engine Docs Doxygen Source Documentation
 
Loading...
Searching...
No Matches
Functions
add_kernels_bf16.c File Reference

Element-wise addition kernels for BF16 tensors. More...

#include "bf16_utils.h"
#include "ckernel_engine.h"
#include <stdint.h>
#include <stddef.h>

Go to the source code of this file.

Functions

void add_backward_bf16 (const uint16_t *d_y, uint16_t *d_a, uint16_t *d_b, size_t n)
 
void add_forward_2d_bf16 (const uint16_t *a, const uint16_t *b, uint16_t *y, int tokens, int dim, int aligned_dim)
 
void add_forward_bf16 (const uint16_t *a, const uint16_t *b, uint16_t *y, size_t n)
 
void add_forward_f32 (const float *a, const float *b, float *y, size_t n)
 
void add_inplace_bf16 (uint16_t *a, const uint16_t *b, size_t n)
 
void add_inplace_f32 (float *a, const float *b, size_t n)
 
void add_scaled_forward_bf16 (const uint16_t *a, const uint16_t *b, uint16_t *y, float alpha, size_t n)
 
void add_scaled_inplace_bf16 (uint16_t *a, const uint16_t *b, float alpha, size_t n)
 
void ck_residual_add_token_major_bf16_storage (const float *a, const float *b, float *out, int tokens, int aligned_embed_dim)
 

Detailed Description

Element-wise addition kernels for BF16 tensors.

CK-ENGINE KERNEL RULES:

  1. NO malloc/free - memory via bump allocator, pointers passed in
  2. NO OpenMP - parallelization at orchestrator/codegen layer
  3. API must define: inputs, outputs, workspace, and memory layouts
  4. Pure computation - deterministic, no side effects

After changes: make test && make llamacpp-parity-full

Used for residual connections in transformer models: residual = x + sublayer_output

Supports:

  • Forward: y = a + b
  • Forward with scale: y = a + alpha * b
  • Backward: d_a = d_y, d_b = d_y (gradient flows through unchanged)
  • In-place: a += b

Definition in file add_kernels_bf16.c.

Function Documentation

◆ add_backward_bf16()

void add_backward_bf16 ( const uint16_t *  d_y,
uint16_t *  d_a,
uint16_t *  d_b,
size_t  n 
)

Definition at line 187 of file add_kernels_bf16.c.

191{
192 if (!d_y || n == 0) {
193 return;
194 }
195
196 size_t i = 0;
197
198 /* Copy to d_a if not in-place */
199 if (d_a && d_a != d_y) {
200#if defined(__AVX512F__)
201 for (; i + 32 <= n; i += 32) {
202 __m512i v0 = _mm512_loadu_si512((const __m512i*)&d_y[i]);
203 __m512i v1 = _mm512_loadu_si512((const __m512i*)&d_y[i + 32]);
204 _mm512_storeu_si512((__m512i*)&d_a[i], v0);
205 _mm512_storeu_si512((__m512i*)&d_a[i + 32], v1);
206 }
207#endif
208 for (; i < n; ++i) {
209 d_a[i] = d_y[i];
210 }
211 }
212
213 /* Copy to d_b if not in-place */
214 i = 0;
215 if (d_b && d_b != d_y) {
216#if defined(__AVX512F__)
217 for (; i + 32 <= n; i += 32) {
218 __m512i v0 = _mm512_loadu_si512((const __m512i*)&d_y[i]);
219 __m512i v1 = _mm512_loadu_si512((const __m512i*)&d_y[i + 32]);
220 _mm512_storeu_si512((__m512i*)&d_b[i], v0);
221 _mm512_storeu_si512((__m512i*)&d_b[i + 32], v1);
222 }
223#endif
224 for (; i < n; ++i) {
225 d_b[i] = d_y[i];
226 }
227 }
228}

◆ add_forward_2d_bf16()

void add_forward_2d_bf16 ( const uint16_t *  a,
const uint16_t *  b,
uint16_t *  y,
int  tokens,
int  dim,
int  aligned_dim 
)

Definition at line 235 of file add_kernels_bf16.c.

241{
242 if (!a || !b || !y || tokens <= 0 || dim <= 0) {
243 return;
244 }
245
246 for (int t = 0; t < tokens; ++t) {
247 const uint16_t *a_row = a + (size_t)t * aligned_dim;
248 const uint16_t *b_row = b + (size_t)t * aligned_dim;
249 uint16_t *y_row = y + (size_t)t * aligned_dim;
250
251 int d = 0;
252
253#if defined(__AVX512F__)
254 for (; d + 16 <= dim; d += 16) {
255 __m512 av = bf16_loadu_cvt_fp32(&a_row[d]);
256 __m512 bv = bf16_loadu_cvt_fp32(&b_row[d]);
257 __m512 yv = _mm512_add_ps(av, bv);
258 fp32_cvt_storeu_bf16(&y_row[d], yv);
259 }
260#endif
261
262 for (; d < dim; ++d) {
263 float af = bf16_to_float(a_row[d]);
264 float bf = bf16_to_float(b_row[d]);
265 y_row[d] = float_to_bf16(af + bf);
266 }
267 }
268}
float_to_bf16
static uint16_t float_to_bf16(float f)
Definition bf16_utils.h:90
bf16_to_float
static float bf16_to_float(uint16_t v)
Definition bf16_utils.h:38

References bf16_to_float(), and float_to_bf16().

◆ add_forward_bf16()

void add_forward_bf16 ( const uint16_t *  a,
const uint16_t *  b,
uint16_t *  y,
size_t  n 
)

Definition at line 38 of file add_kernels_bf16.c.

42{
43 if (!a || !b || !y || n == 0) {
44 return;
45 }
46
47 size_t i = 0;
48
49#if defined(__AVX512F__)
50 /* AVX-512: Process 16 bf16 elements at a time */
51 for (; i + 16 <= n; i += 16) {
52 __m512 av = bf16_loadu_cvt_fp32(&a[i]);
53 __m512 bv = bf16_loadu_cvt_fp32(&b[i]);
54 __m512 yv = _mm512_add_ps(av, bv);
55 fp32_cvt_storeu_bf16(&y[i], yv);
56 }
57#endif
58
59 /* Scalar fallback */
60 for (; i < n; ++i) {
61 float af = bf16_to_float(a[i]);
62 float bf = bf16_to_float(b[i]);
63 y[i] = float_to_bf16(af + bf);
64 }
65}

References bf16_to_float(), and float_to_bf16().

◆ add_forward_f32()

void add_forward_f32 ( const float *  a,
const float *  b,
float *  y,
size_t  n 
)

Element-wise add: y = a + b

Test:

test_add.py::TestAddForward::test_add_forward_f32

test_add.py::TestAddForward::test_add_inplace_f32

test_multi_layer_parity.py::TestMultiLayerParity::test_residual_add

Element-wise addition of two vectors.

After changes: make test

Definition at line 284 of file add_kernels_bf16.c.

288{
289 if (!a || !b || !y || n == 0) {
290 return;
291 }
292
293 size_t i = 0;
294
295#if defined(__AVX512F__)
296 for (; i + 16 <= n; i += 16) {
297 __m512 av = _mm512_loadu_ps(&a[i]);
298 __m512 bv = _mm512_loadu_ps(&b[i]);
299 __m512 yv = _mm512_add_ps(av, bv);
300 _mm512_storeu_ps(&y[i], yv);
301 }
302#endif
303
304#if defined(__AVX2__)
305 for (; i + 8 <= n; i += 8) {
306 __m256 av = _mm256_loadu_ps(&a[i]);
307 __m256 bv = _mm256_loadu_ps(&b[i]);
308 __m256 yv = _mm256_add_ps(av, bv);
309 _mm256_storeu_ps(&y[i], yv);
310 }
311#endif
312
313 for (; i < n; ++i) {
314 y[i] = a[i] + b[i];
315 }
316}

◆ add_inplace_bf16()

void add_inplace_bf16 ( uint16_t *  a,
const uint16_t *  b,
size_t  n 
)

Definition at line 119 of file add_kernels_bf16.c.

122{
123 if (!a || !b || n == 0) {
124 return;
125 }
126
127 size_t i = 0;
128
129#if defined(__AVX512F__)
130 for (; i + 16 <= n; i += 16) {
131 __m512 av = bf16_loadu_cvt_fp32(&a[i]);
132 __m512 bv = bf16_loadu_cvt_fp32(&b[i]);
133 __m512 yv = _mm512_add_ps(av, bv);
134 fp32_cvt_storeu_bf16(&a[i], yv);
135 }
136#endif
137
138 for (; i < n; ++i) {
139 float af = bf16_to_float(a[i]);
140 float bf = bf16_to_float(b[i]);
141 a[i] = float_to_bf16(af + bf);
142 }
143}

References bf16_to_float(), and float_to_bf16().

◆ add_inplace_f32()

void add_inplace_f32 ( float *  a,
const float *  b,
size_t  n 
)

Definition at line 318 of file add_kernels_bf16.c.

321{
322 if (!a || !b || n == 0) {
323 return;
324 }
325
326 size_t i = 0;
327
328#if defined(__AVX512F__)
329 for (; i + 16 <= n; i += 16) {
330 __m512 av = _mm512_loadu_ps(&a[i]);
331 __m512 bv = _mm512_loadu_ps(&b[i]);
332 __m512 yv = _mm512_add_ps(av, bv);
333 _mm512_storeu_ps(&a[i], yv);
334 }
335#endif
336
337#if defined(__AVX2__)
338 for (; i + 8 <= n; i += 8) {
339 __m256 av = _mm256_loadu_ps(&a[i]);
340 __m256 bv = _mm256_loadu_ps(&b[i]);
341 __m256 yv = _mm256_add_ps(av, bv);
342 _mm256_storeu_ps(&a[i], yv);
343 }
344#endif
345
346 for (; i < n; ++i) {
347 a[i] = a[i] + b[i];
348 }
349}

Referenced by mega_fused_outproj_mlp_prefill().

◆ add_scaled_forward_bf16()

void add_scaled_forward_bf16 ( const uint16_t *  a,
const uint16_t *  b,
uint16_t *  y,
float  alpha,
size_t  n 
)

Definition at line 86 of file add_kernels_bf16.c.

91{
92 if (!a || !b || !y || n == 0) {
93 return;
94 }
95
96 size_t i = 0;
97
98#if defined(__AVX512F__)
99 __m512 alpha_v = _mm512_set1_ps(alpha);
100 for (; i + 16 <= n; i += 16) {
101 __m512 av = bf16_loadu_cvt_fp32(&a[i]);
102 __m512 bv = bf16_loadu_cvt_fp32(&b[i]);
103 __m512 yv = _mm512_fmadd_ps(bv, alpha_v, av); /* a + alpha * b */
104 fp32_cvt_storeu_bf16(&y[i], yv);
105 }
106#endif
107
108 for (; i < n; ++i) {
109 float af = bf16_to_float(a[i]);
110 float bf = bf16_to_float(b[i]);
111 y[i] = float_to_bf16(af + alpha * bf);
112 }
113}

References bf16_to_float(), and float_to_bf16().

◆ add_scaled_inplace_bf16()

void add_scaled_inplace_bf16 ( uint16_t *  a,
const uint16_t *  b,
float  alpha,
size_t  n 
)

Definition at line 149 of file add_kernels_bf16.c.

153{
154 if (!a || !b || n == 0) {
155 return;
156 }
157
158 size_t i = 0;
159
160#if defined(__AVX512F__)
161 __m512 alpha_v = _mm512_set1_ps(alpha);
162 for (; i + 16 <= n; i += 16) {
163 __m512 av = bf16_loadu_cvt_fp32(&a[i]);
164 __m512 bv = bf16_loadu_cvt_fp32(&b[i]);
165 __m512 yv = _mm512_fmadd_ps(bv, alpha_v, av);
166 fp32_cvt_storeu_bf16(&a[i], yv);
167 }
168#endif
169
170 for (; i < n; ++i) {
171 float af = bf16_to_float(a[i]);
172 float bf = bf16_to_float(b[i]);
173 a[i] = float_to_bf16(af + alpha * bf);
174 }
175}

References bf16_to_float(), and float_to_bf16().

◆ ck_residual_add_token_major_bf16_storage()

void ck_residual_add_token_major_bf16_storage ( const float *  a,
const float *  b,
float *  out,
int  tokens,
int  aligned_embed_dim 
)

Definition at line 67 of file add_kernels_bf16.c.

72{
73 const size_t count = (size_t)tokens * (size_t)aligned_embed_dim;
74 for (size_t i = 0; i < count; ++i) {
75 const float av = bf16_to_float(float_to_bf16(a[i]));
76 const float bv = bf16_to_float(float_to_bf16(b[i]));
77 out[i] = bf16_to_float(float_to_bf16(av + bv));
78 }
79}

References bf16_to_float(), and float_to_bf16().