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

Macros
#define	_GNU_SOURCE

Typedefs
typedef float(*	ck_recurrent_libm_f32_fn) (float)

Functions
static void	ck_bind_recurrent_llama_libm (void)

void	recurrent_dt_gate_backward (const float d_gate, const float alpha, const float dt_bias, const float a, float d_alpha, float d_dt_bias, float *d_a, int rows, int dim)

void	recurrent_dt_gate_expanded_forward (const float alpha, const float dt_bias, const float a, float gate, int rows, int num_heads, int state_dim)

void	recurrent_dt_gate_forward (const float alpha, const float dt_bias, const float a, float gate, int rows, int num_heads, int state_dim)

void	recurrent_dt_gate_forward_pytorch_fp32 (const float alpha, const float dt_bias, const float a, float gate, int rows, int num_heads, int state_dim)

static float	recurrent_sigmoid (float x)

void	recurrent_sigmoid_forward_ggml (const float x, float out, int rows, int dim)

void	recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output (const float x, float out, int rows, int dim)

void	recurrent_silu_backward (const float d_out, const float x, float *d_x, int rows, int dim)

void	recurrent_silu_forward (const float x, float out, int rows, int dim)

void	recurrent_silu_forward_ggml (const float x, float out, int rows, int dim)

void	recurrent_silu_forward_pytorch_bf16_input_fp32_output (const float x, float out, int rows, int dim)

void	recurrent_silu_forward_pytorch_bf16_storage (const float x, float out, int rows, int dim)

static float	recurrent_softplus (float x)

Variables
static void *	ck_recurrent_libm_handle = NULL

static pthread_once_t	ck_recurrent_libm_once = PTHREAD_ONCE_INIT

static ck_recurrent_libm_f32_fn	ck_recurrent_llama_expf = NULL

static ck_recurrent_libm_f32_fn	ck_recurrent_llama_logf = NULL

static ck_recurrent_libm_f32_fn	ck_recurrent_pytorch_log1pf = NULL

Macro Definition Documentation

◆ _GNU_SOURCE

#define _GNU_SOURCE

Definition at line 2 of file recurrent_gate_kernels.c.

Typedef Documentation

◆ ck_recurrent_libm_f32_fn

typedef float(* ck_recurrent_libm_f32_fn) (float)

Definition at line 17 of file recurrent_gate_kernels.c.

Function Documentation

◆ ck_bind_recurrent_llama_libm()

static void ck_bind_recurrent_llama_libm ( void )

static

Definition at line 24 of file recurrent_gate_kernels.c.

{
    ck_recurrent_libm_handle = dlopen("libm.so.6", RTLD_NOW | RTLD_LOCAL);
    if (ck_recurrent_libm_handle) {
        ck_recurrent_llama_expf =
            (ck_recurrent_libm_f32_fn)dlsym(ck_recurrent_libm_handle, "expf");
        ck_recurrent_llama_logf =
            (ck_recurrent_libm_f32_fn)dlsym(ck_recurrent_libm_handle, "logf");
        ck_recurrent_pytorch_log1pf =
            (ck_recurrent_libm_f32_fn)dlsym(ck_recurrent_libm_handle, "log1pf");
    }
    if (!ck_recurrent_llama_expf || !ck_recurrent_llama_logf) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: llama.cpp recurrent softplus "
                "requires expf/logf from libm.so.6\n");
        abort();
    }
}

References ck_recurrent_libm_handle, ck_recurrent_llama_expf, ck_recurrent_llama_logf, and ck_recurrent_pytorch_log1pf.

Referenced by recurrent_dt_gate_forward_pytorch_fp32(), and recurrent_softplus().

◆ recurrent_dt_gate_backward()

void recurrent_dt_gate_backward	(	const float *	d_gate,
		const float *	alpha,
		const float *	dt_bias,
		const float *	a,
		float *	d_alpha,
		float *	d_dt_bias,
		float *	d_a,
		int	rows,
		int	dim
	)

