94 lines
3.9 KiB
C
94 lines
3.9 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/**
|
|
* @file fm_kernel_gemm_f32.h
|
|
* @brief Dense float32 GEMM / GEMV via MPS (MetalPerformanceShaders), single
|
|
* and batched. No custom shader — MPS supplies the kernels, this unit owns
|
|
* the descriptors, the stride validation and the encode.
|
|
*
|
|
* All matrices are row-major float32.
|
|
*
|
|
* @status None
|
|
* @issues None
|
|
* @todo None
|
|
*/
|
|
|
|
#ifndef FOUNDRY_METAL_KERNEL_GEMM_F32_H
|
|
#define FOUNDRY_METAL_KERNEL_GEMM_F32_H
|
|
|
|
#include "metal/fm_metal.h"
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
/* C = alpha * A @ B + beta * C
|
|
* A: [M, K], B: [K, N], C: [M, N] — all row-major float32. */
|
|
void fm_sgemm(fm_buffer *A, fm_buffer *B, fm_buffer *C,
|
|
int M, int N, int K, float alpha, float beta);
|
|
|
|
/* y = alpha * A @ x + beta * y
|
|
* A: [M, K] row-major, x: [K], y: [M] — all float32. */
|
|
void fm_sgemv(fm_buffer *A, fm_buffer *x, fm_buffer *y,
|
|
int M, int K, float alpha, float beta);
|
|
|
|
/**
|
|
* Batch of independent GEMMs in a single MPS submission:
|
|
*
|
|
* C[b] = alpha * A[b] @ B[b] + beta * C[b], for b in [0, batch)
|
|
*
|
|
* All `batch` products share the same M, N, K — this is the uniform-shape
|
|
* batch the chunked delta-net prefill needs (chunk x head), not a ragged one.
|
|
* Encoding one batched kernel instead of `batch` separate ones is the point:
|
|
* it collapses the per-call MPS object construction and dispatch overhead,
|
|
* which dominates when the individual matrices are small.
|
|
*
|
|
* Memory layout (row-major float32, all three buffers):
|
|
*
|
|
* A[b] occupies rows M x cols K starting at byte b * a_matrix_stride_bytes
|
|
* B[b] occupies rows K x cols N starting at byte b * b_matrix_stride_bytes
|
|
* C[b] occupies rows M x cols N starting at byte b * c_matrix_stride_bytes
|
|
*
|
|
* Within a matrix, element (i, j) is at row_bytes * i + sizeof(float) * j,
|
|
* with row_bytes = K, N and N floats respectively — i.e. each matrix is
|
|
* densely packed, only the spacing *between* batch elements is adjustable.
|
|
* A stride of 0 means "packed": M*K, K*N and M*N floats respectively.
|
|
* A non-zero stride must be a multiple of that matrix's row_bytes and at
|
|
* least as large as the packed size; anything else is rejected with a log
|
|
* message and no dispatch. Strides let a caller batch over slices of a
|
|
* larger tensor without copying, and let A and B advance at different rates.
|
|
*
|
|
* Keep hot paths packed. A fully packed batch is encoded as one batched
|
|
* MPS kernel; a padded one falls back to one kernel per batch element in
|
|
* the same command buffer (still a single submission, same results, but
|
|
* without the batched kernel). The fallback exists because MPS sizes a
|
|
* batch from the packed extent and silently drops the tail of a padded
|
|
* one — see the comment in fm_kernel_gemm_f32.m.
|
|
*
|
|
* @param A Buffer holding the batch of [M, K] left matrices
|
|
* @param B Buffer holding the batch of [K, N] right matrices
|
|
* @param C Buffer holding the batch of [M, N] result matrices
|
|
* @param batch Number of products, >= 1
|
|
* @param M, N, K Shape shared by every product, all >= 1
|
|
* @param a_matrix_stride_bytes Byte distance between A[b] and A[b+1], 0 = packed
|
|
* @param b_matrix_stride_bytes Byte distance between B[b] and B[b+1], 0 = packed
|
|
* @param c_matrix_stride_bytes Byte distance between C[b] and C[b+1], 0 = packed
|
|
* @param alpha, beta Scalars applied as in the formula above
|
|
*
|
|
* Thread-safe: no (shares the global command buffer)
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void fm_sgemm_batched(fm_buffer *A, fm_buffer *B, fm_buffer *C,
|
|
int batch, int M, int N, int K,
|
|
size_t a_matrix_stride_bytes,
|
|
size_t b_matrix_stride_bytes,
|
|
size_t c_matrix_stride_bytes,
|
|
float alpha, float beta);
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|
|
|
|
#endif /* FOUNDRY_METAL_KERNEL_GEMM_F32_H */
|