gemm_microkernel.c File Reference

GEMM Microkernel - High-Performance Register-Blocked Matrix Multiplication. More...

#include "ckernel_engine.h"
#include "cpu_features.h"
#include <string.h>
#include <stdlib.h>
#include <stdio.h>

Go to the source code of this file.

Macros
#define	GEMM_BACKEND   "Native"
#define	KC   (get_gemm_params()->KC)
#define	MC   (get_gemm_params()->MC)
#define	MR   (MR_FIXED)
#define	MR_FIXED   4
#define	NC   (get_gemm_params()->NC)
#define	NR   (NR_FIXED)
#define	NR_FIXED   4
#define	PACK_THRESHOLD   256

Functions
const char *	gemm_get_backend (void)
static void	gemm_init_threads (void)
void	gemm_microkernel (const float *A, const float *B, float *C, int M, int N, int K, int B_transposed)
void	gemm_microkernel_blocked (const float *A, const float *B, float *C, int M, int N, int K)
void	gemm_microkernel_blocked_bt (const float *A, const float *B, float *C, int M, int N, int K)
static void	gemm_microkernel_edge (int m, int n, int K, const float *A, int lda, const float *B, int ldb, float *C, int ldc, int first_k)
void	gemm_microkernel_packed (const float *A, const float *B, float *C, int M, int N, int K)
static void	gemm_microkernel_sequential (const float *A, const float *B, float *C, int M, int N, int K)
static void	pack_a_panel (const float *A, int lda, float *Ap, int mc, int kc, int mr)
static void	pack_b_panel (const float *B, int ldb, float *Bp, int kc, int nc, int nr)

Variables
static int	g_threads_initialized = 0

Detailed Description

GEMM Microkernel - High-Performance Register-Blocked Matrix Multiplication.

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

This file implements optimized GEMM microkernels with multiple backends:

1. USE_MKL: Intel MKL cblas_sgemm (best performance on Intel CPUs)
2. USE_ONEDNN: Intel oneDNN matmul primitive (Apache 2.0 licensed)
3. Native: Our own AVX-512/AVX2/AVX microkernels (no dependencies)

Build with: make USE_MKL=1 # Use Intel MKL make USE_ONEDNN=1 # Use Intel oneDNN make # Use native kernels

Layout: C[M,N] = A[M,K] @ B[K,N] (row-major)

Definition in file gemm_microkernel.c.

Macro Definition Documentation

GEMM_BACKEND

#define GEMM_BACKEND   "Native"

Definition at line 45 of file gemm_microkernel.c.

KC

#define KC   (get_gemm_params()->KC)

Definition at line 230 of file gemm_microkernel.c.

MC

#define MC   (get_gemm_params()->MC)

Definition at line 228 of file gemm_microkernel.c.

MR

#define MR   (MR_FIXED)

Definition at line 226 of file gemm_microkernel.c.

MR_FIXED

#define MR_FIXED   4

Definition at line 221 of file gemm_microkernel.c.

NC

#define NC   (get_gemm_params()->NC)

Definition at line 229 of file gemm_microkernel.c.

NR

#define NR   (NR_FIXED)

Definition at line 227 of file gemm_microkernel.c.

NR_FIXED

#define NR_FIXED   4

Definition at line 222 of file gemm_microkernel.c.

PACK_THRESHOLD

#define PACK_THRESHOLD   256

Definition at line 1132 of file gemm_microkernel.c.

Function Documentation

gemm_get_backend()

const char* gemm_get_backend

(

void

)

Definition at line 1160 of file gemm_microkernel.c.

                                   {
    return GEMM_BACKEND;
}

References GEMM_BACKEND.

gemm_init_threads()

static void gemm_init_threads ( void  )

static

Definition at line 915 of file gemm_microkernel.c.

                                    {
    if (g_threads_initialized) return;
 
#ifdef _OPENMP
    const CPUInfo* cpu = get_cpu_info();
    int physical_cores = cpu->num_cores;
 
    // Only use physical cores - hyperthreading hurts compute-bound GEMM
    if (physical_cores > 0) {
        int current_max = omp_get_max_threads();
        // Only reduce if we have more threads than physical cores
        if (current_max > physical_cores) {
            omp_set_num_threads(physical_cores);
        }
    }
#endif
    g_threads_initialized = 1;
}

