Fused GEMM Kernels with activations. More...

#include "ckernel_engine.h"
#include <math.h>

Functions
static float	fast_gelu_scalar (float x)

void	gemm_bias_gelu_fused (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_bias_relu_fused (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_bias_silu_fused (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_swiglu_fused (const float x, const float W_gate, const float W_up, const float b_gate, const float b_up, float output, int M, int N, int K)

Detailed Description

Fused GEMM Kernels with activations.

CK-ENGINE KERNEL RULES:

NO malloc/free - memory via bump allocator, pointers passed in
NO OpenMP - parallelization at orchestrator/codegen layer
API must define: inputs, outputs, workspace, and memory layouts
Pure computation - deterministic, no side effects

After changes: make test && make llamacpp-parity-full

GEMM operations fused with activations (ReLU, GELU, SiLU) and SwiGLU. The key benefit: intermediate results stay in registers, avoiding DRAM round-trips between operations.

Supported operations:

gemm_bias_relu_fused: C = ReLU(A @ B^T + bias)
gemm_bias_gelu_fused: C = GELU(A @ B^T + bias)
gemm_bias_silu_fused: C = SiLU(A @ B^T + bias)
gemm_swiglu_fused: C = SiLU(x @ W_gate) * (x @ W_up)

All kernels support:

AVX1 SIMD (256-bit vectors, no FMA)
OpenMP parallelization
Scalar fallback

Definition in file gemm_fused_kernels.c.

Function Documentation

◆ fast_gelu_scalar()

static float fast_gelu_scalar ( float x )

inlinestatic

Definition at line 74 of file gemm_fused_kernels.c.

                                              {
    float sx = 1.702f * x;
    float sig = 1.0f / (1.0f + expf(-sx));
    return x * sig;
}

Referenced by gemm_bias_gelu_fused().

◆ gemm_bias_gelu_fused()

void gemm_bias_gelu_fused	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 131 of file gemm_fused_kernels.c.

{
#if defined(__AVX__)
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            __m256 sum_vec = _mm256_setzero_ps();
            int k;
            for (k = 0; k <= K - 8; k += 8) {
                __m256 a_vec = _mm256_loadu_ps(&A[i * K + k]);
                __m256 b_vec = _mm256_loadu_ps(&B[j * K + k]);
                __m256 prod = _mm256_mul_ps(a_vec, b_vec);
                sum_vec = _mm256_add_ps(sum_vec, prod);
            }
            float sum = hsum256_ps_fused(sum_vec);
            for (; k < K; k++) {
                sum += A[i * K + k] * B[j * K + k];
            }
            sum += bias[j];
            C[i * N + j] = fast_gelu_scalar(sum);
        }
    }
#else
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            float sum = 0.0f;
            for (int k = 0; k < K; k++) {
                sum += A[i * K + k] * B[j * K + k];
            }
            sum += bias[j];
            C[i * N + j] = fast_gelu_scalar(sum);
        }
    }
#endif
}

References C, and fast_gelu_scalar().

◆ gemm_bias_relu_fused()

void gemm_bias_relu_fused	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 84 of file gemm_fused_kernels.c.

{
#if defined(__AVX__)
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            __m256 sum_vec = _mm256_setzero_ps();
            int k;
            for (k = 0; k <= K - 8; k += 8) {
                __m256 a_vec = _mm256_loadu_ps(&A[i * K + k]);
                __m256 b_vec = _mm256_loadu_ps(&B[j * K + k]);
                __m256 prod = _mm256_mul_ps(a_vec, b_vec);
                sum_vec = _mm256_add_ps(sum_vec, prod);
            }
            float sum = hsum256_ps_fused(sum_vec);
            for (; k < K; k++) {
                sum += A[i * K + k] * B[j * K + k];
            }
            // Fused: add bias and ReLU while still in register
            sum += bias[j];
            C[i * N + j] = sum > 0.0f ? sum : 0.0f;
        }
    }
#else
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            float sum = 0.0f;
            for (int k = 0; k < K; k++) {
                sum += A[i * K + k] * B[j * K + k];
            }
            sum += bias[j];
            C[i * N + j] = sum > 0.0f ? sum : 0.0f;
        }
    }
