gemm_kernels_bf16.c File Reference

Optimized BF16 GEMM Kernels for AVX-512. More...

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "bf16_utils.h"
#include "ckernel_engine.h"
#include "ck_threadpool.h"

Go to the source code of this file.

Macros
#define	BLK_K   256
#define	BLK_M   64
#define	BLK_N   64

Functions
	__attribute__ ((unused))
static void	ck_bf16_convert_work (int ith, int nth, void *opaque)
static void	ck_bf16_round_work (int ith, int nth, void *opaque)
int	ck_gemm_bf16_amx_available (void)
static void	ck_gemm_bf16_amx_work (int ith, int nth, void *opaque)
int	ck_gemm_bf16_fp32out_amx_raw (const uint16_t *A, const uint16_t *B, float *C, int M, int N, int K, int accumulate)
static void	ck_gemm_bf16_native_work (int ith, int nth, void *opaque)
static void	ck_gemm_nt_bf16_accurate_storage_rows (int begin, int end, void *opaque)
static void	ck_gemm_nt_bf16_exact_rows (int begin, int end, void *opaque)
static void	ck_gemm_nt_bf16_storage_exact_rows (int begin, int end, void *opaque)
static void	ck_gemv_bf16_rows (int begin, int end, void *opaque)
static void	ck_gemv_bf16_storage_rows (int begin, int end, void *opaque)
static int	ck_min_i (int a, int b)
static void	ck_patch_projection_bf16_native_work (int ith, int nth, void *opaque)
static void	ck_pytorch_brgemm_fault (const char *message, int M, int N, int K)
static void	ck_pytorch_brgemm_workspace_sizes (int M, int N, int K, size_t *input_bytes, size_t *output_bytes, size_t *bias_bytes)
void	gemm_backward_bf16_mixed (const uint16_t *d_output, const uint16_t *input, const uint16_t *weight, float *d_input, float *d_weight, float *d_bias, int tokens, int in_dim, int out_dim)
void	gemm_bf16_fp32out (const uint16_t *A, const uint16_t *B, const float *bias, float *C, int M, int N, int K)
void	gemm_blocked_serial_bf16 (const uint16_t *A, const uint16_t *B, const uint16_t *bias, uint16_t *C, int M, int N, int K)
void	gemm_nn_bf16 (const uint16_t *A, const uint16_t *B, const uint16_t *bias, uint16_t *C, int M, int N, int K)
void	gemm_nt_bf16 (const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
void	gemm_nt_bf16_accurate_storage (const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
void	gemm_nt_bf16_accurate_storage_parallel_dispatch (const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
void	gemm_nt_bf16_accurate_storage_row_range (const float *A, const void *B, const float *bias, float *C, int M, int N, int K, int row_begin, int row_end)
void	gemm_nt_bf16_amx_bf16_storage (const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
void	gemm_nt_bf16_amx_bf16_storage_workspace (const float *A, const void *B, const float *bias, float *C, int M, int N, int K, uint16_t *a_bf16, size_t a_bf16_bytes)
void	gemm_nt_bf16_bf16_storage (const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
void	gemm_nt_bf16_bf16_storage_parallel_dispatch (const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
void	gemm_nt_bf16_bf16_storage_row_range (const float *A, const void *B, const float *bias, float *C, int M, int N, int K, int row_begin, int row_end)
void	gemm_nt_bf16_native_bf16_storage (const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
void	gemm_nt_bf16_parallel_dispatch (const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
void	gemm_nt_bf16_prefill_shape_safe_bf16_storage (const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
void	gemm_nt_bf16_prefill_shape_safe_bf16_storage_workspace (const float *A, const void *B, const float *bias, float *C, int M, int N, int K, uint16_t *a_bf16, size_t a_bf16_bytes)
void	gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage (const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
void	gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage_workspace (const float *A, const void *B, const float *bias, float *C, int M, int N, int K, uint16_t *input_bf16, size_t input_bf16_bytes, uint16_t *output_bf16, size_t output_bf16_bytes, uint16_t *bias_bf16, size_t bias_bf16_bytes)
void	gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage (const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
static void	gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_impl (const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
static void	gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_workspace_impl (const float *A, const void *B, const float *bias, float *C, int M, int N, int K, uint16_t *input_bf16, size_t input_bf16_bytes, uint16_t *output_bf16, size_t output_bf16_bytes, uint16_t *bias_bf16, size_t bias_bf16_bytes)
void	gemm_nt_bf16_row_range (const float *A, const void *B, const float *bias, float *C, int M, int N, int K, int row_begin, int row_end)
void	gemm_tn_bf16 (const uint16_t *A, const uint16_t *B, const uint16_t *bias, uint16_t *C, int M, int N, int K)
void	gemv_bf16 (float *y, const void *W, const float *x, int M, int K)
void	gemv_bf16_bf16_storage (float *y, const void *W, const float *x, int M, int K)
void	gemv_bf16_bf16_storage_parallel_dispatch (float *y, const void *W, const float *x, int M, int K)
static void	gemv_bf16_bf16_storage_row_range (float *y, const uint16_t *w, const float *x, int M, int K, int row_begin, int row_end)
void	gemv_bf16_parallel_dispatch (float *y, const void *W, const float *x, int M, int K)
static void	gemv_bf16_row_range (float *y, const uint16_t *w, const float *x, int M, int K, int row_begin, int row_end)
void	patch_projection_bf16_pytorch_onednn_conv3d_storage (const float *input, const void *weights, const float *bias, float *output, int batch, int out_channels, int in_channels, int temporal, int patch_h, int patch_w)
void	patch_projection_image_bf16_native_storage (const float *image, const void *weights_t0, const void *weights_t1, const float *bias, float *output, int channels, int image_h, int image_w, int patch_size, int out_channels, int merge_size)
void	patch_projection_image_bf16_pytorch_onednn_conv3d_storage (const float *image, const void *weights_t0, const void *weights_t1, const float *bias, float *output, int channels, int image_h, int image_w, int patch_size, int out_channels, int merge_size)

Detailed Description

Optimized BF16 GEMM Kernels for AVX-512.

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

Layout: A: [M x K] row-major (BF16) B: [N x K] row-major, stored as [out x in] (BF16) C: [M x N] row-major (BF16 or FP32)

Key optimizations:

1. AVX-512 BF16 instructions (VDPBF16PS) when available
2. Cache blocking for L1/L2 efficiency
3. Vectorized BF16<->FP32 conversion
4. CKE thread-pool row parallelization

Definition in file gemm_kernels_bf16.c.

Macro Definition Documentation

BLK_K

#define BLK_K   256

Definition at line 63 of file gemm_kernels_bf16.c.

BLK_M

#define BLK_M   64

Definition at line 61 of file gemm_kernels_bf16.c.

BLK_N

#define BLK_N   64

Definition at line 62 of file gemm_kernels_bf16.c.

Function Documentation

__attribute__()

__attribute__

(

(unused)

)

Definition at line 71 of file gemm_kernels_bf16.c.

{
    for (int i = 0; i < M; ++i) {
        for (int j = 0; j < N; ++j) {
            float sum = bias ? bf16_to_float(bias[j]) : 0.0f;
            const size_t a_row = (size_t)i * (size_t)K;
            const size_t b_row = (size_t)j * (size_t)K;
            for (int k = 0; k < K; ++k) {
                sum += bf16_to_float(A[a_row + k]) * bf16_to_float(B[b_row + k]);
            }
            C[(size_t)i * (size_t)N + j] = float_to_bf16(sum);
        }
    }
}

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

ck_bf16_convert_work()

static void ck_bf16_convert_work ( int  ith, int  nth, void *  opaque  )

static

Definition at line 1107 of file gemm_kernels_bf16.c.

{
    ck_bf16_convert_args_t *args = (ck_bf16_convert_args_t *)opaque;
    const size_t begin = args->count * (size_t)ith / (size_t)nth;
    const size_t end = args->count * (size_t)(ith + 1) / (size_t)nth;
    for (size_t i = begin; i < end; ++i) args->dst[i] = float_to_bf16(args->src[i]);
}

References end, and float_to_bf16().

Referenced by gemm_nt_bf16_amx_bf16_storage_workspace().

ck_bf16_round_work()

static void ck_bf16_round_work ( int  ith, int  nth, void *  opaque  )

static

Definition at line 1193 of file gemm_kernels_bf16.c.

{
    ck_bf16_round_args_t *args = (ck_bf16_round_args_t *)opaque;
    const size_t begin = args->count * (size_t)ith / (size_t)nth;
    const size_t end = args->count * (size_t)(ith + 1) / (size_t)nth;
    for (size_t i = begin; i < end; ++i) {
        args->values[i] = bf16_to_float(float_to_bf16(args->values[i]));
    }
}

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

Referenced by gemm_nt_bf16_amx_bf16_storage_workspace().

ck_gemm_bf16_amx_available()

int ck_gemm_bf16_amx_available

(

void

)

Definition at line 1203 of file gemm_kernels_bf16.c.

{
#if HAVE_AMX_BF16
    return ck_amx_request_xtile_data();
#else
    return 0;
#endif
}

Referenced by ck_attention_full_bf16_sdpa_amx_range(), gemm_nt_bf16_prefill_shape_safe_bf16_storage(), and gemm_nt_bf16_prefill_shape_safe_bf16_storage_workspace().

ck_gemm_bf16_amx_work()

static void ck_gemm_bf16_amx_work ( int  ith, int  nth, void *  opaque  )

static

Definition at line 1126 of file gemm_kernels_bf16.c.

{
#if HAVE_AMX_BF16
    ck_gemm_bf16_amx_args_t *args = (ck_gemm_bf16_amx_args_t *)opaque;
    if (!ck_amx_request_xtile_data()) {
        __atomic_store_n(&args->failed, 1, __ATOMIC_RELAXED);
        return;
    }
    ck_amx_config_bf16_16x16x32();
    uint16_t b_tile[16 * 32];
    const int mb = args->M / 16;
    const int nb = args->N / 16;
    const int m_groups = (mb + 3) / 4;
    const int jobs = m_groups * nb;
 
    for (int job = ith; job < jobs; job += nth) {
        const int m_group = job / nb;
        const int j = (job % nb) * 16;
        const int group_blocks = (mb - m_group * 4 < 4) ? mb - m_group * 4 : 4;
        _tile_zero(2);
        if (group_blocks > 1) _tile_zero(3);
        if (group_blocks > 2) _tile_zero(4);
        if (group_blocks > 3) _tile_zero(5);
        for (int k = 0; k < args->K; k += 32) {
            ck_pack_bf16_ktile_pairs_16x16(b_tile, args->B, args->K, j, k);
            _tile_loadd(1, b_tile, 32 * (int)sizeof(uint16_t));
            for (int g = 0; g < group_blocks; ++g) {
                const int i = (m_group * 4 + g) * 16;
                _tile_loadd(0, args->A + (size_t)i * args->K + k,
                            args->K * (int)sizeof(uint16_t));
                switch (g) {
                    case 0: _tile_dpbf16ps(2, 0, 1); break;
                    case 1: _tile_dpbf16ps(3, 0, 1); break;
                    case 2: _tile_dpbf16ps(4, 0, 1); break;
                    default: _tile_dpbf16ps(5, 0, 1); break;
                }
            }
        }
        for (int g = 0; g < group_blocks; ++g) {
            const int i = (m_group * 4 + g) * 16;
            float *tile_dst = args->C + (size_t)i * args->N + j;
            const int tile_stride = args->N * (int)sizeof(float);
            switch (g) {
                case 0: _tile_stored(2, tile_dst, tile_stride); break;
                case 1: _tile_stored(3, tile_dst, tile_stride); break;
                case 2: _tile_stored(4, tile_dst, tile_stride); break;
                default: _tile_stored(5, tile_dst, tile_stride); break;
            }
            if (args->bias) {
                for (int ii = 0; ii < 16; ++ii) {
                    float *row = args->C + (size_t)(i + ii) * args->N + j;
                    for (int jj = 0; jj < 16; ++jj) row[jj] += args->bias[j + jj];
                }
            }
        }
    }
    _tile_release();
#else
    (void)ith; (void)nth; (void)opaque;
#endif
}

Referenced by gemm_nt_bf16_amx_bf16_storage_workspace().

ck_gemm_bf16_fp32out_amx_raw()

int ck_gemm_bf16_fp32out_amx_raw	(	const uint16_t *	A,
		const uint16_t *	B,
		float *	C,
		int	M,
		int	N,
		int	K,
		int	accumulate
	)

Definition at line 1212 of file gemm_kernels_bf16.c.

{
#if HAVE_AMX_BF16
    if (!A || !B || !C || M <= 0 || N <= 0 || K <= 0 ||
        (M % 16) != 0 || (N % 16) != 0 || (K % 2) != 0 ||
        !ck_amx_request_xtile_data()) {
        return 0;
    }
    /* Match oneDNN BRGEMM: largest even divisor of K no greater than 32. */
    int k_block = K < 32 ? K : 32;
    while (k_block > 2 && K % k_block != 0) k_block -= 2;
    ck_amx_config_bf16_16x16_kblock(k_block);
    uint16_t b_tile[16 * 32];
    for (int i = 0; i < M; i += 16) {
        for (int j = 0; j < N; j += 16) {
            if (accumulate) {
                _tile_loadd(2, C + (size_t)i * (size_t)N + (size_t)j,
                            N * (int)sizeof(float));
            } else {
                _tile_zero(2);
            }
            for (int k = 0; k < K; k += k_block) {
                memset(b_tile, 0, sizeof(b_tile));
                for (int kp = 0; kp < k_block / 2; ++kp) {
                    const int k0 = k + kp * 2;
                    for (int nn = 0; nn < 16; ++nn) {
                        b_tile[(size_t)kp * 32u + (size_t)nn * 2u] =
                            B[(size_t)(j + nn) * (size_t)K + (size_t)k0];
                        b_tile[(size_t)kp * 32u + (size_t)nn * 2u + 1u] =
                            B[(size_t)(j + nn) * (size_t)K + (size_t)k0 + 1u];
                    }
                }
                _tile_loadd(0, A + (size_t)i * (size_t)K + (size_t)k,
                            K * (int)sizeof(uint16_t));
                _tile_loadd(1, b_tile, 32 * (int)sizeof(uint16_t));
                _tile_dpbf16ps(2, 0, 1);
            }
            _tile_stored(2, C + (size_t)i * (size_t)N + (size_t)j,
                         N * (int)sizeof(float));
        }
    }
    _tile_release();
    return 1;
#else
    (void)A; (void)B; (void)C; (void)M; (void)N; (void)K; (void)accumulate;
    return 0;
#endif
}

References C.

Referenced by ck_attention_full_bf16_sdpa_amx_range().

ck_gemm_bf16_native_work()

static void ck_gemm_bf16_native_work ( int  ith, int  nth, void *  opaque  )

static

Definition at line 976 of file gemm_kernels_bf16.c.

{
    ck_gemm_bf16_native_args_t *args = (ck_gemm_bf16_native_args_t *)opaque;
    const int N = args->N;
    const int K = args->K;
    enum { ROW_TILE = 4 };
    uint16_t *a_bf16 = (uint16_t *)alloca(
        (size_t)ROW_TILE * (size_t)K * sizeof(uint16_t));
 
    for (int row0 = ith * ROW_TILE; row0 < args->M; row0 += nth * ROW_TILE) {
        const int rows = args->M - row0 < ROW_TILE ? args->M - row0 : ROW_TILE;
        for (int r = 0; r < rows; ++r) {
            const float *src = args->A + (size_t)(row0 + r) * (size_t)K;
            uint16_t *ar = a_bf16 + (size_t)r * (size_t)K;
            for (int k = 0; k < K; ++k) ar[k] = float_to_bf16(src[k]);
        }
 
#if HAVE_NATIVE_BF16
        int j = 0;
        for (; j + 4 <= N; j += 4) {
            const uint16_t *b0 = args->B + (size_t)(j + 0) * K;
            const uint16_t *b1 = args->B + (size_t)(j + 1) * K;
            const uint16_t *b2 = args->B + (size_t)(j + 2) * K;
            const uint16_t *b3 = args->B + (size_t)(j + 3) * K;
            __m512 acc[ROW_TILE][4];
            for (int r = 0; r < rows; ++r) {
                for (int lane = 0; lane < 4; ++lane) acc[r][lane] = _mm512_setzero_ps();
            }
            int k = 0;
            for (; k <= K - 32; k += 32) {
                const __m512bh bv[4] = {
                    load_bf16x32(b0 + k), load_bf16x32(b1 + k),
                    load_bf16x32(b2 + k), load_bf16x32(b3 + k)
                };
                for (int r = 0; r < rows; ++r) {
                    const __m512bh av =
                        load_bf16x32(a_bf16 + (size_t)r * (size_t)K + k);
                    for (int lane = 0; lane < 4; ++lane) {
                        acc[r][lane] = _mm512_dpbf16_ps(acc[r][lane], av, bv[lane]);
                    }
                }
            }
            for (int r = 0; r < rows; ++r) {
                float sums[4];
                for (int lane = 0; lane < 4; ++lane) {
                    sums[lane] = _mm512_reduce_add_ps(acc[r][lane]);
                }
                const uint16_t *ar = a_bf16 + (size_t)r * (size_t)K;
                for (int tail = k; tail < K; ++tail) {
                    const float av = bf16_to_float(ar[tail]);
                    sums[0] += av * bf16_to_float(b0[tail]);
                    sums[1] += av * bf16_to_float(b1[tail]);
                    sums[2] += av * bf16_to_float(b2[tail]);
                    sums[3] += av * bf16_to_float(b3[tail]);
                }
                float *dst = args->C + (size_t)(row0 + r) * (size_t)N;
                for (int lane = 0; lane < 4; ++lane) {
                    if (args->bias) sums[lane] += args->bias[j + lane];
                    dst[j + lane] = bf16_to_float(float_to_bf16(sums[lane]));
                }
            }
        }
        for (; j < N; ++j) {
            const uint16_t *b = args->B + (size_t)j * K;
            __m512 acc[ROW_TILE];
            for (int r = 0; r < rows; ++r) acc[r] = _mm512_setzero_ps();
            int k = 0;
            for (; k <= K - 32; k += 32) {
                const __m512bh bv = load_bf16x32(b + k);
                for (int r = 0; r < rows; ++r) {
                    acc[r] = _mm512_dpbf16_ps(
                        acc[r], load_bf16x32(a_bf16 + (size_t)r * (size_t)K + k), bv);
                }
            }
            for (int r = 0; r < rows; ++r) {
                const uint16_t *ar = a_bf16 + (size_t)r * (size_t)K;
                float sum = _mm512_reduce_add_ps(acc[r]);
                for (int tail = k; tail < K; ++tail) {
                    sum += bf16_to_float(ar[tail]) * bf16_to_float(b[tail]);
                }
                if (args->bias) sum += args->bias[j];
                args->C[(size_t)(row0 + r) * (size_t)N + j] =
                    bf16_to_float(float_to_bf16(sum));
            }
        }
#else
        for (int r = 0; r < rows; ++r) {
            const uint16_t *ar = a_bf16 + (size_t)r * (size_t)K;
            float *dst = args->C + (size_t)(row0 + r) * (size_t)N;
            for (int j = 0; j < N; ++j) {
                const uint16_t *b = args->B + (size_t)j * K;
                float sum = args->bias ? args->bias[j] : 0.0f;
                for (int k = 0; k < K; ++k) {
                    sum += bf16_to_float(ar[k]) * bf16_to_float(b[k]);
                }
                dst[j] = bf16_to_float(float_to_bf16(sum));
            }
        }
#endif
    }
}

References bf16_to_float(), and float_to_bf16().

Referenced by gemm_nt_bf16_native_bf16_storage().

ck_gemm_nt_bf16_accurate_storage_rows()

static void ck_gemm_nt_bf16_accurate_storage_rows ( int  begin, int  end, void *  opaque  )

static

Definition at line 2119 of file gemm_kernels_bf16.c.

{
    const ck_gemm_nt_bf16_exact_args_t *args =
        (const ck_gemm_nt_bf16_exact_args_t *)opaque;
    gemm_nt_bf16_accurate_storage_row_range(
        args->A, args->B, args->bias, args->C,
        args->M, args->N, args->K, begin, end);
}

References end, and gemm_nt_bf16_accurate_storage_row_range().

Referenced by gemm_nt_bf16_accurate_storage_parallel_dispatch().

ck_gemm_nt_bf16_exact_rows()

static void ck_gemm_nt_bf16_exact_rows ( int  begin, int  end, void *  opaque  )

static

Definition at line 775 of file gemm_kernels_bf16.c.

{
    const ck_gemm_nt_bf16_exact_args_t *args =
        (const ck_gemm_nt_bf16_exact_args_t *)opaque;
    gemm_nt_bf16_row_range(
        args->A, args->B, args->bias, args->C,
        args->M, args->N, args->K, begin, end);
}

References end, and gemm_nt_bf16_row_range().

Referenced by gemm_nt_bf16_parallel_dispatch().

ck_gemm_nt_bf16_storage_exact_rows()

static void ck_gemm_nt_bf16_storage_exact_rows ( int  begin, int  end, void *  opaque  )

static

Definition at line 2110 of file gemm_kernels_bf16.c.

{
    const ck_gemm_nt_bf16_exact_args_t *args =
        (const ck_gemm_nt_bf16_exact_args_t *)opaque;
    gemm_nt_bf16_bf16_storage_row_range(
        args->A, args->B, args->bias, args->C,
        args->M, args->N, args->K, begin, end);
}

References end, and gemm_nt_bf16_bf16_storage_row_range().

Referenced by gemm_nt_bf16_bf16_storage_parallel_dispatch().

ck_gemv_bf16_rows()

static void ck_gemv_bf16_rows ( int  begin, int  end, void *  opaque  )

static

Definition at line 646 of file gemm_kernels_bf16.c.

{
    const ck_gemv_bf16_args_t *args = (const ck_gemv_bf16_args_t *)opaque;
    gemv_bf16_row_range(
        args->y, args->w, args->x, args->M, args->K, begin, end);
}

References end, and gemv_bf16_row_range().

Referenced by gemv_bf16_parallel_dispatch().

ck_gemv_bf16_storage_rows()

static void ck_gemv_bf16_storage_rows ( int  begin, int  end, void *  opaque  )

static

Definition at line 688 of file gemm_kernels_bf16.c.

{
    const ck_gemv_bf16_args_t *args = (const ck_gemv_bf16_args_t *)opaque;
    gemv_bf16_bf16_storage_row_range(
        args->y, args->w, args->x, args->M, args->K, begin, end);
}

References end, and gemv_bf16_bf16_storage_row_range().

Referenced by gemv_bf16_bf16_storage_parallel_dispatch().

ck_min_i()

static int ck_min_i ( int  a, int  b  )

inlinestatic

Definition at line 65 of file gemm_kernels_bf16.c.

{ return a < b ? a : b; }

ck_patch_projection_bf16_native_work()

static void ck_patch_projection_bf16_native_work ( int  ith, int  nth, void *  opaque  )

static

Definition at line 1922 of file gemm_kernels_bf16.c.

{
    ck_patch_projection_bf16_native_args_t *args =
        (ck_patch_projection_bf16_native_args_t *)opaque;
    const int begin = args->batch * ith / nth;
    const int end = args->batch * (ith + 1) / nth;
    const int patch_area = args->patch_size * args->patch_size;
    const int half_k = args->channels * patch_area;
    const int tiles_per_row = args->grid_w / args->merge_size;
    const int tile_area = args->merge_size * args->merge_size;
 
    for (int tok = begin; tok < end; ++tok) {
        const int tile = tok / tile_area;
        const int within = tok % tile_area;
        const int patch_y =
            (tile / tiles_per_row) * args->merge_size + within / args->merge_size;
        const int patch_x =
            (tile % tiles_per_row) * args->merge_size + within % args->merge_size;
        for (int n = 0; n < args->out_channels; ++n) {
            float sum = args->bias
                ? bf16_to_float(float_to_bf16(args->bias[n]))
                : 0.0f;
#if defined(__AVX512BF16__) && defined(__AVX512VL__)
            if (args->patch_size == 16) {
                __m256 acc = _mm256_setzero_ps();
                for (int c = 0; c < args->channels; ++c) {
                    for (int t = 0; t < 2; ++t) {
                        const uint16_t *weights = t == 0
                            ? args->weights_t0 : args->weights_t1;
                        const uint16_t *weight_plane = weights +
                            (size_t)n * (size_t)half_k +
                            (size_t)c * (size_t)patch_area;
                        for (int py = 0; py < 16; ++py) {
                            const float *src = args->image +
                                ((size_t)c * (size_t)args->image_h +
                                 (size_t)(patch_y * 16 + py)) *
                                    (size_t)args->image_w +
                                (size_t)(patch_x * 16);
                            const __m256bh image_bf16 = _mm256_cvtne2ps_pbh(
                                _mm256_loadu_ps(src + 8), _mm256_loadu_ps(src));
                            const __m256bh weight_bf16 = (__m256bh)_mm256_loadu_si256(
                                (const __m256i *)(weight_plane + (size_t)py * 16u));
                            acc = _mm256_dpbf16_ps(acc, image_bf16, weight_bf16);
                        }
                    }
                }
                float lanes[8];
                _mm256_storeu_ps(lanes, acc);
                const float sum01 = lanes[0] + lanes[1];
                const float sum23 = lanes[2] + lanes[3];
                const float sum45 = lanes[4] + lanes[5];
                const float sum67 = lanes[6] + lanes[7];
                sum += (sum01 + sum23) + (sum45 + sum67);
            } else
#endif
            {
                for (int c = 0; c < args->channels; ++c) {
                    for (int t = 0; t < 2; ++t) {
                        const uint16_t *weights = t == 0
                            ? args->weights_t0 : args->weights_t1;
                        const uint16_t *weight_plane = weights +
                            (size_t)n * (size_t)half_k +
                            (size_t)c * (size_t)patch_area;
                        for (int py = 0; py < args->patch_size; ++py) {
                            const float *src = args->image +
                                ((size_t)c * (size_t)args->image_h +
                                 (size_t)(patch_y * args->patch_size + py)) *
                                    (size_t)args->image_w +
                                (size_t)(patch_x * args->patch_size);
                            for (int px = 0; px < args->patch_size; ++px) {
                                const float value = bf16_to_float(float_to_bf16(src[px]));
                                const float weight = bf16_to_float(
                                    weight_plane[(size_t)py *
                                        (size_t)args->patch_size + (size_t)px]);
                                sum += value * weight;
                            }
                        }
                    }
                }
            }
            args->output[(size_t)tok * (size_t)args->out_channels + (size_t)n] =
                bf16_to_float(float_to_bf16(sum));
        }
    }
}

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

Referenced by patch_projection_image_bf16_native_storage().

ck_pytorch_brgemm_fault()

static void ck_pytorch_brgemm_fault ( const char *  message, int  M, int  N, int  K  )

static

Definition at line 31 of file gemm_kernels_bf16.c.

{
    fprintf(stderr, "HARD KERNEL CONTRACT FAULT: PyTorch oneDNN BF16 BRGEMM %s "
                    "(M=%d N=%d K=%d)\n", message, M, N, K);
    abort();
}

Referenced by ck_pytorch_brgemm_workspace_sizes(), gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_impl(), gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_workspace_impl(), patch_projection_bf16_pytorch_onednn_conv3d_storage(), and patch_projection_image_bf16_pytorch_onednn_conv3d_storage().

ck_pytorch_brgemm_workspace_sizes()

static void ck_pytorch_brgemm_workspace_sizes ( int  M, int  N, int  K, size_t *  input_bytes, size_t *  output_bytes, size_t *  bias_bytes  )

static

Definition at line 1430 of file gemm_kernels_bf16.c.

{
    if (M <= 0 || N <= 0 || K <= 0 ||
        (size_t)M > SIZE_MAX / (size_t)K ||
        (size_t)M > SIZE_MAX / (size_t)N) {
        ck_pytorch_brgemm_fault("workspace size overflow", M, N, K);
    }
    const size_t input_count = (size_t)M * (size_t)K;
    const size_t output_count = (size_t)M * (size_t)N;
    if (input_count > SIZE_MAX / sizeof(uint16_t) ||
        output_count > SIZE_MAX / sizeof(uint16_t) ||
        (size_t)N > SIZE_MAX / sizeof(uint16_t)) {
        ck_pytorch_brgemm_fault("workspace byte size overflow", M, N, K);
    }
    *input_bytes = input_count * sizeof(uint16_t);
    *output_bytes = output_count * sizeof(uint16_t);
    *bias_bytes = (size_t)N * sizeof(uint16_t);
}

References ck_pytorch_brgemm_fault().

Referenced by gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_impl(), and gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_workspace_impl().

gemm_backward_bf16_mixed()

void gemm_backward_bf16_mixed	(	const uint16_t *	d_output,
		const uint16_t *	input,
		const uint16_t *	weight,
		float *	d_input,
		float *	d_weight,
		float *	d_bias,
		int	tokens,
		int	in_dim,
		int	out_dim
	)

Definition at line 2194 of file gemm_kernels_bf16.c.

{
    if (!d_output || !input || !weight || tokens <= 0 || in_dim <= 0 || out_dim <= 0) {
        return;
    }
 
    if (d_input) {
        for (int t = 0; t < tokens; ++t) {
            for (int i = 0; i < in_dim; ++i) {
                float sum = 0.0f;
                for (int o = 0; o < out_dim; ++o) {
                    const float dy = bf16_to_float(d_output[(size_t)t * (size_t)out_dim + (size_t)o]);
                    const float w = bf16_to_float(weight[(size_t)o * (size_t)in_dim + (size_t)i]);
                    sum += dy * w;
                }
                d_input[(size_t)t * (size_t)in_dim + (size_t)i] = sum;
            }
        }
    }
 
    if (d_weight) {
        for (int o = 0; o < out_dim; ++o) {
            for (int i = 0; i < in_dim; ++i) {
                float sum = 0.0f;
                for (int t = 0; t < tokens; ++t) {
                    const float dy = bf16_to_float(d_output[(size_t)t * (size_t)out_dim + (size_t)o]);
                    const float x = bf16_to_float(input[(size_t)t * (size_t)in_dim + (size_t)i]);
                    sum += dy * x;
                }
                d_weight[(size_t)o * (size_t)in_dim + (size_t)i] = sum;
            }
        }
    }
 
    if (d_bias) {
        for (int o = 0; o < out_dim; ++o) {
            float sum = 0.0f;
            for (int t = 0; t < tokens; ++t) {
                sum += bf16_to_float(d_output[(size_t)t * (size_t)out_dim + (size_t)o]);
            }
            d_bias[o] = sum;
        }
    }
}

References bf16_to_float().

gemm_bf16_fp32out()

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

Definition at line 460 of file gemm_kernels_bf16.c.

{
    if (!A || !B || !C || M <= 0 || N <= 0 || K <= 0) {
        return;
    }
 
#if HAVE_NATIVE_BF16
#if HAVE_AMX_BF16
    const char *amx_env = getenv("CK_BF16_AMX");
    if (amx_env && amx_env[0] == '1' &&
        (M % 16) == 0 && (N % 16) == 0 && (K % 32) == 0 &&
        M >= 16 && N >= 16 && K >= 32 && ck_amx_request_xtile_data()) {
        gemm_bf16_fp32out_amx(A, B, bias, C, M, N, K);
        return;
    }
#endif
 
    #pragma omp parallel for schedule(dynamic)
    for (int i = 0; i < M; ++i) {
        const uint16_t *a_row = A + (size_t)i * K;
        int j = 0;
 
        for (; j + 4 <= N; j += 4) {
            const uint16_t *b0 = B + (size_t)(j + 0) * K;
            const uint16_t *b1 = B + (size_t)(j + 1) * K;
            const uint16_t *b2 = B + (size_t)(j + 2) * K;
            const uint16_t *b3 = B + (size_t)(j + 3) * K;
            __m512 acc0 = _mm512_setzero_ps();
            __m512 acc1 = _mm512_setzero_ps();
            __m512 acc2 = _mm512_setzero_ps();
            __m512 acc3 = _mm512_setzero_ps();
 
            int k = 0;
            for (; k <= K - 32; k += 32) {
                const __m512bh a_vec = load_bf16x32(a_row + k);
                acc0 = _mm512_dpbf16_ps(acc0, a_vec, load_bf16x32(b0 + k));
                acc1 = _mm512_dpbf16_ps(acc1, a_vec, load_bf16x32(b1 + k));
                acc2 = _mm512_dpbf16_ps(acc2, a_vec, load_bf16x32(b2 + k));
                acc3 = _mm512_dpbf16_ps(acc3, a_vec, load_bf16x32(b3 + k));
            }
 
            float s0 = _mm512_reduce_add_ps(acc0);
            float s1 = _mm512_reduce_add_ps(acc1);
            float s2 = _mm512_reduce_add_ps(acc2);
            float s3 = _mm512_reduce_add_ps(acc3);
            for (; k < K; ++k) {
                const float a = bf16_to_float(a_row[k]);
                s0 += a * bf16_to_float(b0[k]);
                s1 += a * bf16_to_float(b1[k]);
                s2 += a * bf16_to_float(b2[k]);
                s3 += a * bf16_to_float(b3[k]);
            }
            if (bias) {
                s0 += bias[j + 0];
                s1 += bias[j + 1];
                s2 += bias[j + 2];
                s3 += bias[j + 3];
            }
            C[(size_t)i * N + (j + 0)] = s0;
            C[(size_t)i * N + (j + 1)] = s1;
            C[(size_t)i * N + (j + 2)] = s2;
            C[(size_t)i * N + (j + 3)] = s3;
        }
 
        for (; j < N; ++j) {
            const uint16_t *b_row = B + (size_t)j * K;
            __m512 sum_vec = _mm512_setzero_ps();
 
            int k = 0;
            for (; k <= K - 32; k += 32) {
                const __m512bh a_vec = load_bf16x32(a_row + k);
                const __m512bh b_vec = load_bf16x32(b_row + k);
                sum_vec = _mm512_dpbf16_ps(sum_vec, a_vec, b_vec);
            }
 
            float sum = _mm512_reduce_add_ps(sum_vec);
            for (; k < K; ++k) {
                sum += bf16_to_float(a_row[k]) * bf16_to_float(b_row[k]);
            }
            if (bias) {
                sum += bias[j];
            }
            C[(size_t)i * N + j] = sum;
        }
    }
#elif defined(__AVX512F__)
    #pragma omp parallel for schedule(dynamic)
    for (int i = 0; i < M; ++i) {
        const uint16_t *a_row = A + (size_t)i * K;
 
        for (int j = 0; j < N; ++j) {
            const uint16_t *b_row = B + (size_t)j * K;
 
            __m512 sum_vec = _mm512_setzero_ps();
 
            int k = 0;
            for (; k <= K - 16; k += 16) {
                __m256i a_bf16 = _mm256_loadu_si256((const __m256i *)(a_row + k));
                __m256i b_bf16 = _mm256_loadu_si256((const __m256i *)(b_row + k));
                sum_vec = bf16_dot16(a_bf16, b_bf16, sum_vec);
            }
 
            float sum = _mm512_reduce_add_ps(sum_vec);
 
            for (; k < K; ++k) {
                sum += bf16_to_float(a_row[k]) * bf16_to_float(b_row[k]);
            }
 
            if (bias) {
                sum += bias[j];
            }
 
            C[(size_t)i * N + j] = sum;
        }
    }
#else
    for (int i = 0; i < M; ++i) {
        for (int j = 0; j < N; ++j) {
            float sum = bias ? bias[j] : 0.0f;
            for (int k = 0; k < K; ++k) {
                sum += bf16_to_float(A[(size_t)i * K + k]) *
                       bf16_to_float(B[(size_t)j * K + k]);
            }
            C[(size_t)i * N + j] = sum;
        }
    }
#endif
}

References bf16_to_float(), and C.

Referenced by mlp_token_parallel_bf16(), and mlp_token_parallel_bf16_fp32act().

gemm_blocked_serial_bf16()

void gemm_blocked_serial_bf16	(	const uint16_t *	A,
		const uint16_t *	B,
		const uint16_t *	bias,
		uint16_t *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 431 of file gemm_kernels_bf16.c.

{
    if (!A || !B || !C || M <= 0 || N <= 0 || K <= 0) {
        return;
    }
 
#if HAVE_NATIVE_BF16
    /* Native BF16 instructions available (Ice Lake / Sapphire Rapids+) */
    gemm_bf16_native(A, B, bias, C, M, N, K);
#elif defined(__AVX512F__)
    /* Use AVX-512F with software BF16 conversion */
    if (M * N > 4096) {
        gemm_bf16_blocked_avx512(A, B, bias, C, M, N, K);
    } else {
        gemm_bf16_avx512(A, B, bias, C, M, N, K);
    }
#else
    /* Scalar fallback */
    gemm_bf16_scalar(A, B, bias, C, M, N, K);
#endif
}

References C.

gemm_nn_bf16()

void gemm_nn_bf16	(	const uint16_t *	A,
		const uint16_t *	B,
		const uint16_t *	bias,
		uint16_t *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 815 of file gemm_kernels_bf16.c.

{
    if (!A || !B || !C || M <= 0 || N <= 0 || K <= 0) {
        return;
    }
 
#if defined(__AVX512F__)
    #pragma omp parallel for
    for (int i = 0; i < M; ++i) {
        /* Initialize row with bias */
        int j = 0;
        for (; j <= N - 16; j += 16) {
            __m512 b_vec = bias ? bf16x16_to_fp32(_mm256_loadu_si256((const __m256i *)(bias + j)))
                                : _mm512_setzero_ps();
            __m256i out = fp32x16_to_bf16(b_vec);
            _mm256_storeu_si256((__m256i *)(C + (size_t)i * N + j), out);
        }
        for (; j < N; ++j) {
            float b = bias ? bf16_to_float(bias[j]) : 0.0f;
            C[(size_t)i * N + j] = float_to_bf16(b);
        }
 
        /* Accumulate: C[i,:] += A[i,k] * B[k,:] */
        for (int k = 0; k < K; ++k) {
            float a_val = bf16_to_float(A[(size_t)i * K + k]);
            __m512 a_broadcast = _mm512_set1_ps(a_val);
 
            j = 0;
            for (; j <= N - 16; j += 16) {
                __m256i b_bf16 = _mm256_loadu_si256((const __m256i *)(B + (size_t)k * N + j));
                __m512 b_fp32 = bf16x16_to_fp32(b_bf16);
 
                __m256i c_bf16 = _mm256_loadu_si256((const __m256i *)(C + (size_t)i * N + j));
                __m512 c_fp32 = bf16x16_to_fp32(c_bf16);
 
                c_fp32 = _mm512_fmadd_ps(a_broadcast, b_fp32, c_fp32);
 
                __m256i c_out = fp32x16_to_bf16(c_fp32);
                _mm256_storeu_si256((__m256i *)(C + (size_t)i * N + j), c_out);
            }
            for (; j < N; ++j) {
                float c_val = bf16_to_float(C[(size_t)i * N + j]);
                c_val += a_val * bf16_to_float(B[(size_t)k * N + j]);
                C[(size_t)i * N + j] = float_to_bf16(c_val);
            }
        }
    }
#else
    /* Scalar fallback */
    for (int i = 0; i < M; ++i) {
        for (int j = 0; j < N; ++j) {
            float sum = bias ? bf16_to_float(bias[j]) : 0.0f;
            for (int k = 0; k < K; ++k) {
                sum += bf16_to_float(A[(size_t)i * K + k]) *
                       bf16_to_float(B[(size_t)k * N + j]);
            }
            C[(size_t)i * N + j] = float_to_bf16(sum);
        }
    }
#endif
}

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

gemm_nt_bf16()

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

Definition at line 755 of file gemm_kernels_bf16.c.

{
    if (M <= 0) return;
    gemm_nt_bf16_row_range(A, B, bias, C, M, N, K, 0, M);
}

References C, and gemm_nt_bf16_row_range().

Referenced by gemm_nt_bf16_parallel_dispatch().

gemm_nt_bf16_accurate_storage()

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

Definition at line 2185 of file gemm_kernels_bf16.c.

{
    gemm_nt_bf16_accurate_storage_parallel_dispatch(A, B, bias, C, M, N, K);
}

References C, and gemm_nt_bf16_accurate_storage_parallel_dispatch().

gemm_nt_bf16_accurate_storage_parallel_dispatch()

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

Definition at line 2161 of file gemm_kernels_bf16.c.

{
    ck_threadpool_t *pool = ck_threadpool_global();
    if (!pool || ck_threadpool_n_threads(pool) <= 1 || M < 2 ||
        (size_t)M * (size_t)N <= 4096) {
        gemm_nt_bf16_accurate_storage_row_range(A, B, bias, C, M, N, K, 0, M);
        return;
    }
 
    ck_gemm_nt_bf16_exact_args_t args = {
        .A = A, .B = B, .bias = bias, .C = C, .M = M, .N = N, .K = K,
    };
    int active = ck_threadpool_n_threads(pool);
    if (active > M) active = M;
    int grain = M / (active * 4);
    if (grain < 1) grain = 1;
    ck_threadpool_parallel_for_n(
        pool, active, 0, M, grain, ck_gemm_nt_bf16_accurate_storage_rows, &args);
}

References C, ck_gemm_nt_bf16_accurate_storage_rows(), ck_threadpool_global(), ck_threadpool_n_threads(), ck_threadpool_parallel_for_n(), and gemm_nt_bf16_accurate_storage_row_range().

Referenced by gemm_nt_bf16_accurate_storage().

gemm_nt_bf16_accurate_storage_row_range()

void gemm_nt_bf16_accurate_storage_row_range	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K,
		int	row_begin,
		int	row_end
	)

Definition at line 2070 of file gemm_kernels_bf16.c.

{
    const uint16_t *w = (const uint16_t *)B;
    if (!A || !w || !C || M <= 0 || N <= 0 || K <= 0 ||
        row_begin < 0 || row_begin >= row_end || row_end > M) {
        return;
    }
    for (int row = row_begin; row < row_end; ++row) {
        const float *a_row = A + (size_t)row * (size_t)K;
        float *dst = C + (size_t)row * (size_t)N;
        for (int col = 0; col < N; ++col) {
            const uint16_t *w_row = w + (size_t)col * (size_t)K;
            /* BF16 products are exact in double. A single FP32 accumulator can
             * round a long dot across a BF16 midpoint before output storage;
             * that changes downstream routing. Keep the accumulation accurate
             * until the declared FP32 bias/output conversion boundary. */
            double even = 0.0;
            double odd = 0.0;
            int k = 0;
            for (; k + 1 < K; k += 2) {
                even += (double)bf16_to_float(float_to_bf16(a_row[k])) *
                        (double)bf16_to_float(w_row[k]);
                odd += (double)bf16_to_float(float_to_bf16(a_row[k + 1])) *
                       (double)bf16_to_float(w_row[k + 1]);
            }
            if (k < K) {
                even += (double)bf16_to_float(float_to_bf16(a_row[k])) *
                        (double)bf16_to_float(w_row[k]);
            }
            const float sum = (float)((even + odd) + (bias ? bias[col] : 0.0f));
            dst[col] = bf16_to_float(float_to_bf16(sum));
        }
    }
}

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

Referenced by ck_gemm_nt_bf16_accurate_storage_rows(), and gemm_nt_bf16_accurate_storage_parallel_dispatch().

gemm_nt_bf16_amx_bf16_storage()

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

Definition at line 1318 of file gemm_kernels_bf16.c.

{
    size_t input_count = 0;
    if (M > 0 && K > 0 && (size_t)M <= SIZE_MAX / (size_t)K) {
        input_count = (size_t)M * (size_t)K;
    }
    uint16_t *workspace = input_count > 0 && input_count <= SIZE_MAX / sizeof(uint16_t)
        ? (uint16_t *)malloc(input_count * sizeof(uint16_t))
        : NULL;
    if (!workspace) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: AMX BF16 compatibility workspace allocation failed\n");
        abort();
    }
    gemm_nt_bf16_amx_bf16_storage_workspace(
        A, B, bias, C, M, N, K, workspace, input_count * sizeof(uint16_t));
    free(workspace);
}

References C, and gemm_nt_bf16_amx_bf16_storage_workspace().

Referenced by gemm_nt_bf16_prefill_shape_safe_bf16_storage().

gemm_nt_bf16_amx_bf16_storage_workspace()

void gemm_nt_bf16_amx_bf16_storage_workspace	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K,
		uint16_t *	a_bf16,
		size_t	a_bf16_bytes
	)

Definition at line 1265 of file gemm_kernels_bf16.c.

{
#if HAVE_AMX_BF16
    if (!A || !B || !C || M < 16 || N < 16 || K < 32 ||
        (M % 16) != 0 || (N % 16) != 0 || (K % 32) != 0) {
        fprintf(stderr,
            "HARD KERNEL CONTRACT FAULT: AMX BF16 GEMM requires non-null buffers "
            "and M%%16=N%%16=K%%32=0 (M=%d N=%d K=%d)\n",
            M, N, K);
        abort();
    }
    const size_t input_count = (size_t)M * K;
    if (!a_bf16 || input_count > SIZE_MAX / sizeof(uint16_t) ||
        a_bf16_bytes < input_count * sizeof(uint16_t)) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: AMX BF16 activation workspace is too small\n");
        abort();
    }
    ck_threadpool_t *pool = ck_threadpool_global();
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    if (active > 24) active = 24;
    ck_bf16_convert_args_t convert = {.src=A, .dst=a_bf16, .count=input_count};
    if (pool && active > 1) ck_threadpool_dispatch_n(pool, active, ck_bf16_convert_work, &convert);
    else ck_bf16_convert_work(0, 1, &convert);
    ck_gemm_bf16_amx_args_t gemm = {
        .A=a_bf16, .B=(const uint16_t *)B, .bias=bias, .C=C,
        .M=M, .N=N, .K=K, .failed=0
    };
    if (pool && active > 1) ck_threadpool_dispatch_n(pool, active, ck_gemm_bf16_amx_work, &gemm);
    else ck_gemm_bf16_amx_work(0, 1, &gemm);
    if (__atomic_load_n(&gemm.failed, __ATOMIC_RELAXED)) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: AMX tile permission request failed\n");
        abort();
    }
    ck_bf16_round_args_t round = {.values=C, .count=(size_t)M * N};
    if (pool && active > 1) ck_threadpool_dispatch_n(pool, active, ck_bf16_round_work, &round);
    else ck_bf16_round_work(0, 1, &round);
    return;
#else
    (void)A; (void)B; (void)bias; (void)C; (void)M; (void)N; (void)K;
    (void)a_bf16; (void)a_bf16_bytes;
    fprintf(stderr,
            "HARD KERNEL CONTRACT FAULT: gemm_nt_bf16_amx_bf16_storage was selected "
            "without AMX BF16 support\n");
    abort();
#endif
}

References C, ck_bf16_convert_work(), ck_bf16_round_work(), ck_gemm_bf16_amx_work(), ck_threadpool_dispatch_n(), ck_threadpool_global(), and ck_threadpool_n_threads().

Referenced by gemm_nt_bf16_amx_bf16_storage(), and gemm_nt_bf16_prefill_shape_safe_bf16_storage_workspace().

gemm_nt_bf16_bf16_storage()

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

Definition at line 2152 of file gemm_kernels_bf16.c.

{
    gemm_nt_bf16_bf16_storage_parallel_dispatch(A, B, bias, C, M, N, K);
}

References C, and gemm_nt_bf16_bf16_storage_parallel_dispatch().

gemm_nt_bf16_bf16_storage_parallel_dispatch()

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

Definition at line 2128 of file gemm_kernels_bf16.c.

{
    ck_threadpool_t *pool = ck_threadpool_global();
    if (!pool || ck_threadpool_n_threads(pool) <= 1 || M < 2 ||
        (size_t)M * (size_t)N <= 4096) {
        gemm_nt_bf16_bf16_storage_row_range(A, B, bias, C, M, N, K, 0, M);
        return;
    }
 
    ck_gemm_nt_bf16_exact_args_t args = {
        .A = A, .B = B, .bias = bias, .C = C, .M = M, .N = N, .K = K,
    };
    int active = ck_threadpool_n_threads(pool);
    if (active > M) active = M;
    int grain = M / (active * 4);
    if (grain < 1) grain = 1;
    ck_threadpool_parallel_for_n(
        pool, active, 0, M, grain, ck_gemm_nt_bf16_storage_exact_rows, &args);
}

References C, ck_gemm_nt_bf16_storage_exact_rows(), ck_threadpool_global(), ck_threadpool_n_threads(), ck_threadpool_parallel_for_n(), and gemm_nt_bf16_bf16_storage_row_range().

Referenced by gemm_nt_bf16_bf16_storage().

gemm_nt_bf16_bf16_storage_row_range()

void gemm_nt_bf16_bf16_storage_row_range	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K,
		int	row_begin,
		int	row_end
	)

Definition at line 2054 of file gemm_kernels_bf16.c.

{
    gemm_nt_bf16_row_range(A, B, bias, C, M, N, K, row_begin, row_end);
    for (int row = row_begin; row < row_end; ++row) {
        float *dst = C + (size_t)row * (size_t)N;
        for (int col = 0; col < N; ++col) {
            dst[col] = bf16_to_float(float_to_bf16(dst[col]));
        }
    }
}

References bf16_to_float(), C, float_to_bf16(), and gemm_nt_bf16_row_range().

Referenced by ck_gemm_nt_bf16_storage_exact_rows(), and gemm_nt_bf16_bf16_storage_parallel_dispatch().

gemm_nt_bf16_native_bf16_storage()

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

Definition at line 1078 of file gemm_kernels_bf16.c.

{
    const uint16_t *weights = (const uint16_t *)B;
    if (!A || !weights || !C || M <= 0 || N <= 0 || K <= 0) return;
 
    ck_gemm_bf16_native_args_t args = {
        .A = A, .B = weights, .bias = bias, .C = C, .M = M, .N = N, .K = K
    };
    ck_threadpool_t *pool = ck_threadpool_global();
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    if (active > M) active = M;
    if (active > 24) active = 24;
    if (!pool || active <= 1 || (size_t)M * (size_t)N <= 4096) {
        ck_gemm_bf16_native_work(0, 1, &args);
        return;
    }
    ck_threadpool_dispatch_n(pool, active, ck_gemm_bf16_native_work, &args);
}

References C, ck_gemm_bf16_native_work(), ck_threadpool_dispatch_n(), ck_threadpool_global(), and ck_threadpool_n_threads().

Referenced by gemm_nt_bf16_prefill_shape_safe_bf16_storage(), and gemm_nt_bf16_prefill_shape_safe_bf16_storage_workspace().

gemm_nt_bf16_parallel_dispatch()

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

Definition at line 784 of file gemm_kernels_bf16.c.

{
    ck_threadpool_t *pool = ck_threadpool_global();
    const char *disabled = getenv("CK_DISABLE_BF16_GEMM_PARALLEL_PREFILL");
    if ((disabled && disabled[0] && strcmp(disabled, "0") != 0) ||
        !pool || ck_threadpool_n_threads(pool) <= 1 || M < 2 ||
        (size_t)M * (size_t)N <= 4096) {
        gemm_nt_bf16(A, B, bias, C, M, N, K);
        return;
    }
 
    ck_gemm_nt_bf16_exact_args_t args = {
        .A = A, .B = B, .bias = bias, .C = C, .M = M, .N = N, .K = K,
    };
    int active = ck_threadpool_n_threads(pool);
    if (active > M) active = M;
    int grain = M / (active * 4);
    if (grain < 1) grain = 1;
    ck_threadpool_parallel_for_n(
        pool, active, 0, M, grain, ck_gemm_nt_bf16_exact_rows, &args);
}

References C, ck_gemm_nt_bf16_exact_rows(), ck_threadpool_global(), ck_threadpool_n_threads(), ck_threadpool_parallel_for_n(), and gemm_nt_bf16().

gemm_nt_bf16_prefill_shape_safe_bf16_storage()

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

Definition at line 1354 of file gemm_kernels_bf16.c.

{
    const int amx_shape = M >= 16 && N >= 16 && K >= 32 &&
                          (M % 16) == 0 && (N % 16) == 0 && (K % 32) == 0;
    if (amx_shape && ck_gemm_bf16_amx_available()) {
        gemm_nt_bf16_amx_bf16_storage(A, B, bias, C, M, N, K);
        return;
    }
    gemm_nt_bf16_native_bf16_storage(A, B, bias, C, M, N, K);
}

References C, ck_gemm_bf16_amx_available(), gemm_nt_bf16_amx_bf16_storage(), and gemm_nt_bf16_native_bf16_storage().

gemm_nt_bf16_prefill_shape_safe_bf16_storage_workspace()

void gemm_nt_bf16_prefill_shape_safe_bf16_storage_workspace	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K,
		uint16_t *	a_bf16,
		size_t	a_bf16_bytes
	)

Definition at line 1340 of file gemm_kernels_bf16.c.

{
    const int amx_shape = M >= 16 && N >= 16 && K >= 32 &&
                          (M % 16) == 0 && (N % 16) == 0 && (K % 32) == 0;
    if (amx_shape && ck_gemm_bf16_amx_available()) {
        gemm_nt_bf16_amx_bf16_storage_workspace(
            A, B, bias, C, M, N, K, a_bf16, a_bf16_bytes);
        return;
    }
    gemm_nt_bf16_native_bf16_storage(A, B, bias, C, M, N, K);
}

References C, ck_gemm_bf16_amx_available(), gemm_nt_bf16_amx_bf16_storage_workspace(), and gemm_nt_bf16_native_bf16_storage().

gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage()

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

Definition at line 1601 of file gemm_kernels_bf16.c.

{
#ifdef USE_ONEDNN
    ck_pytorch_brgemm_require_version(
        3, 12, 0, "80afa71049cd69a3df32adcccb623b12cd7baa22",
        "gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage", M, N, K);
#endif
    gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_impl(
        A, B, bias, C, M, N, K);
}

References C, and gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_impl().

gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage_workspace()

void gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage_workspace	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K,
		uint16_t *	input_bf16,
		size_t	input_bf16_bytes,
		uint16_t *	output_bf16,
		size_t	output_bf16_bytes,
		uint16_t *	bias_bf16,
		size_t	bias_bf16_bytes
	)

Definition at line 1614 of file gemm_kernels_bf16.c.

{
#ifdef USE_ONEDNN
    ck_pytorch_brgemm_require_version(
        3, 12, 0, "80afa71049cd69a3df32adcccb623b12cd7baa22",
        "gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage_workspace",
        M, N, K);
#endif
    gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_workspace_impl(
        A, B, bias, C, M, N, K,
        input_bf16, input_bf16_bytes, output_bf16, output_bf16_bytes,
        bias_bf16, bias_bf16_bytes);
}

References C, and gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_workspace_impl().

Referenced by ck_attention_muse_eager_bf16_storage_impl().

gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage()

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

Definition at line 1586 of file gemm_kernels_bf16.c.

{
#ifdef USE_ONEDNN
    ck_pytorch_brgemm_require_version(
        3, 7, 1, "8d263e693366ef8db40acc569cc7d8edf644556d",
        "gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage", M, N, K);
#endif
    gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_impl(
        A, B, bias, C, M, N, K);
}

References C, and gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_impl().

gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_impl()

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

static

Definition at line 1559 of file gemm_kernels_bf16.c.

{
    size_t input_bytes = 0;
    size_t output_bytes = 0;
    size_t bias_bytes = 0;
    ck_pytorch_brgemm_workspace_sizes(
        M, N, K, &input_bytes, &output_bytes, &bias_bytes);
    uint16_t *input_bf16 = (uint16_t *)malloc(input_bytes);
    uint16_t *output_bf16 = (uint16_t *)malloc(output_bytes);
    uint16_t *bias_bf16 = bias ? (uint16_t *)malloc(bias_bytes) : NULL;
    if (!input_bf16 || !output_bf16 || (bias && !bias_bf16)) {
        free(bias_bf16);
        free(output_bf16);
        free(input_bf16);
        ck_pytorch_brgemm_fault("workspace allocation failed", M, N, K);
    }
    gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_workspace_impl(
        A, B, bias, C, M, N, K,
        input_bf16, input_bytes, output_bf16, output_bytes,
        bias_bf16, bias_bytes);
    free(bias_bf16);
    free(output_bf16);
    free(input_bf16);
}

References C, ck_pytorch_brgemm_fault(), ck_pytorch_brgemm_workspace_sizes(), and gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_workspace_impl().

Referenced by gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage(), and gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage().

gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_workspace_impl()

static void gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_workspace_impl ( const float *  A, const void *  B, const float *  bias, float *  C, int  M, int  N, int  K, uint16_t *  input_bf16, size_t  input_bf16_bytes, uint16_t *  output_bf16, size_t  output_bf16_bytes, uint16_t *  bias_bf16, size_t  bias_bf16_bytes  )

static

Definition at line 1451 of file gemm_kernels_bf16.c.

{
#ifdef USE_ONEDNN
    if (!A || !B || !C || M <= 0 || N <= 0 || K <= 0) {
        ck_pytorch_brgemm_fault("received an invalid tensor contract", M, N, K);
    }
    if (ck_pytorch_brgemm_init_status != 0) {
        ck_pytorch_brgemm_fault("could not initialize oneDNN", M, N, K);
    }
 
    size_t required_input_bytes = 0;
    size_t required_output_bytes = 0;
    size_t required_bias_bytes = 0;
    ck_pytorch_brgemm_workspace_sizes(
        M, N, K, &required_input_bytes, &required_output_bytes,
        &required_bias_bytes);
    const size_t input_count = required_input_bytes / sizeof(*input_bf16);
    const size_t output_count = required_output_bytes / sizeof(*output_bf16);
    if (!input_bf16 || input_bf16_bytes < required_input_bytes ||
        !output_bf16 || output_bf16_bytes < required_output_bytes ||
        (bias && (!bias_bf16 || bias_bf16_bytes < required_bias_bytes))) {
        ck_pytorch_brgemm_fault("received an undersized workspace", M, N, K);
    }
    for (size_t i = 0; i < input_count; ++i) input_bf16[i] = float_to_bf16(A[i]);
    if (bias) {
        for (int j = 0; j < N; ++j) bias_bf16[j] = float_to_bf16(bias[j]);
        for (int i = 0; i < M; ++i) {
            memcpy(output_bf16 + (size_t)i * (size_t)N,
                   bias_bf16, (size_t)N * sizeof(*bias_bf16));
        }
    }
 
    dnnl_memory_desc_t src_md = NULL, weights_md = NULL, dst_md = NULL;
    dnnl_primitive_attr_t attr = NULL;
    dnnl_post_ops_t post_ops = NULL;
    dnnl_primitive_desc_t primitive_desc = NULL;
    dnnl_primitive_t primitive = NULL;
    dnnl_memory_t src_mem = NULL, weights_mem = NULL, dst_mem = NULL;
    dnnl_dims_t src_dims = {M, K};
    dnnl_dims_t weights_dims = {K, N};
    dnnl_dims_t dst_dims = {M, N};
    dnnl_dims_t src_strides = {K, 1};
    dnnl_dims_t weights_strides = {1, K};
    dnnl_dims_t dst_strides = {N, 1};
    dnnl_status_t status = dnnl_success;
 
#define CK_DNNL(call) do { status = (call); if (status != dnnl_success) goto cleanup; } while (0)
    pthread_mutex_lock(&ck_pytorch_brgemm_lock);
    CK_DNNL(dnnl_memory_desc_create_with_strides(&src_md, 2, src_dims, dnnl_bf16, src_strides));
    CK_DNNL(dnnl_memory_desc_create_with_strides(
        &weights_md, 2, weights_dims, dnnl_bf16, weights_strides));
    CK_DNNL(dnnl_memory_desc_create_with_strides(&dst_md, 2, dst_dims, dnnl_bf16, dst_strides));
    if (bias) {
        CK_DNNL(dnnl_primitive_attr_create(&attr));
        CK_DNNL(dnnl_post_ops_create(&post_ops));
        CK_DNNL(dnnl_post_ops_append_sum(post_ops, 1.0f, 0, dnnl_bf16));
        CK_DNNL(dnnl_primitive_attr_set_post_ops(attr, post_ops));
    }
    CK_DNNL(dnnl_matmul_primitive_desc_create(
        &primitive_desc, ck_pytorch_brgemm_engine, src_md, weights_md, NULL, dst_md, attr));
    CK_DNNL(dnnl_primitive_create(&primitive, primitive_desc));
    CK_DNNL(dnnl_memory_create(&src_mem, src_md, ck_pytorch_brgemm_engine, input_bf16));
    CK_DNNL(dnnl_memory_create(
        &weights_mem, weights_md, ck_pytorch_brgemm_engine, (void *)B));
    CK_DNNL(dnnl_memory_create(&dst_mem, dst_md, ck_pytorch_brgemm_engine, output_bf16));
    dnnl_exec_arg_t args[] = {
        {DNNL_ARG_SRC, src_mem},
        {DNNL_ARG_WEIGHTS, weights_mem},
        {DNNL_ARG_DST, dst_mem},
    };
    CK_DNNL(dnnl_primitive_execute(
        primitive, ck_pytorch_brgemm_stream, (int)(sizeof(args) / sizeof(args[0])), args));
    CK_DNNL(dnnl_stream_wait(ck_pytorch_brgemm_stream));
 
cleanup:
    if (dst_mem) dnnl_memory_destroy(dst_mem);
    if (weights_mem) dnnl_memory_destroy(weights_mem);
    if (src_mem) dnnl_memory_destroy(src_mem);
    if (primitive) dnnl_primitive_destroy(primitive);
    if (primitive_desc) dnnl_primitive_desc_destroy(primitive_desc);
    if (dst_md) dnnl_memory_desc_destroy(dst_md);
    if (weights_md) dnnl_memory_desc_destroy(weights_md);
    if (src_md) dnnl_memory_desc_destroy(src_md);
    if (post_ops) dnnl_post_ops_destroy(post_ops);
    if (attr) dnnl_primitive_attr_destroy(attr);
    pthread_mutex_unlock(&ck_pytorch_brgemm_lock);
#undef CK_DNNL
 
    if (status != dnnl_success) {
        ck_pytorch_brgemm_fault("execution failed", M, N, K);
    }
    for (size_t i = 0; i < output_count; ++i) C[i] = bf16_to_float(output_bf16[i]);
#else
    (void)A; (void)B; (void)bias; (void)C; (void)M; (void)N; (void)K;
    (void)input_bf16; (void)input_bf16_bytes;
    (void)output_bf16; (void)output_bf16_bytes;
    (void)bias_bf16; (void)bias_bf16_bytes;
    fprintf(stderr, "HARD KERNEL CONTRACT FAULT: PyTorch oneDNN BF16 BRGEMM was "
                    "selected without USE_ONEDNN=1\n");
    abort();
#endif
}

References bf16_to_float(), C, ck_pytorch_brgemm_fault(), ck_pytorch_brgemm_workspace_sizes(), and float_to_bf16().

Referenced by gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage_workspace(), and gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_impl().

gemm_nt_bf16_row_range()

void gemm_nt_bf16_row_range	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K,
		int	row_begin,
		int	row_end
	)

Definition at line 727 of file gemm_kernels_bf16.c.

{
    const uint16_t *w = (const uint16_t *)B;
    if (!A || !w || !C || M <= 0 || N <= 0 || K <= 0 ||
        row_begin < 0 || row_begin >= row_end || row_end > M) {
        return;
    }
 
    for (int i = row_begin; i < row_end; ++i) {
        const float *a_row = A + (size_t)i * (size_t)K;
        float *c_row = C + (size_t)i * (size_t)N;
        for (int j = 0; j < N; ++j) {
            const uint16_t *w_row = w + (size_t)j * (size_t)K;
            float sum = bias ? bias[j] : 0.0f;
            for (int k = 0; k < K; ++k) {
                const float ab = bf16_to_float(float_to_bf16(a_row[k]));
                sum += ab * bf16_to_float(w_row[k]);
            }
            c_row[j] = sum;
        }
    }
}

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

Referenced by ck_gemm_nt_bf16_exact_rows(), gemm_nt_bf16(), and gemm_nt_bf16_bf16_storage_row_range().

gemm_tn_bf16()

void gemm_tn_bf16	(	const uint16_t *	A,
		const uint16_t *	B,
		const uint16_t *	bias,
		uint16_t *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 882 of file gemm_kernels_bf16.c.

{
    if (!A || !B || !C || M <= 0 || N <= 0 || K <= 0) {
        return;
    }
 
    /* A is [K x M], we want A.T which is [M x K] */
    /* B is [K x N] */
    /* C is [M x N] */
 
#if defined(__AVX512F__)
    /* Initialize C with bias */
    #pragma omp parallel for
    for (int i = 0; i < M; ++i) {
        for (int j = 0; j < N; ++j) {
            float b = bias ? bf16_to_float(bias[j]) : 0.0f;
            C[(size_t)i * N + j] = float_to_bf16(b);
        }
    }
 
    /* Accumulate: C[i,j] += sum_k A[k,i] * B[k,j] */
    #pragma omp parallel for
    for (int i = 0; i < M; ++i) {
        for (int j = 0; j < N; ++j) {
            __m512 sum_vec = _mm512_setzero_ps();
 
            int k = 0;
            for (; k <= K - 16; k += 16) {
                /* Gather A[k:k+16, i] - strided access */
                __m512 a_fp32 = _mm512_setzero_ps();
                for (int kk = 0; kk < 16; ++kk) {
                    float val = bf16_to_float(A[(size_t)(k + kk) * M + i]);
                    a_fp32 = _mm512_mask_mov_ps(a_fp32, 1 << kk, _mm512_set1_ps(val));
                }
 
                /* Note: B has stride N, so we need to gather element by element */
                __m512 b_fp32 = _mm512_setzero_ps();
                for (int kk = 0; kk < 16; ++kk) {
                    float val = bf16_to_float(B[(size_t)(k + kk) * N + j]);
                    b_fp32 = _mm512_mask_mov_ps(b_fp32, 1 << kk, _mm512_set1_ps(val));
                }
 
                sum_vec = _mm512_fmadd_ps(a_fp32, b_fp32, sum_vec);
            }
 
            float sum = _mm512_reduce_add_ps(sum_vec);
 
            for (; k < K; ++k) {
                sum += bf16_to_float(A[(size_t)k * M + i]) *
                       bf16_to_float(B[(size_t)k * N + j]);
            }
 
            float old_val = bf16_to_float(C[(size_t)i * N + j]);
            C[(size_t)i * N + j] = float_to_bf16(old_val + sum);
        }
    }
#else
    for (int i = 0; i < M; ++i) {
        for (int j = 0; j < N; ++j) {
            float sum = bias ? bf16_to_float(bias[j]) : 0.0f;
            for (int k = 0; k < K; ++k) {
                sum += bf16_to_float(A[(size_t)k * M + i]) *
                       bf16_to_float(B[(size_t)k * N + j]);
            }
            C[(size_t)i * N + j] = float_to_bf16(sum);
        }
    }
#endif
}

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

gemv_bf16()

void gemv_bf16	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Definition at line 629 of file gemm_kernels_bf16.c.

{
    gemv_bf16_row_range(
        y, (const uint16_t *)W, x, M, K, 0, M);
}

References gemv_bf16_row_range().

Referenced by gemv_bf16_parallel_dispatch().

gemv_bf16_bf16_storage()

void gemv_bf16_bf16_storage	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Definition at line 719 of file gemm_kernels_bf16.c.

{
    gemv_bf16_bf16_storage_parallel_dispatch(y, W, x, M, K);
}

References gemv_bf16_bf16_storage_parallel_dispatch().

gemv_bf16_bf16_storage_parallel_dispatch()

void gemv_bf16_bf16_storage_parallel_dispatch	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Definition at line 695 of file gemm_kernels_bf16.c.

{
    ck_threadpool_t *pool = ck_threadpool_global();
    if (!pool || ck_threadpool_n_threads(pool) <= 1 || M < 2 ||
        (size_t)M * (size_t)K <= 65536) {
        gemv_bf16_bf16_storage_row_range(
            y, (const uint16_t *)W, x, M, K, 0, M);
        return;
    }
 
    ck_gemv_bf16_args_t args = {
        .y = y, .w = (const uint16_t *)W, .x = x, .M = M, .K = K,
    };
    int active = ck_threadpool_n_threads(pool);
    if (active > M) active = M;
    int grain = M / (active * 4);
    if (grain < 1) grain = 1;
    ck_threadpool_parallel_for_n(
        pool, active, 0, M, grain, ck_gemv_bf16_storage_rows, &args);
}

References ck_gemv_bf16_storage_rows(), ck_threadpool_global(), ck_threadpool_n_threads(), ck_threadpool_parallel_for_n(), and gemv_bf16_bf16_storage_row_range().

Referenced by gemv_bf16_bf16_storage().

gemv_bf16_bf16_storage_row_range()

static void gemv_bf16_bf16_storage_row_range ( float *  y, const uint16_t *  w, const float *  x, int  M, int  K, int  row_begin, int  row_end  )

static

Definition at line 676 of file gemm_kernels_bf16.c.

{
    gemv_bf16_row_range(y, w, x, M, K, row_begin, row_end);
    for (int row = row_begin; row < row_end; ++row) {
        y[row] = bf16_to_float(float_to_bf16(y[row]));
    }
}

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

Referenced by ck_gemv_bf16_storage_rows(), and gemv_bf16_bf16_storage_parallel_dispatch().

gemv_bf16_parallel_dispatch()

void gemv_bf16_parallel_dispatch	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Definition at line 653 of file gemm_kernels_bf16.c.

{
    ck_threadpool_t *pool = ck_threadpool_global();
    if (!pool || ck_threadpool_n_threads(pool) <= 1 || M < 2 ||
        (size_t)M * (size_t)K <= 65536) {
        gemv_bf16(y, W, x, M, K);
        return;
    }
 
    ck_gemv_bf16_args_t args = {
        .y = y, .w = (const uint16_t *)W, .x = x, .M = M, .K = K,
    };
    int active = ck_threadpool_n_threads(pool);
    if (active > M) active = M;
    int grain = M / (active * 4);
    if (grain < 1) grain = 1;
    ck_threadpool_parallel_for_n(
        pool, active, 0, M, grain, ck_gemv_bf16_rows, &args);
}

References ck_gemv_bf16_rows(), ck_threadpool_global(), ck_threadpool_n_threads(), ck_threadpool_parallel_for_n(), and gemv_bf16().

gemv_bf16_row_range()

static void gemv_bf16_row_range ( float *  y, const uint16_t *  w, const float *  x, int  M, int  K, int  row_begin, int  row_end  )

static

Definition at line 607 of file gemm_kernels_bf16.c.

{
    if (!y || !w || !x || M <= 0 || K <= 0 ||
        row_begin < 0 || row_begin >= row_end || row_end > M) {
        return;
    }
 
    for (int i = row_begin; i < row_end; ++i) {
        const uint16_t *w_row = w + (size_t)i * (size_t)K;
        float sum = 0.0f;
        for (int k = 0; k < K; ++k) {
            const float xb = bf16_to_float(float_to_bf16(x[k]));
            sum += xb * bf16_to_float(w_row[k]);
        }
        y[i] = sum;
    }
}

References bf16_to_float(), and float_to_bf16().

Referenced by ck_gemv_bf16_rows(), gemv_bf16(), and gemv_bf16_bf16_storage_row_range().

patch_projection_bf16_pytorch_onednn_conv3d_storage()

void patch_projection_bf16_pytorch_onednn_conv3d_storage	(	const float *	input,
		const void *	weights,
		const float *	bias,
		float *	output,
		int	batch,
		int	out_channels,
		int	in_channels,
		int	temporal,
		int	patch_h,
		int	patch_w
	)

Definition at line 1633 of file gemm_kernels_bf16.c.

{
#ifdef USE_ONEDNN
    if (!input || !weights || !bias || !output || batch <= 0 ||
        out_channels <= 0 || in_channels <= 0 || temporal <= 0 ||
        patch_h <= 0 || patch_w <= 0) {
        ck_pytorch_brgemm_fault("invalid Conv3D patch contract", batch,
                                out_channels, in_channels * temporal * patch_h * patch_w);
    }
    pthread_once(&ck_pytorch_brgemm_once, ck_pytorch_brgemm_init);
    if (ck_pytorch_brgemm_init_status != 0) {
        ck_pytorch_brgemm_fault("could not initialize oneDNN Conv3D", batch,
                                out_channels, in_channels * temporal * patch_h * patch_w);
    }
    ck_pytorch_brgemm_require_version(
        3, 7, 1, "8d263e693366ef8db40acc569cc7d8edf644556d",
        "patch_projection_bf16_pytorch_onednn_conv3d_storage",
        batch, out_channels, in_channels * temporal * patch_h * patch_w);
 
    const size_t input_count = (size_t)batch * (size_t)in_channels *
        (size_t)temporal * (size_t)patch_h * (size_t)patch_w;
    const size_t output_count = (size_t)batch * (size_t)out_channels;
    uint16_t *input_bf16 = (uint16_t *)malloc(input_count * sizeof(*input_bf16));
    uint16_t *bias_bf16 = (uint16_t *)malloc((size_t)out_channels * sizeof(*bias_bf16));
    uint16_t *output_bf16 = (uint16_t *)malloc(output_count * sizeof(*output_bf16));
    if (!input_bf16 || !bias_bf16 || !output_bf16) {
        free(output_bf16);
        free(bias_bf16);
        free(input_bf16);
        ck_pytorch_brgemm_fault("Conv3D workspace allocation failed", batch,
                                out_channels, in_channels * temporal * patch_h * patch_w);
    }
    for (size_t i = 0; i < input_count; ++i) input_bf16[i] = float_to_bf16(input[i]);
    for (int i = 0; i < out_channels; ++i) bias_bf16[i] = float_to_bf16(bias[i]);
 
    dnnl_dims_t src_dims = {batch, in_channels, temporal, patch_h, patch_w};
    dnnl_dims_t weight_dims = {
        out_channels, in_channels, temporal, patch_h, patch_w};
    dnnl_dims_t bias_dims = {out_channels};
    dnnl_dims_t dst_dims = {batch, out_channels, 1, 1, 1};
    dnnl_dims_t strides = {temporal, patch_h, patch_w};
    dnnl_dims_t dilates = {0, 0, 0};
    dnnl_dims_t padding = {0, 0, 0};
 
    dnnl_memory_desc_t user_src_md = NULL, user_weight_md = NULL;
    dnnl_memory_desc_t bias_md = NULL, user_dst_md = NULL;
    dnnl_memory_desc_t any_src_md = NULL, any_weight_md = NULL, any_dst_md = NULL;
    dnnl_primitive_desc_t conv_pd = NULL;
    dnnl_primitive_t conv = NULL;
    dnnl_primitive_desc_t reorder_pd = NULL;
    dnnl_primitive_t reorder = NULL;
    dnnl_exec_arg_t reorder_args[2];
    dnnl_memory_t user_src = NULL, user_weight = NULL, bias_mem = NULL, user_dst = NULL;
    dnnl_memory_t conv_src = NULL, conv_weight = NULL, conv_dst = NULL;
    dnnl_status_t status = dnnl_success;
 
#define CK_DNNL_CONV(call) do { status = (call); if (status != dnnl_success) goto cleanup_conv; } while (0)
    pthread_mutex_lock(&ck_pytorch_brgemm_lock);
    CK_DNNL_CONV(dnnl_memory_desc_create_with_tag(
        &user_src_md, 5, src_dims, dnnl_bf16, dnnl_ncdhw));
    CK_DNNL_CONV(dnnl_memory_desc_create_with_tag(
        &user_weight_md, 5, weight_dims, dnnl_bf16, dnnl_oidhw));
    CK_DNNL_CONV(dnnl_memory_desc_create_with_tag(
        &bias_md, 1, bias_dims, dnnl_bf16, dnnl_x));
    CK_DNNL_CONV(dnnl_memory_desc_create_with_tag(
        &user_dst_md, 5, dst_dims, dnnl_bf16, dnnl_ncdhw));
    CK_DNNL_CONV(dnnl_memory_desc_create_with_tag(
        &any_src_md, 5, src_dims, dnnl_bf16, dnnl_format_tag_any));
    CK_DNNL_CONV(dnnl_memory_desc_create_with_tag(
        &any_weight_md, 5, weight_dims, dnnl_bf16, dnnl_format_tag_any));
    CK_DNNL_CONV(dnnl_memory_desc_create_with_tag(
        &any_dst_md, 5, dst_dims, dnnl_bf16, dnnl_format_tag_any));
    CK_DNNL_CONV(dnnl_convolution_forward_primitive_desc_create(
        &conv_pd, ck_pytorch_brgemm_engine, dnnl_forward_training,
        dnnl_convolution_direct, any_src_md, any_weight_md, bias_md, any_dst_md,
        strides, dilates, padding, padding, NULL));
 
    const_dnnl_memory_desc_t conv_src_md =
        dnnl_primitive_desc_query_md(conv_pd, dnnl_query_src_md, 0);
    const_dnnl_memory_desc_t conv_weight_md =
        dnnl_primitive_desc_query_md(conv_pd, dnnl_query_weights_md, 0);
    const_dnnl_memory_desc_t conv_dst_md =
        dnnl_primitive_desc_query_md(conv_pd, dnnl_query_dst_md, 0);
    CK_DNNL_CONV(dnnl_memory_create(
        &user_src, user_src_md, ck_pytorch_brgemm_engine, input_bf16));
    CK_DNNL_CONV(dnnl_memory_create(
        &user_weight, user_weight_md, ck_pytorch_brgemm_engine, (void *)weights));
    CK_DNNL_CONV(dnnl_memory_create(
        &bias_mem, bias_md, ck_pytorch_brgemm_engine, bias_bf16));
    CK_DNNL_CONV(dnnl_memory_create(
        &user_dst, user_dst_md, ck_pytorch_brgemm_engine, output_bf16));
    CK_DNNL_CONV(dnnl_memory_create(
        &conv_src, conv_src_md, ck_pytorch_brgemm_engine, DNNL_MEMORY_ALLOCATE));
    CK_DNNL_CONV(dnnl_memory_create(
        &conv_weight, conv_weight_md, ck_pytorch_brgemm_engine, DNNL_MEMORY_ALLOCATE));
    CK_DNNL_CONV(dnnl_memory_create(
        &conv_dst, conv_dst_md, ck_pytorch_brgemm_engine, DNNL_MEMORY_ALLOCATE));
 
    CK_DNNL_CONV(dnnl_reorder_primitive_desc_create(
        &reorder_pd, user_src_md, ck_pytorch_brgemm_engine,
        conv_src_md, ck_pytorch_brgemm_engine, NULL));
    CK_DNNL_CONV(dnnl_primitive_create(&reorder, reorder_pd));
    reorder_args[0] = (dnnl_exec_arg_t){DNNL_ARG_FROM, user_src};
    reorder_args[1] = (dnnl_exec_arg_t){DNNL_ARG_TO, conv_src};
    CK_DNNL_CONV(dnnl_primitive_execute(
        reorder, ck_pytorch_brgemm_stream, 2, reorder_args));
    dnnl_primitive_destroy(reorder); reorder = NULL;
    dnnl_primitive_desc_destroy(reorder_pd); reorder_pd = NULL;
 
    CK_DNNL_CONV(dnnl_reorder_primitive_desc_create(
        &reorder_pd, user_weight_md, ck_pytorch_brgemm_engine,
        conv_weight_md, ck_pytorch_brgemm_engine, NULL));
    CK_DNNL_CONV(dnnl_primitive_create(&reorder, reorder_pd));
    reorder_args[0] = (dnnl_exec_arg_t){DNNL_ARG_FROM, user_weight};
    reorder_args[1] = (dnnl_exec_arg_t){DNNL_ARG_TO, conv_weight};
    CK_DNNL_CONV(dnnl_primitive_execute(
        reorder, ck_pytorch_brgemm_stream, 2, reorder_args));
    dnnl_primitive_destroy(reorder); reorder = NULL;
    dnnl_primitive_desc_destroy(reorder_pd); reorder_pd = NULL;
 
    CK_DNNL_CONV(dnnl_primitive_create(&conv, conv_pd));
    dnnl_exec_arg_t conv_args[] = {
        {DNNL_ARG_SRC, conv_src},
        {DNNL_ARG_WEIGHTS, conv_weight},
        {DNNL_ARG_BIAS, bias_mem},
        {DNNL_ARG_DST, conv_dst},
    };
    CK_DNNL_CONV(dnnl_primitive_execute(
        conv, ck_pytorch_brgemm_stream,
        (int)(sizeof(conv_args) / sizeof(conv_args[0])), conv_args));
 
    CK_DNNL_CONV(dnnl_reorder_primitive_desc_create(
        &reorder_pd, conv_dst_md, ck_pytorch_brgemm_engine,
        user_dst_md, ck_pytorch_brgemm_engine, NULL));
    CK_DNNL_CONV(dnnl_primitive_create(&reorder, reorder_pd));
    reorder_args[0] = (dnnl_exec_arg_t){DNNL_ARG_FROM, conv_dst};
    reorder_args[1] = (dnnl_exec_arg_t){DNNL_ARG_TO, user_dst};
    CK_DNNL_CONV(dnnl_primitive_execute(
        reorder, ck_pytorch_brgemm_stream, 2, reorder_args));
    CK_DNNL_CONV(dnnl_stream_wait(ck_pytorch_brgemm_stream));
 
cleanup_conv:
    if (reorder) dnnl_primitive_destroy(reorder);
    if (reorder_pd) dnnl_primitive_desc_destroy(reorder_pd);
    if (conv) dnnl_primitive_destroy(conv);
    if (conv_dst) dnnl_memory_destroy(conv_dst);
    if (conv_weight) dnnl_memory_destroy(conv_weight);
    if (conv_src) dnnl_memory_destroy(conv_src);
    if (user_dst) dnnl_memory_destroy(user_dst);
    if (bias_mem) dnnl_memory_destroy(bias_mem);
    if (user_weight) dnnl_memory_destroy(user_weight);
    if (user_src) dnnl_memory_destroy(user_src);
    if (conv_pd) dnnl_primitive_desc_destroy(conv_pd);
    if (any_dst_md) dnnl_memory_desc_destroy(any_dst_md);
    if (any_weight_md) dnnl_memory_desc_destroy(any_weight_md);
    if (any_src_md) dnnl_memory_desc_destroy(any_src_md);
    if (user_dst_md) dnnl_memory_desc_destroy(user_dst_md);
    if (bias_md) dnnl_memory_desc_destroy(bias_md);
    if (user_weight_md) dnnl_memory_desc_destroy(user_weight_md);
    if (user_src_md) dnnl_memory_desc_destroy(user_src_md);
    pthread_mutex_unlock(&ck_pytorch_brgemm_lock);
#undef CK_DNNL_CONV
 
    if (status != dnnl_success) {
        free(output_bf16);
        free(bias_bf16);
        free(input_bf16);
        ck_pytorch_brgemm_fault("oneDNN Conv3D execution failed", batch,
                                out_channels, in_channels * temporal * patch_h * patch_w);
    }
    for (size_t i = 0; i < output_count; ++i) output[i] = bf16_to_float(output_bf16[i]);
    free(output_bf16);
    free(bias_bf16);
    free(input_bf16);
#else
    (void)input; (void)weights; (void)bias; (void)output; (void)batch;
    (void)out_channels; (void)in_channels; (void)temporal; (void)patch_h; (void)patch_w;
    fprintf(stderr, "HARD KERNEL CONTRACT FAULT: PyTorch oneDNN BF16 Conv3D "
                    "was selected without USE_ONEDNN=1\n");
    abort();
#endif
}

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

Referenced by patch_projection_image_bf16_pytorch_onednn_conv3d_storage().

patch_projection_image_bf16_native_storage()

void patch_projection_image_bf16_native_storage	(	const float *	image,
		const void *	weights_t0,
		const void *	weights_t1,
		const float *	bias,
		float *	output,
		int	channels,
		int	image_h,
		int	image_w,
		int	patch_size,
		int	out_channels,
		int	merge_size
	)

Definition at line 2009 of file gemm_kernels_bf16.c.

{
    if (!image || !weights_t0 || !weights_t1 || !output || channels <= 0 ||
        image_h <= 0 || image_w <= 0 || patch_size <= 0 || out_channels <= 0 ||
        merge_size <= 0 || image_h % patch_size != 0 ||
        image_w % patch_size != 0) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: invalid native BF16 image patch projection\n");
        abort();
    }
    const int grid_h = image_h / patch_size;
    const int grid_w = image_w / patch_size;
    if (grid_h % merge_size != 0 || grid_w % merge_size != 0) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: native BF16 patch grid is not merge aligned\n");
        abort();
    }
    ck_patch_projection_bf16_native_args_t args = {
        .image = image,
        .weights_t0 = (const uint16_t *)weights_t0,
        .weights_t1 = (const uint16_t *)weights_t1,
        .bias = bias,
        .output = output,
        .channels = channels,
        .image_h = image_h,
        .image_w = image_w,
        .patch_size = patch_size,
        .out_channels = out_channels,
        .merge_size = merge_size,
        .grid_w = grid_w,
        .batch = grid_h * grid_w,
    };
    ck_threadpool_t *pool = ck_threadpool_global();
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    if (active > args.batch) active = args.batch;
    if (active > 24) active = 24;
    if (pool && active > 1) {
        ck_threadpool_dispatch_n(
            pool, active, ck_patch_projection_bf16_native_work, &args);
    } else {
        ck_patch_projection_bf16_native_work(0, 1, &args);
    }
}

References ck_patch_projection_bf16_native_work(), ck_threadpool_dispatch_n(), ck_threadpool_global(), and ck_threadpool_n_threads().

patch_projection_image_bf16_pytorch_onednn_conv3d_storage()

void patch_projection_image_bf16_pytorch_onednn_conv3d_storage	(	const float *	image,
		const void *	weights_t0,
		const void *	weights_t1,
		const float *	bias,
		float *	output,
		int	channels,
		int	image_h,
		int	image_w,
		int	patch_size,
		int	out_channels,
		int	merge_size
	)

Definition at line 1819 of file gemm_kernels_bf16.c.

{
#ifdef USE_ONEDNN
    if (!image || !weights_t0 || !weights_t1 || !bias || !output ||
        channels <= 0 || image_h <= 0 || image_w <= 0 || patch_size <= 0 ||
        out_channels <= 0 || merge_size <= 0 || image_h % patch_size != 0 ||
        image_w % patch_size != 0) {
        ck_pytorch_brgemm_fault("invalid image patch projection contract",
                                image_h, image_w, patch_size);
    }
    const int grid_h = image_h / patch_size;
    const int grid_w = image_w / patch_size;
    if (grid_h % merge_size != 0 || grid_w % merge_size != 0) {
        ck_pytorch_brgemm_fault("patch grid is not merge-tile aligned",
                                grid_h, grid_w, merge_size);
    }
    const int batch = grid_h * grid_w;
    const int temporal = 2;
    const int half_k = channels * patch_size * patch_size;
    const int full_k = temporal * half_k;
    float *patches = (float *)malloc((size_t)batch * (size_t)full_k * sizeof(*patches));
    uint16_t *weights = (uint16_t *)malloc(
        (size_t)out_channels * (size_t)full_k * sizeof(*weights));
    if (!patches || !weights) {
        free(weights);
        free(patches);
        ck_pytorch_brgemm_fault("image patch projection workspace allocation failed",
                                batch, out_channels, full_k);
    }
 
    for (int tok = 0; tok < batch; ++tok) {
        const int tiles_per_row = grid_w / merge_size;
        const int tile_area = merge_size * merge_size;
        const int tile = tok / tile_area;
        const int within = tok % tile_area;
        const int patch_y = (tile / tiles_per_row) * merge_size + within / merge_size;
        const int patch_x = (tile % tiles_per_row) * merge_size + within % merge_size;
        float *dst = patches + (size_t)tok * (size_t)full_k;
        for (int c = 0; c < channels; ++c) {
            for (int t = 0; t < temporal; ++t) {
                for (int py = 0; py < patch_size; ++py) {
                    const float *src = image +
                        ((size_t)c * (size_t)image_h +
                         (size_t)(patch_y * patch_size + py)) * (size_t)image_w +
                        (size_t)(patch_x * patch_size);
                    memcpy(dst, src, (size_t)patch_size * sizeof(*dst));
                    dst += patch_size;
                }
            }
        }
    }
 
    const uint16_t *w0 = (const uint16_t *)weights_t0;
    const uint16_t *w1 = (const uint16_t *)weights_t1;
    for (int n = 0; n < out_channels; ++n) {
        uint16_t *dst = weights + (size_t)n * (size_t)full_k;
        for (int c = 0; c < channels; ++c) {
            const size_t channel_offset =
                (size_t)n * (size_t)half_k +
                (size_t)c * (size_t)patch_size * (size_t)patch_size;
            const size_t plane_bytes =
                (size_t)patch_size * (size_t)patch_size * sizeof(*dst);
            memcpy(dst, w0 + channel_offset, plane_bytes);
            dst += patch_size * patch_size;
            memcpy(dst, w1 + channel_offset, plane_bytes);
            dst += patch_size * patch_size;
        }
    }
 
    patch_projection_bf16_pytorch_onednn_conv3d_storage(
        patches, weights, bias, output, batch, out_channels, channels,
        temporal, patch_size, patch_size);
    free(weights);
    free(patches);
#else
    (void)image; (void)weights_t0; (void)weights_t1; (void)bias; (void)output;
    (void)channels; (void)image_h; (void)image_w; (void)patch_size;
    (void)out_channels; (void)merge_size;
    fprintf(stderr, "HARD KERNEL CONTRACT FAULT: PyTorch oneDNN BF16 image "
                    "patch projection was selected without USE_ONEDNN=1\n");
    abort();
#endif
}

References ck_pytorch_brgemm_fault(), and patch_projection_bf16_pytorch_onednn_conv3d_storage().