gelu_kernels.c File Reference

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

#include <math.h>
#include <stddef.h>
#include <pthread.h>
#include <dlfcn.h>
#include <stdio.h>
#include <stdlib.h>
#include "bf16_utils.h"
#include "ckernel_quant.h"

Go to the source code of this file.

Macros
#define	_GNU_SOURCE

Typedefs
typedef void(*	ck_gelu_ggml_cpu_init_fn) (void)
typedef float(*	ck_gelu_ggml_fp16_to_fp32_fn) (ck_half)
typedef ck_half(*	ck_gelu_ggml_fp32_to_fp16_fn) (float)
typedef float(*	ck_gelu_math_f32_fn) (float)
typedef double(*	ck_gelu_math_f64_fn) (double)

Functions
static void	ck_gelu_ggml_runtime_init (void)
static void	ck_gelu_ggml_table_init (void)
static void	ck_gelu_reference_math_init (void)
static ck_gelu_math_f64_fn	ck_gelu_system_erf (void)
static ck_gelu_math_f32_fn	ck_gelu_system_tanhf (void)
static float	ck_gelu_tanh_f32 (float x)
static float	ck_gelu_tanh_ggml_reference_f32 (float x)
static void	ck_gelu_try_bind_runtime (void *handle)
void	gelu_backward_exact (const float *input, const float *d_output, float *d_input, size_t n)
void	gelu_backward_fast (const float *input, const float *d_output, float *d_input, size_t n)
void	gelu_backward_scalar (const float *input, const float *d_output, float *d_input, size_t n)
void	gelu_erf_bf16_storage (float *data, size_t n)
void	gelu_erf_fp64_f32_inplace (float *data, size_t n)
void	gelu_exact_inplace (float *data, size_t n)
void	gelu_fast_inplace (float *data, size_t n)
void	gelu_ggml_inplace (float *data, size_t n)
void	gelu_ggml_native_inplace (float *data, size_t n)
void	gelu_pytorch_erf_f32_inplace (float *data, size_t n)
void	gelu_pytorch_erf_sleef_bf16_storage (float *data, size_t n)
void	gelu_pytorch_tanh_bf16_storage (float *data, size_t n)

Variables
static pthread_once_t	ck_gelu_ggml_runtime_once = PTHREAD_ONCE_INIT
static ck_half	ck_gelu_ggml_table_f16 [1u<< 16]
static pthread_once_t	ck_gelu_ggml_table_once = PTHREAD_ONCE_INIT
static ck_gelu_math_f64_fn	ck_gelu_reference_erf = NULL
static pthread_once_t	ck_gelu_reference_math_once = PTHREAD_ONCE_INIT
static ck_gelu_math_f32_fn	ck_gelu_reference_tanhf = NULL
static ck_gelu_ggml_fp16_to_fp32_fn	ck_gelu_runtime_fp16_to_fp32 = NULL
static ck_gelu_ggml_fp32_to_fp16_fn	ck_gelu_runtime_fp32_to_fp16 = NULL
static void *	ck_gelu_runtime_handle = NULL
static int	ck_gelu_runtime_ready = 0
static const ck_half *	ck_gelu_runtime_table_f16 = NULL

Detailed Description

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)

Definition in file gelu_kernels.c.

Macro Definition Documentation

_GNU_SOURCE

#define _GNU_SOURCE

Definition at line 19 of file gelu_kernels.c.

Typedef Documentation

ck_gelu_ggml_cpu_init_fn

typedef void(* ck_gelu_ggml_cpu_init_fn) (void)

Definition at line 48 of file gelu_kernels.c.

ck_gelu_ggml_fp16_to_fp32_fn

typedef float(* ck_gelu_ggml_fp16_to_fp32_fn) (ck_half)

Definition at line 50 of file gelu_kernels.c.

ck_gelu_ggml_fp32_to_fp16_fn

typedef ck_half(* ck_gelu_ggml_fp32_to_fp16_fn) (float)

Definition at line 49 of file gelu_kernels.c.

ck_gelu_math_f32_fn

typedef float(* ck_gelu_math_f32_fn) (float)

