relu_kernels.c File Reference

ReLU activation kernels with SIMD (SSE/AVX/AVX512) More...

#include <math.h>
#include <stddef.h>
#include <stdint.h>

Go to the source code of this file.

Functions
void	relu2_backward (const float *input, const float *d_output, float *d_input, size_t n)
void	relu2_forward (const float *input, float *output, size_t n)
void	relu_backward (const float *input, const float *d_output, float *d_input, size_t n)
void	relu_forward (const float *input, float *output, size_t n)
void	relu_forward_inplace (float *data, size_t n)

Detailed Description

ReLU activation kernels with SIMD (SSE/AVX/AVX512)

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

ReLU: y = max(0, x)

Definition in file relu_kernels.c.

Function Documentation

relu2_backward()

void relu2_backward	(	const float *	input,
		const float *	d_output,
		float *	d_input,
		size_t	n
	)

Definition at line 147 of file relu_kernels.c.

{
    size_t i = 0;
 
#if defined(__AVX512F__)
    const __m512 vzero = _mm512_setzero_ps();
    const __m512 vtwo = _mm512_set1_ps(2.0f);
    for (; i + 15 < n; i += 16) {
        const __m512 vx = _mm512_loadu_ps(input + i);
        const __m512 vdy = _mm512_loadu_ps(d_output + i);
        const __mmask16 mask = _mm512_cmp_ps_mask(vx, vzero, _CMP_GT_OQ);
        const __m512 vdx = _mm512_mul_ps(_mm512_mul_ps(vtwo, vx), vdy);
        _mm512_storeu_ps(d_input + i, _mm512_maskz_mov_ps(mask, vdx));
    }
#elif defined(__AVX2__) || defined(__AVX__)
    const __m256 vzero = _mm256_setzero_ps();
    const __m256 vtwo = _mm256_set1_ps(2.0f);
    for (; i + 7 < n; i += 8) {
        const __m256 vx = _mm256_loadu_ps(input + i);
        const __m256 vdy = _mm256_loadu_ps(d_output + i);
        const __m256 mask = _mm256_cmp_ps(vx, vzero, _CMP_GT_OQ);
        const __m256 vdx = _mm256_mul_ps(_mm256_mul_ps(vtwo, vx), vdy);
        _mm256_storeu_ps(d_input + i, _mm256_and_ps(mask, vdx));
    }
#endif
 
    for (; i < n; ++i) {
        const float x = input[i];
        d_input[i] = (x > 0.0f) ? d_output[i] * 2.0f * x : 0.0f;
    }
}

References mask.

relu2_forward()

void relu2_forward	(	const float *	input,
		float *	output,
		size_t	n
	)

Definition at line 120 of file relu_kernels.c.

{
    size_t i = 0;
 
#if defined(__AVX512F__)
    const __m512 vzero = _mm512_setzero_ps();
    for (; i + 15 < n; i += 16) {
        const __m512 vx = _mm512_loadu_ps(input + i);
        const __m512 vr = _mm512_max_ps(vx, vzero);
        _mm512_storeu_ps(output + i, _mm512_mul_ps(vr, vr));
    }
#elif defined(__AVX2__) || defined(__AVX__)
    const __m256 vzero = _mm256_setzero_ps();
    for (; i + 7 < n; i += 8) {
        const __m256 vx = _mm256_loadu_ps(input + i);
        const __m256 vr = _mm256_max_ps(vx, vzero);
        _mm256_storeu_ps(output + i, _mm256_mul_ps(vr, vr));
    }
#endif
 
    for (; i < n; ++i) {
        const float x = input[i];
        output[i] = (x > 0.0f) ? x * x : 0.0f;
    }
}

relu_backward()

void relu_backward	(	const float *	input,
		const float *	d_output,
		float *	d_input,
		size_t	n
	)

Definition at line 84 of file relu_kernels.c.

{
    size_t i = 0;
 
#if defined(__AVX512F__)
    __m512 vzero = _mm512_setzero_ps();
    for (; i + 15 < n; i += 16) {
        __m512 vx = _mm512_loadu_ps(input + i);
        __m512 vdy = _mm512_loadu_ps(d_output + i);
        __mmask16 mask = _mm512_cmp_ps_mask(vx, vzero, _CMP_GT_OQ);
        __m512 vdx = _mm512_maskz_mov_ps(mask, vdy);
        _mm512_storeu_ps(d_input + i, vdx);
    }
#elif defined(__AVX2__) || defined(__AVX__)
    __m256 vzero = _mm256_setzero_ps();
    for (; i + 7 < n; i += 8) {
        __m256 vx = _mm256_loadu_ps(input + i);
        __m256 vdy = _mm256_loadu_ps(d_output + i);
        // Result is all 1s (0xFFFFFFFF) if true, 0 if false.
        __m256 mask = _mm256_cmp_ps(vx, vzero, _CMP_GT_OQ);
        __m256 vdx = _mm256_and_ps(mask, vdy);
        _mm256_storeu_ps(d_input + i, vdx);
    }
#endif
 
    // Scalar fallback
    for (; i < n; ++i) {
        d_input[i] = (input[i] > 0.0f) ? d_output[i] : 0.0f;
    }
}

References mask.

relu_forward()

void relu_forward	(	const float *	input,
		float *	output,
		size_t	n
	)

Definition at line 26 of file relu_kernels.c.

{
    size_t i = 0;
 
#if defined(__AVX512F__)
    __m512 vzero = _mm512_setzero_ps();
    for (; i + 15 < n; i += 16) {
        __m512 vx = _mm512_loadu_ps(input + i);
        __m512 vy = _mm512_max_ps(vx, vzero);
        _mm512_storeu_ps(output + i, vy);
    }
#elif defined(__AVX2__) || defined(__AVX__)
    __m256 vzero = _mm256_setzero_ps();
    for (; i + 7 < n; i += 8) {
        __m256 vx = _mm256_loadu_ps(input + i);
        __m256 vy = _mm256_max_ps(vx, vzero);
        _mm256_storeu_ps(output + i, vy);
    }
#endif
 
    // Scalar fallback
    for (; i < n; ++i) {
        float x = input[i];
        output[i] = (x > 0.0f) ? x : 0.0f;
    }
}

relu_forward_inplace()

void relu_forward_inplace	(	float *	data,
		size_t	n
	)

Definition at line 54 of file relu_kernels.c.

{
    size_t i = 0;
 
#if defined(__AVX512F__)
    __m512 vzero = _mm512_setzero_ps();
    for (; i + 15 < n; i += 16) {
        __m512 vx = _mm512_loadu_ps(data + i);
        __m512 vy = _mm512_max_ps(vx, vzero);
        _mm512_storeu_ps(data + i, vy);
    }
#elif defined(__AVX2__) || defined(__AVX__)
    __m256 vzero = _mm256_setzero_ps();
    for (; i + 7 < n; i += 8) {
        __m256 vx = _mm256_loadu_ps(data + i);
        __m256 vy = _mm256_max_ps(vx, vzero);
        _mm256_storeu_ps(data + i, vy);
    }
#endif
 
    // Scalar fallback
    for (; i < n; ++i) {
        float x = data[i];
        if (x < 0.0f) {
            data[i] = 0.0f;
        }
    }
}