Definition at line 104 of file recurrent_gate_kernels.c.

                                         {
    for (int col = 0; col < dim; ++col) {
        d_dt_bias[col] = 0.0f;
        d_a[col] = 0.0f;
    }
 
    for (int row = 0; row < rows; ++row) {
        const float *d_gate_row = d_gate + (size_t) row * (size_t) dim;
        const float *alpha_row = alpha + (size_t) row * (size_t) dim;
        float *d_alpha_row = d_alpha + (size_t) row * (size_t) dim;
        for (int col = 0; col < dim; ++col) {
            const float x = alpha_row[col] + dt_bias[col];
            const float sp = recurrent_softplus(x);
            const float sig = recurrent_sigmoid(x);
            const float d_out = d_gate_row[col];
            d_a[col] += d_out * sp;
            {
                const float d_x = d_out * a[col] * sig;
                d_alpha_row[col] = d_x;
                d_dt_bias[col] += d_x;
            }
        }
    }
}

References recurrent_sigmoid(), and recurrent_softplus().

◆ recurrent_dt_gate_expanded_forward()

void recurrent_dt_gate_expanded_forward	(	const float *	alpha,
		const float *	dt_bias,
		const float *	a,
		float *	gate,
		int	rows,
		int	num_heads,
		int	state_dim
	)

Definition at line 83 of file recurrent_gate_kernels.c.

                                                       {
    for (int row = 0; row < rows; ++row) {
        const float *alpha_row = alpha + (size_t) row * (size_t) num_heads;
        float *gate_row = gate + (size_t) row * (size_t) num_heads * (size_t) state_dim;
        for (int h = 0; h < num_heads; ++h) {
            const float sp = recurrent_softplus(alpha_row[h] + dt_bias[h]);
            const float *a_head = a + (size_t) h * (size_t) state_dim;
            float *gate_head = gate_row + (size_t) h * (size_t) state_dim;
            for (int col = 0; col < state_dim; ++col) {
                gate_head[col] = sp * a_head[col];
            }
        }
    }
}

References recurrent_softplus().

◆ recurrent_dt_gate_forward()

void recurrent_dt_gate_forward	(	const float *	alpha,
		const float *	dt_bias,
		const float *	a,
		float *	gate,
		int	rows,
		int	num_heads,
		int	state_dim
	)

Definition at line 65 of file recurrent_gate_kernels.c.

                                              {
    const int dim = num_heads * state_dim;
    for (int row = 0; row < rows; ++row) {
        const float *alpha_row = alpha + (size_t) row * (size_t) dim;
        float *gate_row = gate + (size_t) row * (size_t) dim;
        for (int col = 0; col < dim; ++col) {
            const float x = alpha_row[col] + dt_bias[col];
            gate_row[col] = recurrent_softplus(x) * a[col];
        }
    }
}

References recurrent_softplus().

Referenced by ck_test_recurrent_dt_gate().

◆ recurrent_dt_gate_forward_pytorch_fp32()

void recurrent_dt_gate_forward_pytorch_fp32	(	const float *	alpha,
		const float *	dt_bias,
		const float *	a,
		float *	gate,
		int	rows,
		int	num_heads,
		int	state_dim
	)

Definition at line 179 of file recurrent_gate_kernels.c.

{
    if (!alpha || !dt_bias || !a || !gate || rows < 0 || num_heads < 0 || state_dim != 1) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: invalid PyTorch FP32 recurrent dt-gate arguments\n");
        abort();
    }