Definition at line 51 of file gelu_kernels.c.

ck_gelu_math_f64_fn

typedef double(* ck_gelu_math_f64_fn) (double)

Definition at line 52 of file gelu_kernels.c.

Function Documentation

ck_gelu_ggml_runtime_init()

static void ck_gelu_ggml_runtime_init ( void  )

static

Definition at line 119 of file gelu_kernels.c.

                                            {
    ck_gelu_try_bind_runtime(RTLD_DEFAULT);
    if (ck_gelu_runtime_ready) {
        return;
    }
 
    ck_gelu_runtime_handle = dlopen("llama.cpp/build/bin/libggml-cpu.so", RTLD_LAZY | RTLD_LOCAL);
    if (ck_gelu_runtime_handle) {
        ck_gelu_try_bind_runtime(ck_gelu_runtime_handle);
    }
}

References ck_gelu_runtime_handle, ck_gelu_runtime_ready, ck_gelu_try_bind_runtime(), and RTLD_DEFAULT.

Referenced by gelu_ggml_inplace().

ck_gelu_ggml_table_init()

static void ck_gelu_ggml_table_init ( void  )

static

Definition at line 131 of file gelu_kernels.c.

                                          {
    for (uint32_t i = 0; i < (1u << 16); ++i) {
        const ck_half x_fp16 = (ck_half) i;
        const float x = ggml_fp16_to_fp32(x_fp16);
        const float y = ck_gelu_tanh_ggml_reference_f32(x);
        ck_gelu_ggml_table_f16[i] = ggml_fp32_to_fp16(y);
    }
}

References ck_gelu_ggml_table_f16, ck_gelu_tanh_ggml_reference_f32(), ggml_fp16_to_fp32, and ggml_fp32_to_fp16.

Referenced by gelu_ggml_inplace(), and gelu_ggml_native_inplace().

ck_gelu_reference_math_init()

static void ck_gelu_reference_math_init ( void  )

static

Definition at line 63 of file gelu_kernels.c.

                                              {
#if defined(__linux__)
    void *handle = dlopen("libm.so.6", RTLD_NOW | RTLD_LOCAL);
    if (handle) {
        ck_gelu_reference_tanhf = (ck_gelu_math_f32_fn) dlsym(handle, "tanhf");
        ck_gelu_reference_erf = (ck_gelu_math_f64_fn) dlsym(handle, "erf");
    }
#endif
}

References ck_gelu_reference_erf, and ck_gelu_reference_tanhf.

Referenced by ck_gelu_system_erf(), and ck_gelu_system_tanhf().

ck_gelu_system_erf()

static ck_gelu_math_f64_fn ck_gelu_system_erf ( void  )

static

Definition at line 78 of file gelu_kernels.c.

                                                    {
    pthread_once(&ck_gelu_reference_math_once, ck_gelu_reference_math_init);
    return ck_gelu_reference_erf;
}

References ck_gelu_reference_erf, ck_gelu_reference_math_init(), and ck_gelu_reference_math_once.

Referenced by gelu_erf_bf16_storage(), and gelu_erf_fp64_f32_inplace().

ck_gelu_system_tanhf()

static ck_gelu_math_f32_fn ck_gelu_system_tanhf ( void  )

static

Definition at line 73 of file gelu_kernels.c.

                                                      {
    pthread_once(&ck_gelu_reference_math_once, ck_gelu_reference_math_init);
    return ck_gelu_reference_tanhf;
}

References ck_gelu_reference_math_init(), ck_gelu_reference_math_once, and ck_gelu_reference_tanhf.

Referenced by ck_gelu_tanh_ggml_reference_f32().

ck_gelu_tanh_f32()

static float ck_gelu_tanh_f32 ( float  x )

inlinestatic

Definition at line 36 of file gelu_kernels.c.

                                              {
    const float sqrt_2_over_pi = 0.7978845608f;
    const float coeff = 0.044715f;
    const float x3 = x * x * x;
    const float inner = sqrt_2_over_pi * (x + coeff * x3);
    return 0.5f * x * (1.0f + tanhf(inner));
}