References g_threads_initialized, get_cpu_info(), and CPUInfo::num_cores.

Referenced by gemm_microkernel_blocked().

gemm_microkernel()

void gemm_microkernel	(	const float *	A,
		const float *	B,
		float *	C,
		int	M,
		int	N,
		int	K,
		int	B_transposed
	)

Definition at line 1134 of file gemm_microkernel.c.

{
    if (B_transposed) {
        gemm_microkernel_blocked_bt(A, B, C, M, N, K);
    } else {
        // Use packed version for large matrices
        if (M >= PACK_THRESHOLD && N >= PACK_THRESHOLD && K >= PACK_THRESHOLD) {
            gemm_microkernel_packed(A, B, C, M, N, K);
        } else {
            gemm_microkernel_blocked(A, B, C, M, N, K);
        }
    }
}

References C, gemm_microkernel_blocked(), gemm_microkernel_blocked_bt(), gemm_microkernel_packed(), and PACK_THRESHOLD.

Referenced by gemm_blocked_serial().

gemm_microkernel_blocked()

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

Definition at line 934 of file gemm_microkernel.c.

{
    const int mr = MR;
    const int nr = NR;
 
    // Use sequential version for small matrices to avoid OpenMP overhead
    // Threshold tuned for typical 4-8 core systems
    if ((size_t)M * N * K <= 512ULL * 512 * 512) {
        gemm_microkernel_sequential(A, B, C, M, N, K);
        return;
    }
 
    // Initialize thread count to physical cores (once)
    gemm_init_threads();
 
    // Zero output first
    #pragma omp parallel for schedule(static)
    for (int i = 0; i < M; i++) {
        memset(&C[i * N], 0, N * sizeof(float));
    }
 
    // Block over K (outermost - for accumulation across all threads)
    for (int k0 = 0; k0 < K; k0 += KC) {
        int kb = (k0 + KC <= K) ? KC : (K - k0);
        int first_k = (k0 == 0);
 
        // Parallelize over M rows - each thread gets a chunk of M
        // This gives better cache locality than tile-level parallelism
        #pragma omp parallel for schedule(static)
        for (int m0 = 0; m0 < M; m0 += mr) {
            int mr_actual = (m0 + mr <= M) ? mr : (M - m0);
 
            // Each thread processes all N tiles for its M rows
            for (int n0 = 0; n0 < N; n0 += nr) {
                int nr_actual = (n0 + nr <= N) ? nr : (N - n0);
 
                const float *A_tile = &A[m0 * K + k0];
                const float *B_tile = &B[k0 * N + n0];
                float *C_tile = &C[m0 * N + n0];
 
                if (mr_actual == mr && nr_actual == nr) {
#if defined(__AVX512F__)
                    gemm_microkernel_6x32_avx512(kb, A_tile, K, B_tile, N, C_tile, N, first_k);
#elif defined(__FMA__)
                    gemm_microkernel_6x16_avx(kb, A_tile, K, B_tile, N, C_tile, N, first_k);
#elif defined(__AVX__)
                    gemm_microkernel_4x16_avx(kb, A_tile, K, B_tile, N, C_tile, N, first_k);
#else
                    gemm_microkernel_edge(mr_actual, nr_actual, kb, A_tile, K, B_tile, N, C_tile, N, first_k);
#endif
                } else {
                    gemm_microkernel_edge(mr_actual, nr_actual, kb, A_tile, K, B_tile, N, C_tile, N, first_k);
                }
            }
        }
    }
}

References C, gemm_init_threads(), gemm_microkernel_edge(), gemm_microkernel_sequential(), KC, MR, and NR.

Referenced by gemm_microkernel(), and gemm_microkernel_packed().

gemm_microkernel_blocked_bt()

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

Definition at line 1058 of file gemm_microkernel.c.