#endif
}

References C.

◆ gemm_bias_silu_fused()

void gemm_bias_silu_fused	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 177 of file gemm_fused_kernels.c.

{
#if defined(__AVX__)
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            __m256 sum_vec = _mm256_setzero_ps();
            int k;
            for (k = 0; k <= K - 8; k += 8) {
                __m256 a_vec = _mm256_loadu_ps(&A[i * K + k]);
                __m256 b_vec = _mm256_loadu_ps(&B[j * K + k]);
                __m256 prod = _mm256_mul_ps(a_vec, b_vec);
                sum_vec = _mm256_add_ps(sum_vec, prod);
            }
            float sum = hsum256_ps_fused(sum_vec);
            for (; k < K; k++) {
                sum += A[i * K + k] * B[j * K + k];
            }
            sum += bias[j];
            // SiLU: x * sigmoid(x)
            float sig = 1.0f / (1.0f + expf(-sum));
            C[i * N + j] = sum * sig;
        }
    }
#else
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            float sum = 0.0f;
            for (int k = 0; k < K; k++) {
                sum += A[i * K + k] * B[j * K + k];
            }
            sum += bias[j];
            float sig = 1.0f / (1.0f + expf(-sum));
            C[i * N + j] = sum * sig;
        }
    }
#endif
}

References C.

◆ gemm_swiglu_fused()

void gemm_swiglu_fused	(	const float *	x,
		const float *	W_gate,
		const float *	W_up,
		const float *	b_gate,
		const float *	b_up,
		float *	output,
		int	M,
		int	N,
		int	K
	)

Definition at line 241 of file gemm_fused_kernels.c.

{
#if defined(__AVX__)
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        const float *x_row = &x[i * K];
        float *out_row = &output[i * N];
 
        for (int j = 0; j < N; j++) {
            const float *w_gate_row = &W_gate[j * K];
            const float *w_up_row = &W_up[j * K];
 
            // Compute both dot products in parallel using SIMD
            __m256 gate_vec = _mm256_setzero_ps();
            __m256 up_vec = _mm256_setzero_ps();
 
            int k;
            for (k = 0; k <= K - 8; k += 8) {
                __m256 x_vec = _mm256_loadu_ps(&x_row[k]);
                __m256 wg_vec = _mm256_loadu_ps(&w_gate_row[k]);
                __m256 wu_vec = _mm256_loadu_ps(&w_up_row[k]);
 
                // gate += x * W_gate
                gate_vec = _mm256_add_ps(gate_vec, _mm256_mul_ps(x_vec, wg_vec));
                // up += x * W_up
                up_vec = _mm256_add_ps(up_vec, _mm256_mul_ps(x_vec, wu_vec));
            }
 
            // Horizontal sum
            float gate = hsum256_ps_fused(gate_vec);
            float up = hsum256_ps_fused(up_vec);
 
            // Scalar remainder
            for (; k < K; k++) {
                gate += x_row[k] * w_gate_row[k];
                up += x_row[k] * w_up_row[k];
            }
 
            // Add biases
            if (b_gate) gate += b_gate[j];
            if (b_up) up += b_up[j];
 
            // SwiGLU: SiLU(gate) * up = gate * sigmoid(gate) * up
            float sig = 1.0f / (1.0f + expf(-gate));
            out_row[j] = gate * sig * up;
        }
    }
#else
    // Scalar fallback
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            float gate = 0.0f;
            float up = 0.0f;
 
            for (int k = 0; k < K; k++) {
                gate += x[i * K + k] * W_gate[j * K + k];
                up += x[i * K + k] * W_up[j * K + k];
            }
 
            if (b_gate) gate += b_gate[j];
            if (b_up) up += b_up[j];
 
            // SwiGLU: SiLU(gate) * up
            float sig = 1.0f / (1.0f + expf(-gate));
            output[i * N + j] = gate * sig * up;
        }
    }
#endif
}

Referenced by ck_mlp_swiglu_forward_fused_token().

Functions