Referenced by gelu_exact_inplace().

ck_gelu_tanh_ggml_reference_f32()

static float ck_gelu_tanh_ggml_reference_f32 ( float  x )

static

Definition at line 86 of file gelu_kernels.c.

                                                      {
    const float gelu_coef_a = 0.044715f;
    const float sqrt_2_over_pi = 0.79788456080286535588f;
    ck_gelu_math_f32_fn reference_tanhf = ck_gelu_system_tanhf();
    const float inner = sqrt_2_over_pi * x * (1.0f + gelu_coef_a * x * x);
    const float tanh_value = reference_tanhf ? reference_tanhf(inner) : tanhf(inner);
    return 0.5f * x * (1.0f + tanh_value);
}

References ck_gelu_system_tanhf().

Referenced by ck_gelu_ggml_table_init().

ck_gelu_try_bind_runtime()

static void ck_gelu_try_bind_runtime ( void *  handle )

static

Definition at line 98 of file gelu_kernels.c.

                                                   {
    ck_gelu_ggml_cpu_init_fn cpu_init_fn =
        (ck_gelu_ggml_cpu_init_fn) dlsym(handle, "ggml_cpu_init");
    ck_gelu_ggml_fp32_to_fp16_fn fp32_to_fp16_fn =
        (ck_gelu_ggml_fp32_to_fp16_fn) dlsym(handle, "ggml_fp32_to_fp16");
    ck_gelu_ggml_fp16_to_fp32_fn fp16_to_fp32_fn =
        (ck_gelu_ggml_fp16_to_fp32_fn) dlsym(handle, "ggml_fp16_to_fp32");
    const ck_half *table =
        (const ck_half *) dlsym(handle, "ggml_table_gelu_f16");
 
    if (!cpu_init_fn || !fp32_to_fp16_fn || !fp16_to_fp32_fn || !table) {
        return;
    }
 
    cpu_init_fn();
    ck_gelu_runtime_table_f16 = table;
    ck_gelu_runtime_fp32_to_fp16 = fp32_to_fp16_fn;
    ck_gelu_runtime_fp16_to_fp32 = fp16_to_fp32_fn;
    ck_gelu_runtime_ready = 1;
}

References ck_gelu_runtime_fp16_to_fp32, ck_gelu_runtime_fp32_to_fp16, ck_gelu_runtime_ready, and ck_gelu_runtime_table_f16.

Referenced by ck_gelu_ggml_runtime_init().

gelu_backward_exact()

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

Definition at line 368 of file gelu_kernels.c.

{
    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
}

References __attribute__().

gelu_backward_fast()

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

Definition at line 818 of file gelu_kernels.c.

{
    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;
    }
 
#else
    // Scalar fallback
    for (size_t i = 0; 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
}

References __attribute__().

Referenced by gelu_backward_fast_bf16().

gelu_backward_scalar()

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

Definition at line 794 of file gelu_kernels.c.

{
    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;
    }
}

Referenced by gelu_backward_exact_bf16().

gelu_erf_bf16_storage()

void gelu_erf_bf16_storage	(	float *	data,
		size_t	n
	)

Definition at line 609 of file gelu_kernels.c.

{
    const double inv_sqrt_2 = 0.707106781186547524400844362104849039;
    ck_gelu_math_f64_fn reference_erf = ck_gelu_system_erf();
    for (size_t i = 0; i < n; ++i) {
        const float x = bf16_to_float(float_to_bf16(data[i]));
        const double scaled = (double)x * inv_sqrt_2;
        const double erf_value = reference_erf ? reference_erf(scaled) : erf(scaled);
        const float output = (float)(0.5 * (double)x * (1.0 + erf_value));
        data[i] = bf16_to_float(float_to_bf16(output));
    }
}

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

gelu_erf_fp64_f32_inplace()

void gelu_erf_fp64_f32_inplace	(	float *	data,
		size_t	n
	)

Definition at line 566 of file gelu_kernels.c.