{
    // Zero output first
    #pragma omp parallel for schedule(static)
    for (int i = 0; i < M; i++) {
        memset(&C[i * N], 0, N * sizeof(float));
    }
 
    const int mr = MR;
    const int nr = NR;
 
    #pragma omp parallel for schedule(dynamic) collapse(2)
    for (int m0 = 0; m0 < M; m0 += MC) {
        for (int n0 = 0; n0 < N; n0 += NC) {
            int mb = (m0 + MC <= M) ? MC : (M - m0);
            int nb = (n0 + NC <= N) ? NC : (N - n0);
 
            for (int k0 = 0; k0 < K; k0 += KC) {
                int kb = (k0 + KC <= K) ? KC : (K - k0);
                int first_k = (k0 == 0);
 
                for (int m1 = 0; m1 < mb; m1 += mr) {
                    int mr_actual = (m1 + mr <= mb) ? mr : (mb - m1);
 
                    for (int n1 = 0; n1 < nb; n1 += nr) {
                        int nr_actual = (n1 + nr <= nb) ? nr : (nb - n1);
 
                        const float *A_tile = &A[(m0 + m1) * K + k0];
                        const float *B_tile = &B[(n0 + n1) * K + k0];
                        float *C_tile = &C[(m0 + m1) * N + (n0 + n1)];
 
                        if (mr_actual == mr && nr_actual == nr) {
#if defined(__AVX512F__)
                            gemm_microkernel_6x32_bt_avx512(kb, A_tile, K, B_tile, K, C_tile, N, first_k);
#else
                            // Scalar fallback for B-transposed
                            for (int i = 0; i < mr; i++) {
                                for (int j = 0; j < nr; j++) {
                                    float sum = first_k ? 0.0f : C_tile[i * N + j];
                                    for (int kk = 0; kk < kb; kk++) {
                                        sum += A_tile[i * K + kk] * B_tile[j * K + kk];
                                    }
                                    C_tile[i * N + j] = sum;
                                }
                            }
#endif
                        } else {
                            // Edge case
                            for (int i = 0; i < mr_actual; i++) {
                                for (int j = 0; j < nr_actual; j++) {
                                    float sum = first_k ? 0.0f : C_tile[i * N + j];
                                    for (int kk = 0; kk < kb; kk++) {
                                        sum += A_tile[i * K + kk] * B_tile[j * K + kk];
                                    }
                                    C_tile[i * N + j] = sum;
                                }
                            }
                        }
                    }
                }
            }
        }
    }
}

References C, KC, MC, MR, NC, and NR.

Referenced by gemm_microkernel().

gemm_microkernel_edge()

static void gemm_microkernel_edge ( int  m, int  n, int  K, const float *  A, int  lda, const float *  B, int  ldb, float *  C, int  ldc, int  first_k  )

static

Definition at line 737 of file gemm_microkernel.c.

{
    for (int i = 0; i < m; i++) {
        for (int j = 0; j < n; j++) {
            float sum = first_k ? 0.0f : C[i * ldc + j];
            for (int k = 0; k < K; k++) {
                sum += A[i * lda + k] * B[k * ldb + j];
            }
            C[i * ldc + j] = sum;
        }
    }
}

References C.

Referenced by gemm_microkernel_blocked(), and gemm_microkernel_sequential().

gemm_microkernel_packed()

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

Definition at line 840 of file gemm_microkernel.c.

{
    // Use tile-parallel blocked version - scales better on many-core systems
    gemm_microkernel_blocked(A, B, C, M, N, K);
}

References C, and gemm_microkernel_blocked().

Referenced by gemm_microkernel().

gemm_microkernel_sequential()

static void gemm_microkernel_sequential ( const float *  A, const float *  B, float *  C, int  M, int  N, int  K  )

static

Definition at line 862 of file gemm_microkernel.c.