#if defined(__AVX512F__)
    pthread_once(&ck_recurrent_sleef_once, ck_bind_recurrent_pytorch_sleef);
    if (!ck_recurrent_pytorch_expf16 || !ck_recurrent_pytorch_log1pf16) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: PyTorch FP32 recurrent dt gate requires "
                "SLEEF Sleef_expf16_u10 and Sleef_log1pf16_u10; set CK_SLEEF_LIBRARY\n");
        abort();
    }
    if ((num_heads & 15) != 0) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: PyTorch AVX-512 recurrent dt gate requires "
                "a head count divisible by 16 (got %d)\n",
                num_heads);
        abort();
    }
    pthread_once(&ck_recurrent_libm_once, ck_bind_recurrent_llama_libm);
    if (!ck_recurrent_pytorch_log1pf) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: PyTorch FP32 recurrent dt gate requires "
                "log1pf from libm.so.6\n");
        abort();
    }
    const __m512 threshold = _mm512_set1_ps(20.0f);
    const int count = rows * num_heads;
    int index = 0;
    for (; index + 32 <= count; index += 32) {
        for (int half = 0; half < 2; ++half) {
            float x_lanes[16] __attribute__((aligned(64)));
            float a_lanes[16] __attribute__((aligned(64)));
            const int base = index + half * 16;
            for (int lane = 0; lane < 16; ++lane) {
                const int head = (base + lane) % num_heads;
                x_lanes[lane] = alpha[base + lane] + dt_bias[head];
                a_lanes[lane] = a[head];
            }
            const __m512 x = _mm512_load_ps(x_lanes);
            const __m512 softplus = _mm512_mask_blend_ps(
                _mm512_cmp_ps_mask(x, threshold, _CMP_GT_OQ),
                ck_recurrent_pytorch_log1pf16(ck_recurrent_pytorch_expf16(x)),
                x);
            _mm512_storeu_ps(
                gate + base,
                _mm512_mul_ps(softplus, _mm512_load_ps(a_lanes)));
        }
    }
    for (; index < count; ++index) {
        const int head = index % num_heads;
        const float x = alpha[index] + dt_bias[head];
        const float softplus = x > 20.0f
            ? x
            : ck_recurrent_pytorch_log1pf(ck_recurrent_llama_expf(x));
        gate[index] = softplus * a[head];
    }
#else
    fprintf(stderr,
            "HARD KERNEL CONTRACT FAULT: PyTorch FP32 recurrent dt gate requires AVX-512\n");
    abort();
#endif
}

References __attribute__(), ck_bind_recurrent_llama_libm(), ck_recurrent_libm_once, ck_recurrent_llama_expf, and ck_recurrent_pytorch_log1pf.

◆ recurrent_sigmoid()

static float recurrent_sigmoid ( float x )

inlinestatic

Definition at line 54 of file recurrent_gate_kernels.c.

                                               {
    if (x >= 0.0f) {
        float z = expf(-x);
        return 1.0f / (1.0f + z);
    }
    {
        float z = expf(x);
        return z / (1.0f + z);
    }
}

Referenced by recurrent_dt_gate_backward(), recurrent_silu_backward(), and recurrent_silu_forward().

◆ recurrent_sigmoid_forward_ggml()

void recurrent_sigmoid_forward_ggml	(	const float *	x,
		float *	out,
		int	rows,
		int	dim
	)

Definition at line 501 of file recurrent_gate_kernels.c.

                                             {
    float (*volatile llama_expf)(float) = expf;
    for (int row = 0; row < rows; ++row) {
        const float *x_row = x + (size_t) row * (size_t) dim;
        float *out_row = out + (size_t) row * (size_t) dim;
        for (int col = 0; col < dim; ++col) {
            out_row[col] = 1.0f / (1.0f + llama_expf(-x_row[col]));
        }
    }
}

Referenced by hyper_connection_mix_quantized(), and recurrent_norm_sigmoid_gate_llama_avx2_forward().

◆ recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output()

void recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output	(	const float *	x,
		float *	out,
		int	rows,
		int	dim
	)

Definition at line 339 of file recurrent_gate_kernels.c.

{
    if (!x || !out || rows < 0 || dim < 0) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: invalid BF16-input FP32-output sigmoid arguments\n");
        abort();
    }
#if defined(__AVX512F__)
    pthread_once(&ck_recurrent_sleef_once, ck_bind_recurrent_pytorch_sleef);
    if (!ck_recurrent_pytorch_expf16) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: PyTorch BF16-input sigmoid requires "
                "SLEEF Sleef_expf16_u10; set CK_SLEEF_LIBRARY\n");
        abort();
    }