{
    const double inv_sqrt_2 = 0.707106781186547524400844362104849039;
    ck_gelu_math_f64_fn reference_erf = ck_gelu_system_erf();
    for (size_t i = 0; i < n; ++i) {
        const float x = data[i];
        const double scaled = (double)x * inv_sqrt_2;
        const double erf_value = reference_erf ? reference_erf(scaled) : erf(scaled);
        data[i] = (float)(0.5 * (double)x * (1.0 + erf_value));
    }
}

References ck_gelu_system_erf().

Referenced by gelu_pytorch_erf_f32_inplace().

gelu_exact_inplace()

void gelu_exact_inplace	(	float *	data,
		size_t	n
	)

Definition at line 557 of file gelu_kernels.c.

{
    for (size_t i = 0; i < n; ++i) {
        data[i] = ck_gelu_tanh_f32(data[i]);
    }
}

References ck_gelu_tanh_f32().

Referenced by gelu_fast_inplace_bf16(), and mlp_token_parallel_exact().

gelu_fast_inplace()

void gelu_fast_inplace	(	float *	data,
		size_t	n
	)

GELU activation forward (fast approximation, in-place)

Test:

test_gelu.py::TestGELUForward::test_gelu_fast_inplace

test_gelu.py::TestGELUForward::test_gelu_vs_exact

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

Definition at line 243 of file gelu_kernels.c.

{
    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
}

References __attribute__().

Referenced by geglu_forward_fp32(), and mlp_token_parallel().

gelu_ggml_inplace()

void gelu_ggml_inplace	(	float *	data,
		size_t	n
	)

Definition at line 753 of file gelu_kernels.c.

{
    pthread_once(&ck_gelu_ggml_runtime_once, ck_gelu_ggml_runtime_init);
    if (ck_gelu_runtime_ready) {
        for (size_t i = 0; i < n; ++i) {
            const float x = data[i];
            if (x <= -10.0f) {
                data[i] = 0.0f;
                continue;
            }
            if (x >= 10.0f) {
                data[i] = x;
                continue;
            }
            const ck_half x_fp16 = ck_gelu_runtime_fp32_to_fp16(x);
            const ck_half y_fp16 = ck_gelu_runtime_table_f16[(uint16_t) x_fp16];
            data[i] = ck_gelu_runtime_fp16_to_fp32(y_fp16);
        }
        return;
    }
 
    pthread_once(&ck_gelu_ggml_table_once, ck_gelu_ggml_table_init);
    for (size_t i = 0; i < n; ++i) {
        const float x = data[i];
        if (x <= -10.0f) {
            data[i] = 0.0f;
            continue;
        }
        if (x >= 10.0f) {
            data[i] = x;
            continue;
        }
        const ck_half x_fp16 = ggml_fp32_to_fp16(x);
        const ck_half y_fp16 = ck_gelu_ggml_table_f16[(uint16_t) x_fp16];
        data[i] = ggml_fp16_to_fp32(y_fp16);
    }
}

References ck_gelu_ggml_runtime_init(), ck_gelu_ggml_runtime_once, ck_gelu_ggml_table_f16, ck_gelu_ggml_table_init(), ck_gelu_ggml_table_once, ck_gelu_runtime_fp16_to_fp32, ck_gelu_runtime_fp32_to_fp16, ck_gelu_runtime_ready, ck_gelu_runtime_table_f16, ggml_fp16_to_fp32, and ggml_fp32_to_fp16.

gelu_ggml_native_inplace()

void gelu_ggml_native_inplace	(	float *	data,
		size_t	n
	)

Definition at line 741 of file gelu_kernels.c.

{
    pthread_once(&ck_gelu_ggml_table_once, ck_gelu_ggml_table_init);
    for (size_t i = 0; i < n; ++i) {
        const float x = data[i];
        if (x <= -10.0f) data[i] = 0.0f;
        else if (x >= 10.0f) data[i] = x;
        else data[i] = ggml_fp16_to_fp32(
            ck_gelu_ggml_table_f16[(uint16_t)ggml_fp32_to_fp16(x)]);
    }
}

References ck_gelu_ggml_table_f16, ck_gelu_ggml_table_init(), ck_gelu_ggml_table_once, ggml_fp16_to_fp32, and ggml_fp32_to_fp16.