{
    // Zero output
    for (int i = 0; i < M; i++) {
        memset(&C[i * N], 0, N * sizeof(float));
    }
 
    const int mr = MR;
    const int nr = NR;
 
    // Block over K
    for (int k0 = 0; k0 < K; k0 += KC) {
        int kb = (k0 + KC <= K) ? KC : (K - k0);
        int first_k = (k0 == 0);
 
        // Loop over tiles
        for (int m0 = 0; m0 < M; m0 += mr) {
            int mr_actual = (m0 + mr <= M) ? mr : (M - m0);
 
            for (int n0 = 0; n0 < N; n0 += nr) {
                int nr_actual = (n0 + nr <= N) ? nr : (N - n0);
 
                const float *A_tile = &A[m0 * K + k0];
                const float *B_tile = &B[k0 * N + n0];
                float *C_tile = &C[m0 * N + n0];
 
                if (mr_actual == mr && nr_actual == nr) {
#if defined(__AVX512F__)
                    gemm_microkernel_6x32_avx512(kb, A_tile, K, B_tile, N, C_tile, N, first_k);
#elif defined(__FMA__)
                    gemm_microkernel_6x16_avx(kb, A_tile, K, B_tile, N, C_tile, N, first_k);
#elif defined(__AVX__)
                    gemm_microkernel_4x16_avx(kb, A_tile, K, B_tile, N, C_tile, N, first_k);
#else
                    gemm_microkernel_edge(mr_actual, nr_actual, kb, A_tile, K, B_tile, N, C_tile, N, first_k);
#endif
                } else {
                    gemm_microkernel_edge(mr_actual, nr_actual, kb, A_tile, K, B_tile, N, C_tile, N, first_k);
                }
            }
        }
    }
}

References C, gemm_microkernel_edge(), KC, MR, and NR.

Referenced by gemm_microkernel_blocked().

pack_a_panel()

static void pack_a_panel ( const float *  A, int  lda, float *  Ap, int  mc, int  kc, int  mr  )

static

Definition at line 761 of file gemm_microkernel.c.

{
    #pragma omp parallel for schedule(static) if(mc > 64)
    for (int i = 0; i < mc; i += mr) {
        int rows = (i + mr <= mc) ? mr : (mc - i);
        float *Ap_panel = &Ap[(i / mr) * mr * kc];
 
        for (int p = 0; p < rows; p++) {
            const float *A_row = &A[(i + p) * lda];
            float *Ap_row = &Ap_panel[p * kc];
 
            // Vectorized copy
            int k = 0;
#if defined(__AVX__)
            for (; k <= kc - 8; k += 8) {
                _mm256_storeu_ps(&Ap_row[k], _mm256_loadu_ps(&A_row[k]));
            }
#endif
            for (; k < kc; k++) {
                Ap_row[k] = A_row[k];
            }
        }
        // Zero pad if partial panel
        for (int p = rows; p < mr; p++) {
            memset(&Ap_panel[p * kc], 0, kc * sizeof(float));
        }
    }
}

pack_b_panel()

static void pack_b_panel ( const float *  B, int  ldb, float *  Bp, int  kc, int  nc, int  nr  )

static

Definition at line 795 of file gemm_microkernel.c.

{
    #pragma omp parallel for schedule(static) if(nc > 128)
    for (int j = 0; j < nc; j += nr) {
        int cols = (j + nr <= nc) ? nr : (nc - j);
        float *Bp_panel = &Bp[(j / nr) * nr * kc];
 
        for (int k = 0; k < kc; k++) {
            const float *B_row = &B[k * ldb + j];
            float *Bp_row = &Bp_panel[k * nr];
 
            // Copy cols and zero-pad
            int c = 0;
#if defined(__AVX512F__)
            for (; c <= cols - 16; c += 16) {
                _mm512_store_ps(&Bp_row[c], _mm512_loadu_ps(&B_row[c]));
            }
#elif defined(__AVX__)
            for (; c <= cols - 8; c += 8) {
                _mm256_store_ps(&Bp_row[c], _mm256_loadu_ps(&B_row[c]));
            }
#endif
            for (; c < cols; c++) {
                Bp_row[c] = B_row[c];
            }
            for (; c < nr; c++) {
                Bp_row[c] = 0.0f;
            }
        }
    }
}

Variable Documentation

g_threads_initialized

int g_threads_initialized = 0

static

Definition at line 912 of file gemm_microkernel.c.

Referenced by gemm_init_threads().