#endif
    const int count = rows * dim;
    int i = 0;
#if defined(__AVX512F__)
    for (; i + 16 <= count; i += 16) {
        float lanes[16] __attribute__((aligned(64)));
        for (int lane = 0; lane < 16; ++lane) {
            lanes[lane] = bf16_to_float(float_to_bf16(x[i + lane]));
        }
        const __m512 values = _mm512_load_ps(lanes);
        const __m512 denominator = _mm512_add_ps(
            _mm512_set1_ps(1.0f),
            ck_recurrent_pytorch_expf16(
                _mm512_sub_ps(_mm512_setzero_ps(), values)));
        _mm512_storeu_ps(out + i,
                         _mm512_div_ps(_mm512_set1_ps(1.0f), denominator));
    }
#endif
    for (; i < count; ++i) {
        const float value = bf16_to_float(float_to_bf16(x[i]));
        out[i] = 1.0f / (1.0f + expf(-value));
    }
}

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

Referenced by attn_gate_sigmoid_mul_pytorch_bf16_storage(), and recurrent_norm_sigmoid_gate_pytorch_bf16_storage().

◆ recurrent_silu_backward()

void recurrent_silu_backward	(	const float *	d_out,
		const float *	x,
		float *	d_x,
		int	rows,
		int	dim
	)

Definition at line 515 of file recurrent_gate_kernels.c.

                                      {
    for (int row = 0; row < rows; ++row) {
        const float *d_out_row = d_out + (size_t) row * (size_t) dim;
        const float *x_row = x + (size_t) row * (size_t) dim;
        float *d_x_row = d_x + (size_t) row * (size_t) dim;
        for (int col = 0; col < dim; ++col) {
            const float xv = x_row[col];
            const float sig = recurrent_sigmoid(xv);
            d_x_row[col] = d_out_row[col] * (sig + xv * sig * (1.0f - sig));
        }
    }
}

References recurrent_sigmoid().

◆ recurrent_silu_forward()

void recurrent_silu_forward	(	const float *	x,
		float *	out,
		int	rows,
		int	dim
	)

Definition at line 137 of file recurrent_gate_kernels.c.

                                     {
    for (int row = 0; row < rows; ++row) {
        const float *x_row = x + (size_t) row * (size_t) dim;
        float *out_row = out + (size_t) row * (size_t) dim;
        for (int col = 0; col < dim; ++col) {
            const float xv = x_row[col];
            out_row[col] = xv * recurrent_sigmoid(xv);
        }
    }
}

References recurrent_sigmoid().

Referenced by ck_test_recurrent_silu().

◆ recurrent_silu_forward_ggml()

void recurrent_silu_forward_ggml	(	const float *	x,
		float *	out,
		int	rows,
		int	dim
	)

Definition at line 468 of file recurrent_gate_kernels.c.

                                          {
    for (int row = 0; row < rows; ++row) {
        const float *x_row = x + (size_t) row * (size_t) dim;
        float *out_row = out + (size_t) row * (size_t) dim;
        int col = 0;
#if defined(__AVX512F__) && defined(__AVX512DQ__)
        for (; col + 16 <= dim; col += 16) {
            const __m512 xv = _mm512_loadu_ps(x_row + col);
            const __m512 neg = _mm512_sub_ps(_mm512_setzero_ps(), xv);
            const __m512 denom = _mm512_add_ps(
                _mm512_set1_ps(1.0f), recurrent_ggml_expf_avx512(neg));
            _mm512_storeu_ps(out_row + col, _mm512_div_ps(xv, denom));
        }
#elif defined(__AVX2__) && defined(__FMA__)
        for (; col + 8 <= dim; col += 8) {
            const __m256 xv = _mm256_loadu_ps(x_row + col);
            const __m256 neg = _mm256_sub_ps(_mm256_setzero_ps(), xv);
            const __m256 denom = _mm256_add_ps(
                _mm256_set1_ps(1.0f), recurrent_ggml_expf_avx2(neg));
            _mm256_storeu_ps(out_row + col, _mm256_div_ps(xv, denom));
        }
#endif
        float (*volatile llama_expf)(float) = expf;
        for (; col < dim; ++col) {
            const float xv = x_row[col];
            out_row[col] = xv / (1.0f + llama_expf(-xv));
        }
    }
}

