spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// simdgroup_matrix GEMM — the perf-critical path (kernel roadmap step 3).4// Structure follows MLX's STEEL kernel (RESEARCH.md §5): device→threadgroup5// staged tiles with +16-byte row padding against bank conflicts, 8x86// simdgroup fragments, f32 accumulators, transposes handled by function7// constants rather than data movement.8//9// Threadgroup layout (BM=64, BN=64, BK=16, WM=2, WN=2):10//   4 simdgroups = 128 threads. Simdgroup (wm, wn) owns a 32x32 output11//   quadrant = 4x4 grid of 8x8 fragments (TM=TN=4 -> 16 f32 accumulators,12//   32 registers/thread).13//   Threadgroup memory: As 64x(16+4) + Bs 16x(64+4) + a 64-float epilogue14//   scratch per simdgroup = 10.5 KB, so ~3 threadgroups stay resident per15//   core. (An earlier version staged the whole 64x64 output block: 27 KB,16//   one threadgroup per core, 3x slower than the 16x16 tiled kernel.)17//18// Staging is transpose-aware: the thread->address map always walks the19// operand's CONTIGUOUS axis, so all four of nn/nt/tn/tt coalesce. Mapping20// idx by the fixed logical axis instead costs ~3x on the nt case that the21// forward pass (X·Wᵀ) issues most.22//23// Tile constants are ENUMERATORS, not `constant constexpr uint`: in MSL24// `constant` is an address-space qualifier (and program-scope variables must25// live there), so `constant constexpr uint TM = 4` declares a VARIABLE and loop26// bounds built from it are not compile-time known. The fragment loops then don't unroll, `acc[i][j]`27// becomes dynamic indexing into an opaque simdgroup_matrix array, and the28// compiler allocas all 16 accumulators (256 B each) to the stack — every MMA29// pays ~1 KB of stack traffic. That cost 3x: 0.8 -> 2.4+ TFLOPs on f32.30// Verify with: metal -S -emit-llvm, then check for `alloca` and count31// air.simdgroup_matrix_8x8_multiply_accumulate calls (want TM*TN*BK/8, not 2).32//33// Edge handling: ALIGNED is a function constant, so fully-tiled matrices get34// the branch-free path with fragments stored straight to device memory, and35// ragged ones get predicated loads plus a per-fragment predicated epilogue.36// The K remainder is handled by zero-filling the staged tile.37#include <metal_stdlib>38#include <metal_simdgroup_matrix>39using namespace metal;4041constant bool TA [[function_constant(0)]];42constant bool TB [[function_constant(1)]];43constant bool ALIGNED [[function_constant(3)]];4445struct MatmulParams {46    uint M, N, K;47    uint lda, ldb;48    uint accumulate;49};5051// Tile geometry as ENUMERATORS, not `constant` variables — see the note above.52enum : uint {53    BM = 64,54    BN = 64,55    BK = 16,56    WM = 2,             // simdgroups along M57    WN = 2,             // simdgroups along N58    NSG = WM * WN,59    THREADS = NSG * 32,60    TM = BM / (8 * WM), // 8x8 fragments per simdgroup along M61    TN = BN / (8 * WN),62    PAD = 4,            // 16 bytes / sizeof(float)63    LDA_S = BK + PAD,   // staged A row stride64    LDB_S = BN + PAD,   // staged B row stride65};6667inline void stage_a(threadgroup float* As, device const float* A,68                    uint row0, uint k0, constant MatmulParams& p, uint tid) {69    for (uint idx = tid; idx < BM * BK; idx += THREADS) {70        // walk A's contiguous axis with consecutive threads71        const uint i = TA ? (idx % BM) : (idx / BK);72        const uint k = TA ? (idx / BM) : (idx % BK);73        const uint gr = row0 + i;74        const uint gk = k0 + k;75        float v = 0.0f;76        if (ALIGNED || (gr < p.M && gk < p.K))77            v = TA ? A[ulong(gk) * p.lda + gr] : A[ulong(gr) * p.lda + gk];78        As[i * LDA_S + k] = v;79    }80}8182inline void stage_b(threadgroup float* Bs, device const float* B,83                    uint k0, uint col0, constant MatmulParams& p, uint tid) {84    for (uint idx = tid; idx < BK * BN; idx += THREADS) {85        const uint k = TB ? (idx % BK) : (idx / BN);86        const uint j = TB ? (idx / BK) : (idx % BN);87        const uint gk = k0 + k;88        const uint gc = col0 + j;89        float v = 0.0f;90        if (ALIGNED || (gk < p.K && gc < p.N))91            v = TB ? B[ulong(gc) * p.ldb + gk] : B[ulong(gk) * p.ldb + gc];92        Bs[k * LDB_S + j] = v;93    }94}9596kernel void matmul_simd_f32(device const float* A [[buffer(0)]],97                            device const float* B [[buffer(1)]],98                            device float*       C [[buffer(2)]],99                            constant MatmulParams& p [[buffer(3)]],100                            uint2 tgid [[threadgroup_position_in_grid]],101                            uint  tid  [[thread_index_in_threadgroup]],102                            uint  lane [[thread_index_in_simdgroup]],103                            uint  sgid [[simdgroup_index_in_threadgroup]]) {104    threadgroup float As[BM * LDA_S];105    threadgroup float Bs[BK * LDB_S];106    threadgroup float frag_scratch[NSG * 64]; // epilogue only, 1 KB total107108    const uint row0 = tgid.y * BM;109    const uint col0 = tgid.x * BN;110111    // this simdgroup's 32x32 quadrant of the 64x64 block112    const uint sg_row = (sgid / WN) * (TM * 8);113    const uint sg_col = (sgid % WN) * (TN * 8);114115    simdgroup_float8x8 acc[TM][TN];116#pragma clang loop unroll(full)117    for (uint i = 0; i < TM; ++i)118#pragma clang loop unroll(full)119        for (uint j = 0; j < TN; ++j)120            acc[i][j] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);121122    for (uint k0 = 0; k0 < p.K; k0 += BK) {123        threadgroup_barrier(mem_flags::mem_threadgroup);124        stage_a(As, A, row0, k0, p, tid);125        stage_b(Bs, B, k0, col0, p, tid);126        threadgroup_barrier(mem_flags::mem_threadgroup);127128#pragma clang loop unroll(full)129        for (uint kk = 0; kk < BK; kk += 8) {130            simdgroup_float8x8 afrag[TM];131            simdgroup_float8x8 bfrag[TN];132            simdgroup_barrier(mem_flags::mem_none);133#pragma clang loop unroll(full)134            for (uint i = 0; i < TM; ++i)135                simdgroup_load(afrag[i], As + (sg_row + i * 8) * LDA_S + kk, LDA_S);136            simdgroup_barrier(mem_flags::mem_none);137#pragma clang loop unroll(full)138            for (uint j = 0; j < TN; ++j)139                simdgroup_load(bfrag[j], Bs + kk * LDB_S + sg_col + j * 8, LDB_S);140            simdgroup_barrier(mem_flags::mem_none);141#pragma clang loop unroll(full)142            for (uint i = 0; i < TM; ++i) {143                // serpentine N order keeps the last-touched b fragment hot144#pragma clang loop unroll(full)145                for (uint jj = 0; jj < TN; ++jj) {146                    const uint j = (i & 1) ? (TN - 1 - jj) : jj;147                    simdgroup_multiply_accumulate(acc[i][j], afrag[i], bfrag[j], acc[i][j]);148                }149            }150        }151    }152153    const uint out_row = row0 + sg_row;154    const uint out_col = col0 + sg_col;155156    // Fast path: whole block in range and overwriting -> fragments go157    // straight to device memory.158    if (ALIGNED && !p.accumulate) {159#pragma clang loop unroll(full)160        for (uint i = 0; i < TM; ++i)161#pragma clang loop unroll(full)162            for (uint j = 0; j < TN; ++j)163                simdgroup_store(acc[i][j],164                                C + ulong(out_row + i * 8) * p.N + out_col + j * 8, p.N);165        return;166    }167168    // Slow path: one 8x8 fragment at a time through simdgroup-local scratch169    // (no threadgroup barrier needed — only this simdgroup touches it), then170    // 2 predicated element writes per lane.171    threadgroup float* scratch = frag_scratch + sgid * 64;172    for (uint i = 0; i < TM; ++i) {173        for (uint j = 0; j < TN; ++j) {174            simdgroup_barrier(mem_flags::mem_threadgroup);175            simdgroup_store(acc[i][j], scratch, 8);176            simdgroup_barrier(mem_flags::mem_threadgroup);177            for (uint e = lane; e < 64; e += 32) {178                const uint gr = out_row + i * 8 + e / 8;179                const uint gc = out_col + j * 8 + e % 8;180                if (!ALIGNED && (gr >= p.M || gc >= p.N)) continue;181                const ulong o = ulong(gr) * p.N + gc;182                const float v = scratch[e];183                C[o] = p.accumulate ? (C[o] + v) : v;184            }185        }186    }187}188