Referenced by geglu_forward_ggml_native().

gelu_pytorch_erf_f32_inplace()

void gelu_pytorch_erf_f32_inplace	(	float *	data,
		size_t	n
	)

Definition at line 580 of file gelu_kernels.c.

{
    gelu_erf_fp64_f32_inplace(data, n);
}

References gelu_erf_fp64_f32_inplace().

gelu_pytorch_erf_sleef_bf16_storage()

void gelu_pytorch_erf_sleef_bf16_storage	(	float *	data,
		size_t	n
	)

Definition at line 682 of file gelu_kernels.c.

{
#if defined(__AVX512F__)
    pthread_once(&ck_pytorch_sleef_once, ck_bind_pytorch_sleef);
    if (!ck_pytorch_sleef_expf16) {
        fprintf(stderr,
                "[CK] PyTorch-exact BF16 GELU requires Sleef_expf16_u10; "
                "set CK_SLEEF_LIBRARY to libtorch_cpu.so or libsleef.so\n");
        abort();
    }
 
    const __m512 alpha = _mm512_set1_ps(0.70710678118654752440f);
    const __m512 half = _mm512_set1_ps(0.5f);
    const __m512 one = _mm512_set1_ps(1.0f);
    const __m512 neg_zero = _mm512_set1_ps(-0.0f);
    const __m512 p = _mm512_set1_ps(0.3275911f);
    const __m512 p1 = _mm512_set1_ps(0.254829592f);
    const __m512 p2 = _mm512_set1_ps(-0.284496736f);
    const __m512 p3 = _mm512_set1_ps(1.421413741f);
    const __m512 p4 = _mm512_set1_ps(-1.453152027f);
    const __m512 p5 = _mm512_set1_ps(1.061405429f);
    size_t i = 0;
    for (; i + 16 <= n; i += 16) {
        __m512 x = _mm512_loadu_ps(data + i);
        __m512 erf_arg = _mm512_mul_ps(x, alpha);
        __m512 sign = _mm512_and_ps(neg_zero, erf_arg);
        __m512 abs_arg = _mm512_abs_ps(erf_arg);
        __m512 t = _mm512_div_ps(one, _mm512_fmadd_ps(p, abs_arg, one));
        __m512 r = _mm512_fmadd_ps(p5, t, p4);
        r = _mm512_fmadd_ps(r, t, p3);
        r = _mm512_fmadd_ps(r, t, p2);
        r = _mm512_fmadd_ps(r, t, p1);
        __m512 arg_sq = _mm512_mul_ps(erf_arg, erf_arg);
        __m512 exp_neg_sq = ck_pytorch_sleef_expf16(_mm512_xor_ps(neg_zero, arg_sq));
        __m512 neg_exp_t = _mm512_mul_ps(_mm512_xor_ps(neg_zero, exp_neg_sq), t);
        __m512 erf_x = _mm512_xor_ps(sign, _mm512_fmadd_ps(neg_exp_t, r, one));
        __m512 y = _mm512_mul_ps(_mm512_mul_ps(x, half), _mm512_add_ps(one, erf_x));
        float lanes[16];
        _mm512_storeu_ps(lanes, y);
        for (size_t lane = 0; lane < 16; ++lane) {
            const uint16_t input_code = float_to_bf16(data[i + lane]);
            const uint16_t output_code = ck_pytorch_gelu_erf_bf16_edge(
                input_code, float_to_bf16(lanes[lane]));
            data[i + lane] = bf16_to_float(output_code);
        }
    }
    for (; i < n; ++i) {
        const float x = bf16_to_float(float_to_bf16(data[i]));
        const float y = (x * 0.5f) * (1.0f + erff(x * 0.70710678118654752440f));
        data[i] = bf16_to_float(float_to_bf16(y));
    }
#else
    (void)data;
    (void)n;
    fprintf(stderr, "[CK] PyTorch-exact BF16 GELU requires an AVX-512 build\n");
    abort();
#endif
}

References bf16_to_float(), and float_to_bf16().

gelu_pytorch_tanh_bf16_storage()

