spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// f32 GEMM, correctness-first (kernel roadmap steps 1–2; the simdgroup_matrix4// version is M5). Transpose variants are compiled via function constants5// TA/TB — four specialized pipelines from one source, no runtime branching6// on the hot path beyond what the compiler folds.7//8// matmul_naive_f32: one thread per output element. Grid (N, M), any9//   threadgroup shape.10// matmul_tiled_f32: 16x16 output tile per threadgroup, 16x16 threads,11//   K consumed in 16-wide steps through padded threadgroup tiles (+1 column12//   to break bank conflicts). Edge tiles zero-fill. Grid = ceil(N/16),13//   ceil(M/16) threadgroups of 16x16.14#include <metal_stdlib>15using namespace metal;1617constant bool TA [[function_constant(0)]];18constant bool TB [[function_constant(1)]];1920struct MatmulParams {21    uint M, N, K;22    uint lda, ldb;   // leading dims of A and B as stored23    uint accumulate; // 1: C += A·B (backward-pass gradients)24};2526inline float load_a(device const float* A, uint i, uint k, uint lda) {27    return TA ? A[k * lda + i] : A[i * lda + k];28}29inline float load_b(device const float* B, uint k, uint j, uint ldb) {30    return TB ? B[j * ldb + k] : B[k * ldb + j];31}3233kernel void matmul_naive_f32(device const float* A [[buffer(0)]],34                             device const float* B [[buffer(1)]],35                             device float*       C [[buffer(2)]],36                             constant MatmulParams& p [[buffer(3)]],37                             uint2 gid [[thread_position_in_grid]]) {38    if (gid.x >= p.N || gid.y >= p.M) return;39    float acc = 0.0f;40    for (uint k = 0; k < p.K; ++k)41        acc += load_a(A, gid.y, k, p.lda) * load_b(B, k, gid.x, p.ldb);42    const uint idx = gid.y * p.N + gid.x;43    C[idx] = p.accumulate ? (C[idx] + acc) : acc;44}4546constant constexpr uint TILE = 16;4748kernel void matmul_tiled_f32(device const float* A [[buffer(0)]],49                             device const float* B [[buffer(1)]],50                             device float*       C [[buffer(2)]],51                             constant MatmulParams& p [[buffer(3)]],52                             uint2 tgid [[threadgroup_position_in_grid]],53                             uint2 tid  [[thread_position_in_threadgroup]]) {54    threadgroup float As[TILE][TILE + 1];55    threadgroup float Bs[TILE][TILE + 1];5657    const uint row = tgid.y * TILE + tid.y;58    const uint col = tgid.x * TILE + tid.x;5960    float acc = 0.0f;61    for (uint kt = 0; kt < p.K; kt += TILE) {62        const uint ka = kt + tid.x;63        const uint kb = kt + tid.y;64        As[tid.y][tid.x] = (row < p.M && ka < p.K) ? load_a(A, row, ka, p.lda) : 0.0f;65        Bs[tid.y][tid.x] = (kb < p.K && col < p.N) ? load_b(B, kb, col, p.ldb) : 0.0f;66        threadgroup_barrier(mem_flags::mem_threadgroup);67        for (uint kk = 0; kk < TILE; ++kk)68            acc = fma(As[tid.y][kk], Bs[kk][tid.x], acc);69        threadgroup_barrier(mem_flags::mem_threadgroup);70    }7172    if (row < p.M && col < p.N) {73        const uint idx = row * p.N + col;74        C[idx] = p.accumulate ? (C[idx] + acc) : acc;75    }76}77