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.

{
    if (!d_y || n == 0) {
        return;
    }
 
    size_t i = 0;
 
    /* Copy to d_a if not in-place */
    if (d_a && d_a != d_y) {
#if defined(__AVX512F__)
        for (; i + 32 <= n; i += 32) {
            __m512i values = _mm512_loadu_si512((const __m512i*)&d_y[i]);
            _mm512_storeu_si512((__m512i*)&d_a[i], values);
        }
#endif
        for (; i < n; ++i) {
            d_a[i] = d_y[i];
        }
    }
 
    /* Copy to d_b if not in-place */
    i = 0;
    if (d_b && d_b != d_y) {
#if defined(__AVX512F__)
        for (; i + 32 <= n; i += 32) {
            __m512i values = _mm512_loadu_si512((const __m512i*)&d_y[i]);
            _mm512_storeu_si512((__m512i*)&d_b[i], values);
        }
#endif
        for (; i < n; ++i) {
            d_b[i] = d_y[i];
        }
    }
}

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 231 of file add_kernels_bf16.c.

{
    if (!a || !b || !y || tokens <= 0 || dim <= 0) {
        return;
    }
 
    for (int t = 0; t < tokens; ++t) {
        const uint16_t *a_row = a + (size_t)t * aligned_dim;
        const uint16_t *b_row = b + (size_t)t * aligned_dim;
        uint16_t *y_row = y + (size_t)t * aligned_dim;
 
        int d = 0;
 
#if defined(__AVX512F__)
        for (; d + 16 <= dim; d += 16) {
            __m512 av = bf16_loadu_cvt_fp32(&a_row[d]);
            __m512 bv = bf16_loadu_cvt_fp32(&b_row[d]);
            __m512 yv = _mm512_add_ps(av, bv);
            fp32_cvt_storeu_bf16(&y_row[d], yv);
        }
#endif
 
        for (; d < dim; ++d) {
            float af = bf16_to_float(a_row[d]);
            float bf = bf16_to_float(b_row[d]);
            y_row[d] = float_to_bf16(af + bf);
        }
    }
}

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.

{
    if (!a || !b || !y || n == 0) {
        return;
    }
 
    size_t i = 0;
 
#if defined(__AVX512F__)
    /* AVX-512: Process 16 bf16 elements at a time */
    for (; i + 16 <= n; i += 16) {
        __m512 av = bf16_loadu_cvt_fp32(&a[i]);
        __m512 bv = bf16_loadu_cvt_fp32(&b[i]);
        __m512 yv = _mm512_add_ps(av, bv);
        fp32_cvt_storeu_bf16(&y[i], yv);
    }
#endif
 
    /* Scalar fallback */
    for (; i < n; ++i) {
        float af = bf16_to_float(a[i]);
        float bf = bf16_to_float(b[i]);
        y[i] = float_to_bf16(af + bf);
    }
}

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 280 of file add_kernels_bf16.c.

{
    if (!a || !b || !y || n == 0) {
        return;
    }
 
    size_t i = 0;
 
#if defined(__AVX512F__)
    for (; i + 16 <= n; i += 16) {
        __m512 av = _mm512_loadu_ps(&a[i]);
        __m512 bv = _mm512_loadu_ps(&b[i]);
        __m512 yv = _mm512_add_ps(av, bv);
        _mm512_storeu_ps(&y[i], yv);
    }
#endif
 
#if defined(__AVX2__)
    for (; i + 8 <= n; i += 8) {
        __m256 av = _mm256_loadu_ps(&a[i]);
        __m256 bv = _mm256_loadu_ps(&b[i]);
        __m256 yv = _mm256_add_ps(av, bv);
        _mm256_storeu_ps(&y[i], yv);
    }
#endif
 
    for (; i < n; ++i) {
        y[i] = a[i] + b[i];
    }
}

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.

{
    if (!a || !b || n == 0) {
        return;
    }
 
    size_t i = 0;
 
#if defined(__AVX512F__)
    for (; i + 16 <= n; i += 16) {
        __m512 av = bf16_loadu_cvt_fp32(&a[i]);
        __m512 bv = bf16_loadu_cvt_fp32(&b[i]);
        __m512 yv = _mm512_add_ps(av, bv);
        fp32_cvt_storeu_bf16(&a[i], yv);
    }
#endif
 
    for (; i < n; ++i) {
        float af = bf16_to_float(a[i]);
        float bf = bf16_to_float(b[i]);
        a[i] = float_to_bf16(af + bf);
    }
}

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 314 of file add_kernels_bf16.c.

{
    if (!a || !b || n == 0) {
        return;
    }
 
    size_t i = 0;
 
#if defined(__AVX512F__)
    for (; i + 16 <= n; i += 16) {
        __m512 av = _mm512_loadu_ps(&a[i]);
        __m512 bv = _mm512_loadu_ps(&b[i]);
        __m512 yv = _mm512_add_ps(av, bv);
        _mm512_storeu_ps(&a[i], yv);
    }
#endif
 
#if defined(__AVX2__)
    for (; i + 8 <= n; i += 8) {
        __m256 av = _mm256_loadu_ps(&a[i]);
        __m256 bv = _mm256_loadu_ps(&b[i]);
        __m256 yv = _mm256_add_ps(av, bv);
        _mm256_storeu_ps(&a[i], yv);
    }
#endif
 
    for (; i < n; ++i) {
        a[i] = a[i] + b[i];
    }
}

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.

{
    if (!a || !b || !y || n == 0) {
        return;
    }
 
    size_t i = 0;
 
#if defined(__AVX512F__)
    __m512 alpha_v = _mm512_set1_ps(alpha);
    for (; i + 16 <= n; i += 16) {
        __m512 av = bf16_loadu_cvt_fp32(&a[i]);
        __m512 bv = bf16_loadu_cvt_fp32(&b[i]);
        __m512 yv = _mm512_fmadd_ps(bv, alpha_v, av);  /* a + alpha * b */
        fp32_cvt_storeu_bf16(&y[i], yv);
    }
#endif
 
    for (; i < n; ++i) {
        float af = bf16_to_float(a[i]);
        float bf = bf16_to_float(b[i]);
        y[i] = float_to_bf16(af + alpha * bf);
    }
}

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.

{
    if (!a || !b || n == 0) {
        return;
    }
 
    size_t i = 0;
 
#if defined(__AVX512F__)
    __m512 alpha_v = _mm512_set1_ps(alpha);
    for (; i + 16 <= n; i += 16) {
        __m512 av = bf16_loadu_cvt_fp32(&a[i]);
        __m512 bv = bf16_loadu_cvt_fp32(&b[i]);
        __m512 yv = _mm512_fmadd_ps(bv, alpha_v, av);
        fp32_cvt_storeu_bf16(&a[i], yv);
    }
#endif
 
    for (; i < n; ++i) {
        float af = bf16_to_float(a[i]);
        float bf = bf16_to_float(b[i]);
        a[i] = float_to_bf16(af + alpha * bf);
    }
}

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.

{
    const size_t count = (size_t)tokens * (size_t)aligned_embed_dim;
    for (size_t i = 0; i < count; ++i) {
        const float av = bf16_to_float(float_to_bf16(a[i]));
        const float bv = bf16_to_float(float_to_bf16(b[i]));
        out[i] = bf16_to_float(float_to_bf16(av + bv));
    }
}

References bf16_to_float(), count, and float_to_bf16().