void gelu_pytorch_tanh_bf16_storage	(	float *	data,
		size_t	n
	)

Definition at line 590 of file gelu_kernels.c.

{
    /* PyTorch's x86 BF16 kernel widens to FP32 and evaluates tanh through
     * SLEEF's vector u10 provider.  That provider saturates at this exact
     * inner-argument boundary; libc tanhf retains a small tail and can round
     * to a different BF16 code before the following projection. */
    const float sleef_tanh_saturation = 8.664339742f;
    for (size_t i = 0; i < n; ++i) {
        const float x = bf16_to_float(float_to_bf16(data[i]));
        const float x3 = x * x * x;
        const float inner = 0.7978845608f * (x + 0.044715f * x3);
        const float tanh_inner = fabsf(inner) > sleef_tanh_saturation
            ? copysignf(1.0f, inner)
            : tanhf(inner);
        const float output = 0.5f * x * (1.0f + tanh_inner);
        data[i] = bf16_to_float(float_to_bf16(output));
    }
}

References bf16_to_float(), and float_to_bf16().

Variable Documentation

ck_gelu_ggml_runtime_once

pthread_once_t ck_gelu_ggml_runtime_once = PTHREAD_ONCE_INIT

static

Definition at line 46 of file gelu_kernels.c.

Referenced by gelu_ggml_inplace().

ck_gelu_ggml_table_f16

ck_half ck_gelu_ggml_table_f16[1u<< 16]

static

Definition at line 44 of file gelu_kernels.c.

Referenced by ck_gelu_ggml_table_init(), gelu_ggml_inplace(), and gelu_ggml_native_inplace().

ck_gelu_ggml_table_once

pthread_once_t ck_gelu_ggml_table_once = PTHREAD_ONCE_INIT

static

Definition at line 45 of file gelu_kernels.c.

Referenced by gelu_ggml_inplace(), and gelu_ggml_native_inplace().

ck_gelu_reference_erf

ck_gelu_math_f64_fn ck_gelu_reference_erf = NULL

static

Definition at line 60 of file gelu_kernels.c.

Referenced by ck_gelu_reference_math_init(), and ck_gelu_system_erf().

ck_gelu_reference_math_once

pthread_once_t ck_gelu_reference_math_once = PTHREAD_ONCE_INIT

static

Definition at line 61 of file gelu_kernels.c.

Referenced by ck_gelu_system_erf(), and ck_gelu_system_tanhf().

ck_gelu_reference_tanhf

ck_gelu_math_f32_fn ck_gelu_reference_tanhf = NULL

static

Definition at line 59 of file gelu_kernels.c.

Referenced by ck_gelu_reference_math_init(), and ck_gelu_system_tanhf().

ck_gelu_runtime_fp16_to_fp32

ck_gelu_ggml_fp16_to_fp32_fn ck_gelu_runtime_fp16_to_fp32 = NULL

static

Definition at line 56 of file gelu_kernels.c.

Referenced by ck_gelu_try_bind_runtime(), and gelu_ggml_inplace().

ck_gelu_runtime_fp32_to_fp16

ck_gelu_ggml_fp32_to_fp16_fn ck_gelu_runtime_fp32_to_fp16 = NULL

static

Definition at line 55 of file gelu_kernels.c.

Referenced by ck_gelu_try_bind_runtime(), and gelu_ggml_inplace().

ck_gelu_runtime_handle

void* ck_gelu_runtime_handle = NULL

static

Definition at line 57 of file gelu_kernels.c.

Referenced by ck_gelu_ggml_runtime_init().

ck_gelu_runtime_ready

int ck_gelu_runtime_ready = 0

static

Definition at line 58 of file gelu_kernels.c.

Referenced by ck_gelu_ggml_runtime_init(), ck_gelu_try_bind_runtime(), and gelu_ggml_inplace().

ck_gelu_runtime_table_f16

const ck_half* ck_gelu_runtime_table_f16 = NULL

static

Definition at line 54 of file gelu_kernels.c.

Referenced by ck_gelu_try_bind_runtime(), and gelu_ggml_inplace().