Referenced by hyper_connection_mix_quantized(), qwen4_ple_gate_conv_inject_impl(), and recurrent_norm_gate_llama_avx2_forward().

◆ recurrent_silu_forward_pytorch_bf16_input_fp32_output()

void recurrent_silu_forward_pytorch_bf16_input_fp32_output	(	const float *	x,
		float *	out,
		int	rows,
		int	dim
	)

Definition at line 299 of file recurrent_gate_kernels.c.

{
    if (!x || !out || rows < 0 || dim < 0) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: invalid BF16-input FP32-output SiLU arguments\n");
        abort();
    }
#if defined(__AVX512F__)
    pthread_once(&ck_recurrent_sleef_once, ck_bind_recurrent_pytorch_sleef);
    if (!ck_recurrent_pytorch_expf16) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: PyTorch BF16-input SiLU requires "
                "SLEEF Sleef_expf16_u10; set CK_SLEEF_LIBRARY\n");
        abort();
    }
#endif
    const int count = rows * dim;
    int i = 0;
#if defined(__AVX512F__)
    for (; i + 16 <= count; i += 16) {
        float lanes[16] __attribute__((aligned(64)));
        for (int lane = 0; lane < 16; ++lane) {
            lanes[lane] = bf16_to_float(float_to_bf16(x[i + lane]));
        }
        const __m512 values = _mm512_load_ps(lanes);
        const __m512 denominator = _mm512_add_ps(
            _mm512_set1_ps(1.0f),
            ck_recurrent_pytorch_expf16(
                _mm512_sub_ps(_mm512_setzero_ps(), values)));
        _mm512_storeu_ps(out + i, _mm512_div_ps(values, denominator));
    }
#endif
    for (; i < count; ++i) {
        const float value = bf16_to_float(float_to_bf16(x[i]));
        out[i] = value / (1.0f + expf(-value));
    }
}

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

Referenced by recurrent_norm_gate_pytorch_bf16_storage().

◆ recurrent_silu_forward_pytorch_bf16_storage()

void recurrent_silu_forward_pytorch_bf16_storage	(	const float *	x,
		float *	out,
		int	rows,
		int	dim
	)

Definition at line 252 of file recurrent_gate_kernels.c.

{
    if (!x || !out || rows < 0 || dim < 0) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: invalid PyTorch BF16 recurrent SiLU arguments\n");
        abort();
    }
#if defined(__AVX512F__)
    pthread_once(&ck_recurrent_sleef_once, ck_bind_recurrent_pytorch_sleef);
    if (!ck_recurrent_pytorch_expf16) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: PyTorch BF16 recurrent SiLU requires "
                "SLEEF Sleef_expf16_u10; set CK_SLEEF_LIBRARY\n");
        abort();
    }
#endif
    for (int row = 0; row < rows; ++row) {
        const float *src = x + (size_t)row * (size_t)dim;
        float *dst = out + (size_t)row * (size_t)dim;
        int col = 0;
#if defined(__AVX512F__)
        for (; col + 16 <= dim; col += 16) {
            float lanes[16] __attribute__((aligned(64)));
            for (int lane = 0; lane < 16; ++lane) {
                lanes[lane] = bf16_to_float(float_to_bf16(src[col + lane]));
            }
            const __m512 values = _mm512_load_ps(lanes);
            const __m512 denominator = _mm512_add_ps(
                _mm512_set1_ps(1.0f),
                ck_recurrent_pytorch_expf16(
                    _mm512_sub_ps(_mm512_setzero_ps(), values)));
            _mm512_store_ps(lanes, _mm512_div_ps(values, denominator));
            for (int lane = 0; lane < 16; ++lane) {
                dst[col + lane] = bf16_to_float(float_to_bf16(lanes[lane]));
            }
        }
#endif
        for (; col < dim; ++col) {
            const float value = bf16_to_float(float_to_bf16(src[col]));
            const float silu = value / (1.0f + expf(-value));
            dst[col] = bf16_to_float(float_to_bf16(silu));
        }
    }
}

