gelu_kernels.c

Go to the documentation of this file.

/**
 * @file gelu_kernels.c
 * @brief GELU 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
 *
 * GELU: y = x * 0.5 * (1 + erf(x / sqrt(2)))
 * Fast approx: y = x * sigmoid(1.702 * x)
 */
 
#include <math.h>
#include <stddef.h>
#include <stdint.h>
#include <string.h>
 
#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
#include <immintrin.h>
#endif
 
#include "bf16_utils.h"
 
/* Fast vectorized exp approximation (same as softmax_kernels.c) */
#if defined(__AVX512F__)
static inline __m512 exp512_fast(__m512 x) {
    // Clamp to avoid overflow/underflow
    x = _mm512_max_ps(x, _mm512_set1_ps(-88.0f));
    x = _mm512_min_ps(x, _mm512_set1_ps(88.0f));
 
    const __m512 log2e = _mm512_set1_ps(1.4426950408889634f);
    const __m512 c1 = _mm512_set1_ps(0.693359375f);
    const __m512 c2 = _mm512_set1_ps(-2.12194440e-4f);
 
    __m512 t = _mm512_mul_ps(x, log2e);
    __m512 ti = _mm512_roundscale_ps(t, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
 
    __m512 rx = _mm512_sub_ps(x, _mm512_mul_ps(ti, c1));
    rx = _mm512_sub_ps(rx, _mm512_mul_ps(ti, c2));
 
    // Polynomial approximation
    const __m512 p0 = _mm512_set1_ps(1.0f);
    const __m512 p1 = _mm512_set1_ps(0.6931471805599453f);
    const __m512 p2 = _mm512_set1_ps(0.24022650695910071f);
    const __m512 p3 = _mm512_set1_ps(0.05550410866482157f);
    const __m512 p4 = _mm512_set1_ps(0.009618129107628477f);
 
    __m512 poly = _mm512_fmadd_ps(p4, rx, p3);
    poly = _mm512_fmadd_ps(poly, rx, p2);
    poly = _mm512_fmadd_ps(poly, rx, p1);
    poly = _mm512_fmadd_ps(poly, rx, p0);
 
    __m512i ti_int = _mm512_cvtps_epi32(ti);
    ti_int = _mm512_add_epi32(ti_int, _mm512_set1_epi32(127));
    ti_int = _mm512_slli_epi32(ti_int, 23);
    __m512 scale = _mm512_castsi512_ps(ti_int);
 
    return _mm512_mul_ps(poly, scale);
}
 
// Fast vectorized tanh: tanh(x) = (exp(2x) - 1) / (exp(2x) + 1)
static inline __m512 tanh512_fast(__m512 x) {
    __m512 two = _mm512_set1_ps(2.0f);
    __m512 one = _mm512_set1_ps(1.0f);
    __m512 exp2x = exp512_fast(_mm512_mul_ps(two, x));
    __m512 num = _mm512_sub_ps(exp2x, one);
    __m512 den = _mm512_add_ps(exp2x, one);
    return _mm512_div_ps(num, den);
}
#endif
 
#if defined(__AVX2__)
static inline __m256 exp256_fast(__m256 x) {
    x = _mm256_max_ps(x, _mm256_set1_ps(-88.0f));
    x = _mm256_min_ps(x, _mm256_set1_ps(88.0f));
 
    const __m256 log2e = _mm256_set1_ps(1.4426950408889634f);
    const __m256 c1 = _mm256_set1_ps(0.693359375f);
    const __m256 c2 = _mm256_set1_ps(-2.12194440e-4f);
 
    __m256 t = _mm256_mul_ps(x, log2e);
    __m256 ti = _mm256_round_ps(t, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
 
    __m256 rx = _mm256_sub_ps(x, _mm256_mul_ps(ti, c1));
    rx = _mm256_sub_ps(rx, _mm256_mul_ps(ti, c2));
 
    const __m256 p0 = _mm256_set1_ps(1.0f);
    const __m256 p1 = _mm256_set1_ps(0.6931471805599453f);
    const __m256 p2 = _mm256_set1_ps(0.24022650695910071f);
    const __m256 p3 = _mm256_set1_ps(0.05550410866482157f);
    const __m256 p4 = _mm256_set1_ps(0.009618129107628477f);
 
    __m256 poly = _mm256_fmadd_ps(p4, rx, p3);
    poly = _mm256_fmadd_ps(poly, rx, p2);
    poly = _mm256_fmadd_ps(poly, rx, p1);
    poly = _mm256_fmadd_ps(poly, rx, p0);
 
    __m256i ti_int = _mm256_cvtps_epi32(ti);
    ti_int = _mm256_add_epi32(ti_int, _mm256_set1_epi32(127));
    ti_int = _mm256_slli_epi32(ti_int, 23);
    __m256 scale = _mm256_castsi256_ps(ti_int);
 
    return _mm256_mul_ps(poly, scale);
}
 
static inline __m256 tanh256_fast(__m256 x) {
    __m256 two = _mm256_set1_ps(2.0f);
    __m256 one = _mm256_set1_ps(1.0f);
    __m256 exp2x = exp256_fast(_mm256_mul_ps(two, x));
    __m256 num = _mm256_sub_ps(exp2x, one);
    __m256 den = _mm256_add_ps(exp2x, one);
    return _mm256_div_ps(num, den);
}
#endif
 
/**
 * GELU activation forward (fast approximation, in-place)
 * @test test_gelu.py::TestGELUForward::test_gelu_fast_inplace
 * @test test_gelu.py::TestGELUForward::test_gelu_vs_exact
 * @test test_parity.py::test_gelu_parity
 *
 * Fast GELU approximation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
 * In-place on contiguous buffer.
 *
 * After changes: make test && make llamacpp-parity-full
 */
void gelu_fast_inplace(float *data, size_t n)
{
    const float sqrt_2_over_pi = 0.7978845608f;
    const float coeff = 0.044715f;
 
#if defined(__AVX512F__)
    const __m512 sqrt_2_pi_vec = _mm512_set1_ps(sqrt_2_over_pi);
    const __m512 coeff_vec = _mm512_set1_ps(coeff);
    const __m512 half_vec = _mm512_set1_ps(0.5f);
    const __m512 one_vec = _mm512_set1_ps(1.0f);
 
    size_t i = 0;
    for (; i + 16 <= n; i += 16) {
        __m512 x = _mm512_loadu_ps(&data[i]);
        __m512 x2 = _mm512_mul_ps(x, x);
        __m512 x3 = _mm512_mul_ps(x2, x);
 
        // inner = sqrt(2/pi) * (x + 0.044715 * x^3)
        __m512 inner = _mm512_fmadd_ps(coeff_vec, x3, x);
        inner = _mm512_mul_ps(sqrt_2_pi_vec, inner);
 
        // result = 0.5 * x * (1 + tanh(inner))
        __m512 tanh_val = tanh512_fast(inner);
        __m512 one_plus_tanh = _mm512_add_ps(one_vec, tanh_val);
        __m512 result = _mm512_mul_ps(half_vec, _mm512_mul_ps(x, one_plus_tanh));
 
        _mm512_storeu_ps(&data[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = data[i];
        float x3 = x * x * x;
        float inner = sqrt_2_over_pi * (x + coeff * x3);
        data[i] = 0.5f * x * (1.0f + tanhf(inner));
    }
 
#elif defined(__AVX2__)
    const __m256 sqrt_2_pi_vec = _mm256_set1_ps(sqrt_2_over_pi);
    const __m256 coeff_vec = _mm256_set1_ps(coeff);
    const __m256 half_vec = _mm256_set1_ps(0.5f);
    const __m256 one_vec = _mm256_set1_ps(1.0f);
 
    size_t i = 0;
    for (; i + 8 <= n; i += 8) {
        __m256 x = _mm256_loadu_ps(&data[i]);
        __m256 x2 = _mm256_mul_ps(x, x);
        __m256 x3 = _mm256_mul_ps(x2, x);
 
        // inner = sqrt(2/pi) * (x + 0.044715 * x^3)
        __m256 inner = _mm256_fmadd_ps(coeff_vec, x3, x);
        inner = _mm256_mul_ps(sqrt_2_pi_vec, inner);
 
        // result = 0.5 * x * (1 + tanh(inner))
        __m256 tanh_val = tanh256_fast(inner);
        __m256 one_plus_tanh = _mm256_add_ps(one_vec, tanh_val);
        __m256 result = _mm256_mul_ps(half_vec, _mm256_mul_ps(x, one_plus_tanh));
 
        _mm256_storeu_ps(&data[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = data[i];
        float x3 = x * x * x;
        float inner = sqrt_2_over_pi * (x + coeff * x3);
        data[i] = 0.5f * x * (1.0f + tanhf(inner));
    }
 
#elif defined(__AVX__)
    // AVX1: Vectorize arithmetic, use scalar tanh
    const __m256 sqrt_2_pi_vec = _mm256_set1_ps(sqrt_2_over_pi);
    const __m256 coeff_vec = _mm256_set1_ps(coeff);
    const __m256 half_vec = _mm256_set1_ps(0.5f);
    const __m256 one_vec = _mm256_set1_ps(1.0f);
 
    size_t i = 0;
    float inner_arr[8] __attribute__((aligned(32)));
    float tanh_arr[8] __attribute__((aligned(32)));
 
    for (; i + 8 <= n; i += 8) {
        __m256 x = _mm256_loadu_ps(&data[i]);
        __m256 x2 = _mm256_mul_ps(x, x);
        __m256 x3 = _mm256_mul_ps(x2, x);
 
        // inner = sqrt(2/pi) * (x + 0.044715 * x^3)
        __m256 coeff_x3 = _mm256_mul_ps(coeff_vec, x3);
        __m256 inner = _mm256_mul_ps(sqrt_2_pi_vec, _mm256_add_ps(x, coeff_x3));
 
        // Compute tanh scalarly
        _mm256_store_ps(inner_arr, inner);
        for (int j = 0; j < 8; ++j) {
            tanh_arr[j] = tanhf(inner_arr[j]);
        }
        __m256 tanh_val = _mm256_load_ps(tanh_arr);
 
        // result = 0.5 * x * (1 + tanh(inner))
        __m256 one_plus_tanh = _mm256_add_ps(one_vec, tanh_val);
        __m256 result = _mm256_mul_ps(half_vec, _mm256_mul_ps(x, one_plus_tanh));
 
        _mm256_storeu_ps(&data[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = data[i];
        float x3 = x * x * x;
        float inner = sqrt_2_over_pi * (x + coeff * x3);
        data[i] = 0.5f * x * (1.0f + tanhf(inner));
    }
 
#else
    // Scalar fallback
    for (size_t i = 0; i < n; ++i) {
        float x = data[i];
        float x3 = x * x * x;
        float inner = sqrt_2_over_pi * (x + coeff * x3);
        data[i] = 0.5f * x * (1.0f + tanhf(inner));
    }
#endif
}
 
// Exact GELU backward using the tanh-based approximation derivative, adapted
// from C-Transformer's backward_gelu. Operates element-wise on contiguous
// buffers.
// Derivative: d/dx GELU(x) = 0.5 * (1 + tanh(g)) + 0.5 * x * sech^2(g) * g'
// where g = sqrt(2/pi) * (x + 0.044715 * x^3)
//       g' = sqrt(2/pi) * (1 + 3 * 0.044715 * x^2)
void gelu_backward_exact(const float *input,
                         const float *d_output,
                         float *d_input,
                         size_t n)
{
    const float sqrt_2_over_pi = 0.7978845608f;
    const float coeff = 0.044715f;
 
#if defined(__AVX512F__)
    const __m512 sqrt_2_pi_vec = _mm512_set1_ps(sqrt_2_over_pi);
    const __m512 coeff_vec = _mm512_set1_ps(coeff);
    const __m512 coeff3_vec = _mm512_set1_ps(3.0f * coeff);
    const __m512 half_vec = _mm512_set1_ps(0.5f);
    const __m512 one_vec = _mm512_set1_ps(1.0f);
 
    size_t i = 0;
    for (; i + 16 <= n; i += 16) {
        __m512 x = _mm512_loadu_ps(&input[i]);
        __m512 dy = _mm512_loadu_ps(&d_output[i]);
 
        __m512 x2 = _mm512_mul_ps(x, x);
        __m512 x3 = _mm512_mul_ps(x2, x);
 
        // g = sqrt(2/pi) * (x + 0.044715 * x^3)
        __m512 g = _mm512_fmadd_ps(coeff_vec, x3, x);
        g = _mm512_mul_ps(sqrt_2_pi_vec, g);
 
        __m512 tanh_g = tanh512_fast(g);
 
        // g' = sqrt(2/pi) * (1 + 3 * 0.044715 * x^2)
        __m512 g_prime = _mm512_fmadd_ps(coeff3_vec, x2, one_vec);
        g_prime = _mm512_mul_ps(sqrt_2_pi_vec, g_prime);
 
        // sech^2(g) = 1 - tanh^2(g)
        __m512 sech2_g = _mm512_fnmadd_ps(tanh_g, tanh_g, one_vec);
 
        // gelu_derivative = 0.5 * (1 + tanh_g) + 0.5 * x * sech2_g * g_prime
        __m512 term1 = _mm512_mul_ps(half_vec, _mm512_add_ps(one_vec, tanh_g));
        __m512 term2 = _mm512_mul_ps(half_vec, _mm512_mul_ps(x, _mm512_mul_ps(sech2_g, g_prime)));
        __m512 gelu_deriv = _mm512_add_ps(term1, term2);
 
        __m512 result = _mm512_mul_ps(dy, gelu_deriv);
        _mm512_storeu_ps(&d_input[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = input[i];
        float x3 = x * x * x;
        float g = sqrt_2_over_pi * (x + coeff * x3);
        float tanh_g = tanhf(g);
        float x2 = x * x;
        float g_prime = sqrt_2_over_pi * (1.0f + 3.0f * coeff * x2);
        float sech2_g = 1.0f - tanh_g * tanh_g;
        float gelu_derivative = 0.5f * (1.0f + tanh_g) + 0.5f * x * sech2_g * g_prime;
        d_input[i] = d_output[i] * gelu_derivative;
    }
 
#elif defined(__AVX2__)
    const __m256 sqrt_2_pi_vec = _mm256_set1_ps(sqrt_2_over_pi);
    const __m256 coeff_vec = _mm256_set1_ps(coeff);
    const __m256 coeff3_vec = _mm256_set1_ps(3.0f * coeff);
    const __m256 half_vec = _mm256_set1_ps(0.5f);
    const __m256 one_vec = _mm256_set1_ps(1.0f);
 
    size_t i = 0;
    for (; i + 8 <= n; i += 8) {
        __m256 x = _mm256_loadu_ps(&input[i]);
        __m256 dy = _mm256_loadu_ps(&d_output[i]);
 
        __m256 x2 = _mm256_mul_ps(x, x);
        __m256 x3 = _mm256_mul_ps(x2, x);
 
        // g = sqrt(2/pi) * (x + 0.044715 * x^3)
        __m256 g = _mm256_fmadd_ps(coeff_vec, x3, x);
        g = _mm256_mul_ps(sqrt_2_pi_vec, g);
 
        __m256 tanh_g = tanh256_fast(g);
 
        // g' = sqrt(2/pi) * (1 + 3 * 0.044715 * x^2)
        __m256 g_prime = _mm256_fmadd_ps(coeff3_vec, x2, one_vec);
        g_prime = _mm256_mul_ps(sqrt_2_pi_vec, g_prime);
 
        // sech^2(g) = 1 - tanh^2(g)
        __m256 sech2_g = _mm256_fnmadd_ps(tanh_g, tanh_g, one_vec);
 
        // gelu_derivative = 0.5 * (1 + tanh_g) + 0.5 * x * sech2_g * g_prime
        __m256 term1 = _mm256_mul_ps(half_vec, _mm256_add_ps(one_vec, tanh_g));
        __m256 term2 = _mm256_mul_ps(half_vec, _mm256_mul_ps(x, _mm256_mul_ps(sech2_g, g_prime)));
        __m256 gelu_deriv = _mm256_add_ps(term1, term2);
 
        __m256 result = _mm256_mul_ps(dy, gelu_deriv);
        _mm256_storeu_ps(&d_input[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = input[i];
        float x3 = x * x * x;
        float g = sqrt_2_over_pi * (x + coeff * x3);
        float tanh_g = tanhf(g);
        float x2 = x * x;
        float g_prime = sqrt_2_over_pi * (1.0f + 3.0f * coeff * x2);
        float sech2_g = 1.0f - tanh_g * tanh_g;
        float gelu_derivative = 0.5f * (1.0f + tanh_g) + 0.5f * x * sech2_g * g_prime;
        d_input[i] = d_output[i] * gelu_derivative;
    }
 
#elif defined(__AVX__)
    // AVX1: Vectorize arithmetic, use scalar tanh
    const __m256 sqrt_2_pi_vec = _mm256_set1_ps(sqrt_2_over_pi);
    const __m256 coeff_vec = _mm256_set1_ps(coeff);
    const __m256 coeff3_vec = _mm256_set1_ps(3.0f * coeff);
    const __m256 half_vec = _mm256_set1_ps(0.5f);
    const __m256 one_vec = _mm256_set1_ps(1.0f);
 
    size_t i = 0;
    float g_arr[8] __attribute__((aligned(32)));
    float tanh_arr[8] __attribute__((aligned(32)));
 
    for (; i + 8 <= n; i += 8) {
        __m256 x = _mm256_loadu_ps(&input[i]);
        __m256 dy = _mm256_loadu_ps(&d_output[i]);
 
        __m256 x2 = _mm256_mul_ps(x, x);
        __m256 x3 = _mm256_mul_ps(x2, x);
 
        // g = sqrt(2/pi) * (x + 0.044715 * x^3)
        __m256 coeff_x3 = _mm256_mul_ps(coeff_vec, x3);
        __m256 g = _mm256_mul_ps(sqrt_2_pi_vec, _mm256_add_ps(x, coeff_x3));
 
        // Compute tanh scalarly
        _mm256_store_ps(g_arr, g);
        for (int j = 0; j < 8; ++j) {
            tanh_arr[j] = tanhf(g_arr[j]);
        }
        __m256 tanh_g = _mm256_load_ps(tanh_arr);
 
        // g' = sqrt(2/pi) * (1 + 3 * 0.044715 * x^2)
        __m256 coeff3_x2 = _mm256_mul_ps(coeff3_vec, x2);
        __m256 g_prime = _mm256_mul_ps(sqrt_2_pi_vec, _mm256_add_ps(one_vec, coeff3_x2));
 
        // sech^2(g) = 1 - tanh^2(g)
        __m256 tanh_g_sq = _mm256_mul_ps(tanh_g, tanh_g);
        __m256 sech2_g = _mm256_sub_ps(one_vec, tanh_g_sq);
 
        // gelu_derivative = 0.5 * (1 + tanh_g) + 0.5 * x * sech2_g * g_prime
        __m256 term1 = _mm256_mul_ps(half_vec, _mm256_add_ps(one_vec, tanh_g));
        __m256 term2 = _mm256_mul_ps(half_vec, _mm256_mul_ps(x, _mm256_mul_ps(sech2_g, g_prime)));
        __m256 gelu_deriv = _mm256_add_ps(term1, term2);
 
        __m256 result = _mm256_mul_ps(dy, gelu_deriv);
        _mm256_storeu_ps(&d_input[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = input[i];
        float x3 = x * x * x;
        float g = sqrt_2_over_pi * (x + coeff * x3);
        float tanh_g = tanhf(g);
        float x2 = x * x;
        float g_prime = sqrt_2_over_pi * (1.0f + 3.0f * coeff * x2);
        float sech2_g = 1.0f - tanh_g * tanh_g;
        float gelu_derivative = 0.5f * (1.0f + tanh_g) + 0.5f * x * sech2_g * g_prime;
        d_input[i] = d_output[i] * gelu_derivative;
    }
 
#else
    // Scalar fallback
    for (size_t i = 0; i < n; ++i) {
        float x = input[i];
 
        float x3 = x * x * x;
        float g = sqrt_2_over_pi * (x + coeff * x3);
        float tanh_g = tanhf(g);
 
        float x2 = x * x;
        float g_prime = sqrt_2_over_pi * (1.0f + 3.0f * coeff * x2);
 
        float sech2_g = 1.0f - tanh_g * tanh_g;
        float gelu_derivative =
            0.5f * (1.0f + tanh_g) + 0.5f * x * sech2_g * g_prime;
 
        d_input[i] = d_output[i] * gelu_derivative;
    }
#endif
}
 
// Scalar-only exact GELU forward using standard library tanhf.
// This is slower than gelu_fast_inplace but provides maximum accuracy.
// Used by BF16 wrapper where conversion overhead dominates anyway.
void gelu_exact_inplace(float *data, size_t n)
{
    const float sqrt_2_over_pi = 0.7978845608f;
    const float coeff = 0.044715f;
 
    for (size_t i = 0; i < n; ++i) {
        float x = data[i];
        float x3 = x * x * x;
        float inner = sqrt_2_over_pi * (x + coeff * x3);
        data[i] = 0.5f * x * (1.0f + tanhf(inner));
    }
}
 
// Scalar-only exact GELU backward using standard library tanhf.
// This is slower than gelu_backward_exact but provides maximum accuracy.
// Used by BF16 wrapper where conversion overhead dominates anyway.
void gelu_backward_scalar(const float *input,
                          const float *d_output,
                          float *d_input,
                          size_t n)
{
    const float sqrt_2_over_pi = 0.7978845608f;
    const float coeff = 0.044715f;
 
    for (size_t i = 0; i < n; ++i) {
        float x = input[i];
        float x3 = x * x * x;
        float g = sqrt_2_over_pi * (x + coeff * x3);
        float tanh_g = tanhf(g);
        float x2 = x * x;
        float g_prime = sqrt_2_over_pi * (1.0f + 3.0f * coeff * x2);
        float sech2_g = 1.0f - tanh_g * tanh_g;
        float gelu_derivative = 0.5f * (1.0f + tanh_g) + 0.5f * x * sech2_g * g_prime;
        d_input[i] = d_output[i] * gelu_derivative;
    }
}
 
// Fast approximate GELU backward, adapted from C-Transformer's backward_gelu_fast.
// Uses sigmoid approximation: GELU(x) ≈ x * sigmoid(1.702 * x)
// Derivative: s * (1 + x * (1 - s) * 1.702) where s = sigmoid(1.702 * x)
void gelu_backward_fast(const float *input,
                        const float *d_output,
                        float *d_input,
                        size_t n)
{
    const float beta = 1.702f;
 
#if defined(__AVX512F__)
    const __m512 beta_vec = _mm512_set1_ps(beta);
    const __m512 one_vec = _mm512_set1_ps(1.0f);
    const __m512 neg_beta_vec = _mm512_set1_ps(-beta);
 
    size_t i = 0;
    for (; i + 16 <= n; i += 16) {
        __m512 x = _mm512_loadu_ps(&input[i]);
        __m512 dy = _mm512_loadu_ps(&d_output[i]);
 
        // s = sigmoid(beta * x) = 1 / (1 + exp(-beta * x))
        __m512 neg_beta_x = _mm512_mul_ps(neg_beta_vec, x);
        __m512 exp_neg = exp512_fast(neg_beta_x);
        __m512 s = _mm512_div_ps(one_vec, _mm512_add_ps(one_vec, exp_neg));
 
        // gelu_derivative = s * (1 + x * (1 - s) * beta)
        __m512 one_minus_s = _mm512_sub_ps(one_vec, s);
        __m512 inner = _mm512_fmadd_ps(_mm512_mul_ps(x, one_minus_s), beta_vec, one_vec);
        __m512 gelu_deriv = _mm512_mul_ps(s, inner);
 
        __m512 result = _mm512_mul_ps(dy, gelu_deriv);
        _mm512_storeu_ps(&d_input[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = input[i];
        float s = 1.0f / (1.0f + expf(-beta * x));
        float gelu_derivative = s * (1.0f + x * (1.0f - s) * beta);
        d_input[i] = d_output[i] * gelu_derivative;
    }
 
#elif defined(__AVX2__)
    const __m256 beta_vec = _mm256_set1_ps(beta);
    const __m256 one_vec = _mm256_set1_ps(1.0f);
    const __m256 neg_beta_vec = _mm256_set1_ps(-beta);
 
    size_t i = 0;
    for (; i + 8 <= n; i += 8) {
        __m256 x = _mm256_loadu_ps(&input[i]);
        __m256 dy = _mm256_loadu_ps(&d_output[i]);
 
        // s = sigmoid(beta * x) = 1 / (1 + exp(-beta * x))
        __m256 neg_beta_x = _mm256_mul_ps(neg_beta_vec, x);
        __m256 exp_neg = exp256_fast(neg_beta_x);
        __m256 s = _mm256_div_ps(one_vec, _mm256_add_ps(one_vec, exp_neg));
 
        // gelu_derivative = s * (1 + x * (1 - s) * beta)
        __m256 one_minus_s = _mm256_sub_ps(one_vec, s);
        __m256 inner = _mm256_fmadd_ps(_mm256_mul_ps(x, one_minus_s), beta_vec, one_vec);
        __m256 gelu_deriv = _mm256_mul_ps(s, inner);
 
        __m256 result = _mm256_mul_ps(dy, gelu_deriv);
        _mm256_storeu_ps(&d_input[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = input[i];
        float s = 1.0f / (1.0f + expf(-beta * x));
        float gelu_derivative = s * (1.0f + x * (1.0f - s) * beta);
        d_input[i] = d_output[i] * gelu_derivative;
    }
 
#elif defined(__AVX__)
    // AVX1: Vectorize arithmetic, use scalar exp
    const __m256 beta_vec = _mm256_set1_ps(beta);
    const __m256 one_vec = _mm256_set1_ps(1.0f);
    const __m256 neg_beta_vec = _mm256_set1_ps(-beta);
 
    size_t i = 0;
    float neg_beta_x_arr[8] __attribute__((aligned(32)));
    float exp_arr[8] __attribute__((aligned(32)));
 
    for (; i + 8 <= n; i += 8) {
        __m256 x = _mm256_loadu_ps(&input[i]);
        __m256 dy = _mm256_loadu_ps(&d_output[i]);
 
        // s = sigmoid(beta * x) = 1 / (1 + exp(-beta * x))
        __m256 neg_beta_x = _mm256_mul_ps(neg_beta_vec, x);
 
        // Compute exp scalarly
        _mm256_store_ps(neg_beta_x_arr, neg_beta_x);
        for (int j = 0; j < 8; ++j) {
            exp_arr[j] = expf(neg_beta_x_arr[j]);
        }
        __m256 exp_neg = _mm256_load_ps(exp_arr);
 
        __m256 s = _mm256_div_ps(one_vec, _mm256_add_ps(one_vec, exp_neg));
 
        // gelu_derivative = s * (1 + x * (1 - s) * beta)
        __m256 one_minus_s = _mm256_sub_ps(one_vec, s);
        __m256 x_one_minus_s = _mm256_mul_ps(x, one_minus_s);
        __m256 x_one_minus_s_beta = _mm256_mul_ps(x_one_minus_s, beta_vec);
        __m256 inner = _mm256_add_ps(one_vec, x_one_minus_s_beta);
        __m256 gelu_deriv = _mm256_mul_ps(s, inner);
 
        __m256 result = _mm256_mul_ps(dy, gelu_deriv);
        _mm256_storeu_ps(&d_input[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = input[i];
        float s = 1.0f / (1.0f + expf(-beta * x));
        float gelu_derivative = s * (1.0f + x * (1.0f - s) * beta);
        d_input[i] = d_output[i] * gelu_derivative;
    }
#endif
}
 
// ============================================================================
// GeGLU Kernels - Gated GELU (used in LLaMA, Mistral, etc.)
// ============================================================================
//
// GeGLU: out = GELU(a) * b where input x = [a, b] along the last dimension
// Input shape: [tokens, 2 * dim]
// Output shape: [tokens, dim]
//
// The input is split along the last dimension:
// - First half: a = x[..., :dim] -> GELU(a)
// - Second half: b = x[..., dim:] -> element-wise multiply
// - out = GELU(a) * b
 
/**
 * GeGLU forward pass (fp32)
 * @test test_geglu.py::TestGeGLU::test_geglu_forward_fp32
 *
 * Computes out = GELU(a) * b where x = [a, b] along last dimension.
 * Input shape: [tokens, 2 * dim], Output shape: [tokens, dim]
 *
 * After changes: make test
 */
void geglu_forward_fp32(const float *x, float *out, int tokens, int dim)
{
    const float sqrt_2_over_pi = 0.7978845608f;
    const float coeff = 0.044715f;
 
    const int inner_dim = dim * 2;
 
#if defined(__AVX512F__)
    const __m512 sqrt_2_pi_vec = _mm512_set1_ps(sqrt_2_over_pi);
    const __m512 coeff_vec = _mm512_set1_ps(coeff);
    const __m512 half_vec = _mm512_set1_ps(0.5f);
    const __m512 one_vec = _mm512_set1_ps(1.0f);
 
    for (int t = 0; t < tokens; ++t) {
        const float *x_ptr = x + (size_t)t * inner_dim;
        float *out_ptr = out + (size_t)t * dim;
 
        int d = 0;
        // Process first half (a) with GELU, second half (b) directly
        for (; d + 32 <= dim; d += 32) {
            // Load a (first half of inner_dim)
            __m512 a0 = _mm512_loadu_ps(&x_ptr[d]);
            __m512 a1 = _mm512_loadu_ps(&x_ptr[d + 16]);
 
            // Compute GELU(a)
            __m512 a0_sq = _mm512_mul_ps(a0, a0);
            __m512 a0_cu = _mm512_mul_ps(a0_sq, a0);
            __m512 a1_sq = _mm512_mul_ps(a1, a1);
            __m512 a1_cu = _mm512_mul_ps(a1_sq, a1);
 
            // inner = sqrt(2/pi) * (a + 0.044715 * a^3)
            __m512 inner0 = _mm512_fmadd_ps(coeff_vec, a0_cu, a0);
            __m512 inner1 = _mm512_fmadd_ps(coeff_vec, a1_cu, a1);
            inner0 = _mm512_mul_ps(sqrt_2_pi_vec, inner0);
            inner1 = _mm512_mul_ps(sqrt_2_pi_vec, inner1);
 
            // tanh(inner)
            __m512 tanh0 = tanh512_fast(inner0);
            __m512 tanh1 = tanh512_fast(inner1);
 
            // GELU = 0.5 * a * (1 + tanh)
            __m512 gelu0 = _mm512_mul_ps(half_vec, _mm512_mul_ps(a0, _mm512_add_ps(one_vec, tanh0)));
            __m512 gelu1 = _mm512_mul_ps(half_vec, _mm512_mul_ps(a1, _mm512_add_ps(one_vec, tanh1)));
 
            // Load b (second half of inner_dim)
            __m512 b0 = _mm512_loadu_ps(&x_ptr[dim + d]);
            __m512 b1 = _mm512_loadu_ps(&x_ptr[dim + d + 16]);
 
            // out = GELU(a) * b
            _mm512_storeu_ps(&out_ptr[d], _mm512_mul_ps(gelu0, b0));
            _mm512_storeu_ps(&out_ptr[d + 16], _mm512_mul_ps(gelu1, b1));
        }
        // Handle remaining
        for (; d < dim; ++d) {
            float a = x_ptr[d];
            float b = x_ptr[dim + d];
            float a3 = a * a * a;
            float inner = sqrt_2_over_pi * (a + coeff * a3);
            float gelu_a = 0.5f * a * (1.0f + tanhf(inner));
            out_ptr[d] = gelu_a * b;
        }
    }
 
#elif defined(__AVX2__)
    const __m256 sqrt_2_pi_vec = _mm256_set1_ps(sqrt_2_over_pi);
    const __m256 coeff_vec = _mm256_set1_ps(coeff);
    const __m256 half_vec = _mm256_set1_ps(0.5f);
    const __m256 one_vec = _mm256_set1_ps(1.0f);
 
    for (int t = 0; t < tokens; ++t) {
        const float *x_ptr = x + (size_t)t * inner_dim;
        float *out_ptr = out + (size_t)t * dim;
 
        int d = 0;
        for (; d + 16 <= dim; d += 16) {
            // Load a
            __m256 a0 = _mm256_loadu_ps(&x_ptr[d]);
            __m256 a1 = _mm256_loadu_ps(&x_ptr[d + 8]);
 
            // GELU(a)
            __m256 a0_sq = _mm256_mul_ps(a0, a0);
            __m256 a0_cu = _mm256_mul_ps(a0_sq, a0);
            __m256 a1_sq = _mm256_mul_ps(a1, a1);
            __m256 a1_cu = _mm256_mul_ps(a1_sq, a1);
 
            __m256 inner0 = _mm256_fmadd_ps(coeff_vec, a0_cu, a0);
            __m256 inner1 = _mm256_fmadd_ps(coeff_vec, a1_cu, a1);
            inner0 = _mm256_mul_ps(sqrt_2_pi_vec, inner0);
            inner1 = _mm256_mul_ps(sqrt_2_pi_vec, inner1);
 
            __m256 tanh0 = tanh256_fast(inner0);
            __m256 tanh1 = tanh256_fast(inner1);
 
            __m256 gelu0 = _mm256_mul_ps(half_vec, _mm256_mul_ps(a0, _mm256_add_ps(one_vec, tanh0)));
            __m256 gelu1 = _mm256_mul_ps(half_vec, _mm256_mul_ps(a1, _mm256_add_ps(one_vec, tanh1)));
 
            // b
            __m256 b0 = _mm256_loadu_ps(&x_ptr[dim + d]);
            __m256 b1 = _mm256_loadu_ps(&x_ptr[dim + d + 8]);
 
            _mm256_storeu_ps(&out_ptr[d], _mm256_mul_ps(gelu0, b0));
            _mm256_storeu_ps(&out_ptr[d + 8], _mm256_mul_ps(gelu1, b1));
        }
        for (; d < dim; ++d) {
            float a = x_ptr[d];
            float b = x_ptr[dim + d];
            float a3 = a * a * a;
            float inner = sqrt_2_over_pi * (a + coeff * a3);
            float gelu_a = 0.5f * a * (1.0f + tanhf(inner));
            out_ptr[d] = gelu_a * b;
        }
    }
 
#elif defined(__AVX__)
    const __m256 sqrt_2_pi_vec = _mm256_set1_ps(sqrt_2_over_pi);
    const __m256 coeff_vec = _mm256_set1_ps(coeff);
    const __m256 half_vec = _mm256_set1_ps(0.5f);
    const __m256 one_vec = _mm256_set1_ps(1.0f);
 
    float inner_arr[8] __attribute__((aligned(32)));
    float tanh_arr[8] __attribute__((aligned(32)));
 
    for (int t = 0; t < tokens; ++t) {
        const float *x_ptr = x + (size_t)t * inner_dim;
        float *out_ptr = out + (size_t)t * dim;
 
        int d = 0;
        for (; d + 8 <= dim; d += 8) {
            __m256 a = _mm256_loadu_ps(&x_ptr[d]);
            __m256 a_sq = _mm256_mul_ps(a, a);
            __m256 a_cu = _mm256_mul_ps(a_sq, a);
 
            __m256 coeff_a_cu = _mm256_mul_ps(coeff_vec, a_cu);
            __m256 inner = _mm256_mul_ps(sqrt_2_pi_vec, _mm256_add_ps(a, coeff_a_cu));
 
            _mm256_store_ps(inner_arr, inner);
            for (int j = 0; j < 8; ++j) {
                tanh_arr[j] = tanhf(inner_arr[j]);
            }
            __m256 tanh_val = _mm256_load_ps(tanh_arr);
 
            __m256 gelu = _mm256_mul_ps(half_vec, _mm256_mul_ps(a, _mm256_add_ps(one_vec, tanh_val)));
            __m256 b = _mm256_loadu_ps(&x_ptr[dim + d]);
 
            _mm256_storeu_ps(&out_ptr[d], _mm256_mul_ps(gelu, b));
        }
        for (; d < dim; ++d) {
            float a = x_ptr[d];
            float b = x_ptr[dim + d];
            float a3 = a * a * a;
            float inner = sqrt_2_over_pi * (a + coeff * a3);
            float gelu_a = 0.5f * a * (1.0f + tanhf(inner));
            out_ptr[d] = gelu_a * b;
        }
    }
 
#else
    // Scalar fallback
    for (int t = 0; t < tokens; ++t) {
        const float *x_ptr = x + (size_t)t * inner_dim;
        float *out_ptr = out + (size_t)t * dim;
 
        for (int d = 0; d < dim; ++d) {
            float a = x_ptr[d];
            float b = x_ptr[dim + d];
            float a3 = a * a * a;
            float inner = sqrt_2_over_pi * (a + coeff * a3);
            float gelu_a = 0.5f * a * (1.0f + tanhf(inner));
            out_ptr[d] = gelu_a * b;
        }
    }
#endif
}
 
/**
 * GeGLU forward pass (bf16)
 * @test test_geglu.py::TestGeGLU::test_geglu_forward_bf16
 *
 * BF16 version: converts to FP32, computes, converts back.
 * Caller provides scratch buffer of size 3 * tokens * dim * sizeof(float).
 *
 * Layout:
 *   - scratch[0 : 2*tokens*dim]     = FP32 input [a, b]
 *   - scratch[2*tokens*dim : ...]   = FP32 output
 *
 * Note: We need separate buffers for input and output to avoid overlap
 * when tokens > 1. The input is 2*dim per token, output is dim per token.
 *
 * After changes: make test
 */
void geglu_forward_bf16(const uint16_t *x, uint16_t *out, int tokens, int dim, float *scratch)
{
    if (!x || !out || !scratch) return;
 
    const size_t fp32_size = (size_t)tokens * (size_t)dim;
    const size_t input_size = fp32_size * 2;  // [a, b] = 2*dim per token
    float *fp32_input = scratch;
    float *fp32_output = scratch + input_size;
 
    // Convert BF16 input to FP32
    bf16_tensor_to_float(x, fp32_input, input_size);
 
    // Run FP32 GeGLU (output goes to separate buffer to avoid overlap)
    geglu_forward_fp32(fp32_input, fp32_output, tokens, dim);
 
    // Convert FP32 output back to BF16
    float_tensor_to_bf16(fp32_output, out, fp32_size);
}
 
/**
 * GeGLU backward pass (fp32)
 * @test test_geglu.py::TestGeGLU::test_geglu_backward_fp32
 *
 * dL/dx given dL/d(out) where out = GELU(a) * b
 * Chain rule:
 *   dL/da = dL/dout * d(GELU)/da * b
 *   dL/db = dL/dout * GELU(a)
 *
 * After changes: make test
 */
void geglu_backward_fp32(const float *x,
                         const float *d_out,
                         float *d_x,
                         int tokens,
                         int dim)
{
    const float sqrt_2_over_pi = 0.7978845608f;
    const float coeff = 0.044715f;
 
    const int inner_dim = dim * 2;
 
    for (int t = 0; t < tokens; ++t) {
        const float *x_ptr = x + (size_t)t * inner_dim;
        const float *d_out_ptr = d_out + (size_t)t * dim;
        float *d_x_ptr = d_x + (size_t)t * inner_dim;
 
        for (int d = 0; d < dim; ++d) {
            float a = x_ptr[d];
            float b = x_ptr[dim + d];
            float dout = d_out_ptr[d];
 
            // GELU(a) derivative components
            float a2 = a * a;
            float a3 = a2 * a;
            float g = sqrt_2_over_pi * (a + coeff * a3);
            float tanh_g = tanhf(g);
            float sech2_g = 1.0f - tanh_g * tanh_g;
            float g_prime = sqrt_2_over_pi * (1.0f + 3.0f * coeff * a2);
 
            // d(GELU)/da = 0.5 * (1 + tanh(g)) + 0.5 * a * sech^2(g) * g'
            float d_gelu = 0.5f * (1.0f + tanh_g) + 0.5f * a * sech2_g * g_prime;
 
            // dL/da = dL/dout * d(GELU)/da * b
            d_x_ptr[d] = dout * d_gelu * b;
 
            // dL/db = dL/dout * GELU(a)
            float gelu_a = 0.5f * a * (1.0f + tanh_g);
            d_x_ptr[dim + d] = dout * gelu_a;
        }
    }
}