References __attribute__(), bf16_to_float(), float_to_bf16(), and silu().

◆ recurrent_softplus()

static float recurrent_softplus ( float x )

inlinestatic

Definition at line 43 of file recurrent_gate_kernels.c.

                                                {
    pthread_once(&ck_recurrent_libm_once, ck_bind_recurrent_llama_libm);
    if (x > 20.0f) {
        return x;
    }
    if (x < -20.0f) {
        return ck_recurrent_llama_expf(x);
    }
    return ck_recurrent_llama_logf(1.0f + ck_recurrent_llama_expf(x));
}

References ck_bind_recurrent_llama_libm(), ck_recurrent_libm_once, ck_recurrent_llama_expf, and ck_recurrent_llama_logf.

Referenced by recurrent_dt_gate_backward(), recurrent_dt_gate_expanded_forward(), and recurrent_dt_gate_forward().

Variable Documentation

◆ ck_recurrent_libm_handle

void* ck_recurrent_libm_handle = NULL

static

Definition at line 21 of file recurrent_gate_kernels.c.

Referenced by ck_bind_recurrent_llama_libm().

◆ ck_recurrent_libm_once

pthread_once_t ck_recurrent_libm_once = PTHREAD_ONCE_INIT

static

Definition at line 22 of file recurrent_gate_kernels.c.

Referenced by recurrent_dt_gate_forward_pytorch_fp32(), and recurrent_softplus().

◆ ck_recurrent_llama_expf

ck_recurrent_libm_f32_fn ck_recurrent_llama_expf = NULL

static

Definition at line 18 of file recurrent_gate_kernels.c.

Referenced by ck_bind_recurrent_llama_libm(), recurrent_dt_gate_forward_pytorch_fp32(), and recurrent_softplus().

◆ ck_recurrent_llama_logf

ck_recurrent_libm_f32_fn ck_recurrent_llama_logf = NULL

static

Definition at line 19 of file recurrent_gate_kernels.c.

Referenced by ck_bind_recurrent_llama_libm(), and recurrent_softplus().

◆ ck_recurrent_pytorch_log1pf

ck_recurrent_libm_f32_fn ck_recurrent_pytorch_log1pf = NULL

static

Definition at line 20 of file recurrent_gate_kernels.c.

Referenced by ck_bind_recurrent_llama_libm(), and recurrent_dt_gate_forward_pytorch_fp32().

Macros

Typedefs

Functions

Variables

Macro Definition Documentation

◆ _GNU_SOURCE

Typedef Documentation

◆ ck_recurrent_libm_f32_fn

Function Documentation

◆ ck_bind_recurrent_llama_libm()

◆ recurrent_dt_gate_backward()

◆ recurrent_dt_gate_expanded_forward()

◆ recurrent_dt_gate_forward()

◆ recurrent_dt_gate_forward_pytorch_fp32()

◆ recurrent_sigmoid()

◆ recurrent_sigmoid_forward_ggml()

◆ recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output()

◆ recurrent_silu_backward()

◆ recurrent_silu_forward()

◆ recurrent_silu_forward_ggml()

◆ recurrent_silu_forward_pytorch_bf16_input_fp32_output()

◆ recurrent_silu_forward_pytorch_bf16_storage()

◆ recurrent_softplus()

Variable Documentation

◆ ck_recurrent_libm_handle

◆ ck_recurrent_libm_once

◆ ck_recurrent_llama_expf

◆ ck_recurrent_llama_logf

◆ ck_recurrent_pytorch_log1pf