// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// simdgroup_matrix GEMM — the perf-critical path (kernel roadmap step 3).
// Structure follows MLX's STEEL kernel (RESEARCH.md §5): device→threadgroup
// staged tiles with +16-byte row padding against bank conflicts, 8x8
// simdgroup fragments, f32 accumulators, transposes handled by function
// constants rather than data movement.
//
// Threadgroup layout (BM=64, BN=64, BK=16, WM=2, WN=2):
//   4 simdgroups = 128 threads. Simdgroup (wm, wn) owns a 32x32 output
//   quadrant = 4x4 grid of 8x8 fragments (TM=TN=4 -> 16 f32 accumulators,
//   32 registers/thread).
//   Threadgroup memory: As 64x(16+4) + Bs 16x(64+4) + a 64-float epilogue
//   scratch per simdgroup = 10.5 KB, so ~3 threadgroups stay resident per
//   core. (An earlier version staged the whole 64x64 output block: 27 KB,
//   one threadgroup per core, 3x slower than the 16x16 tiled kernel.)
//
// Staging is transpose-aware: the thread->address map always walks the
// operand's CONTIGUOUS axis, so all four of nn/nt/tn/tt coalesce. Mapping
// idx by the fixed logical axis instead costs ~3x on the nt case that the
// forward pass (X·Wᵀ) issues most.
//
// Tile constants are ENUMERATORS, not `constant constexpr uint`: in MSL
// `constant` is an address-space qualifier (and program-scope variables must
// live there), so `constant constexpr uint TM = 4` declares a VARIABLE and loop
// bounds built from it are not compile-time known. The fragment loops then don't unroll, `acc[i][j]`
// becomes dynamic indexing into an opaque simdgroup_matrix array, and the
// compiler allocas all 16 accumulators (256 B each) to the stack — every MMA
// pays ~1 KB of stack traffic. That cost 3x: 0.8 -> 2.4+ TFLOPs on f32.
// Verify with: metal -S -emit-llvm, then check for `alloca` and count
// air.simdgroup_matrix_8x8_multiply_accumulate calls (want TM*TN*BK/8, not 2).
//
// Edge handling: ALIGNED is a function constant, so fully-tiled matrices get
// the branch-free path with fragments stored straight to device memory, and
// ragged ones get predicated loads plus a per-fragment predicated epilogue.
// The K remainder is handled by zero-filling the staged tile.
#include <metal_stdlib>
#include <metal_simdgroup_matrix>
using namespace metal;

constant bool TA [[function_constant(0)]];
constant bool TB [[function_constant(1)]];
constant bool ALIGNED [[function_constant(3)]];

struct MatmulParams {
    uint M, N, K;
    uint lda, ldb;
    uint accumulate;
};

// Tile geometry as ENUMERATORS, not `constant` variables — see the note above.
enum : uint {
    BM = 64,
    BN = 64,
    BK = 16,
    WM = 2,             // simdgroups along M
    WN = 2,             // simdgroups along N
    NSG = WM * WN,
    THREADS = NSG * 32,
    TM = BM / (8 * WM), // 8x8 fragments per simdgroup along M
    TN = BN / (8 * WN),
    PAD = 4,            // 16 bytes / sizeof(float)
    LDA_S = BK + PAD,   // staged A row stride
    LDB_S = BN + PAD,   // staged B row stride
};

inline void stage_a(threadgroup float* As, device const float* A,
                    uint row0, uint k0, constant MatmulParams& p, uint tid) {
    for (uint idx = tid; idx < BM * BK; idx += THREADS) {
        // walk A's contiguous axis with consecutive threads
        const uint i = TA ? (idx % BM) : (idx / BK);
        const uint k = TA ? (idx / BM) : (idx % BK);
        const uint gr = row0 + i;
        const uint gk = k0 + k;
        float v = 0.0f;
        if (ALIGNED || (gr < p.M && gk < p.K))
            v = TA ? A[ulong(gk) * p.lda + gr] : A[ulong(gr) * p.lda + gk];
        As[i * LDA_S + k] = v;
    }
}

inline void stage_b(threadgroup float* Bs, device const float* B,
                    uint k0, uint col0, constant MatmulParams& p, uint tid) {
    for (uint idx = tid; idx < BK * BN; idx += THREADS) {
        const uint k = TB ? (idx % BK) : (idx / BN);
        const uint j = TB ? (idx / BK) : (idx % BN);
        const uint gk = k0 + k;
        const uint gc = col0 + j;
        float v = 0.0f;
        if (ALIGNED || (gk < p.K && gc < p.N))
            v = TB ? B[ulong(gc) * p.ldb + gk] : B[ulong(gk) * p.ldb + gc];
        Bs[k * LDB_S + j] = v;
    }
}

kernel void matmul_simd_f32(device const float* A [[buffer(0)]],
                            device const float* B [[buffer(1)]],
                            device float*       C [[buffer(2)]],
                            constant MatmulParams& p [[buffer(3)]],
                            uint2 tgid [[threadgroup_position_in_grid]],
                            uint  tid  [[thread_index_in_threadgroup]],
                            uint  lane [[thread_index_in_simdgroup]],
                            uint  sgid [[simdgroup_index_in_threadgroup]]) {
    threadgroup float As[BM * LDA_S];
    threadgroup float Bs[BK * LDB_S];
    threadgroup float frag_scratch[NSG * 64]; // epilogue only, 1 KB total

    const uint row0 = tgid.y * BM;
    const uint col0 = tgid.x * BN;

    // this simdgroup's 32x32 quadrant of the 64x64 block
    const uint sg_row = (sgid / WN) * (TM * 8);
    const uint sg_col = (sgid % WN) * (TN * 8);

    simdgroup_float8x8 acc[TM][TN];
#pragma clang loop unroll(full)
    for (uint i = 0; i < TM; ++i)
#pragma clang loop unroll(full)
        for (uint j = 0; j < TN; ++j)
            acc[i][j] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);

    for (uint k0 = 0; k0 < p.K; k0 += BK) {
        threadgroup_barrier(mem_flags::mem_threadgroup);
        stage_a(As, A, row0, k0, p, tid);
        stage_b(Bs, B, k0, col0, p, tid);
        threadgroup_barrier(mem_flags::mem_threadgroup);

#pragma clang loop unroll(full)
        for (uint kk = 0; kk < BK; kk += 8) {
            simdgroup_float8x8 afrag[TM];
            simdgroup_float8x8 bfrag[TN];
            simdgroup_barrier(mem_flags::mem_none);
#pragma clang loop unroll(full)
            for (uint i = 0; i < TM; ++i)
                simdgroup_load(afrag[i], As + (sg_row + i * 8) * LDA_S + kk, LDA_S);
            simdgroup_barrier(mem_flags::mem_none);
#pragma clang loop unroll(full)
            for (uint j = 0; j < TN; ++j)
                simdgroup_load(bfrag[j], Bs + kk * LDB_S + sg_col + j * 8, LDB_S);
            simdgroup_barrier(mem_flags::mem_none);
#pragma clang loop unroll(full)
            for (uint i = 0; i < TM; ++i) {
                // serpentine N order keeps the last-touched b fragment hot
#pragma clang loop unroll(full)
                for (uint jj = 0; jj < TN; ++jj) {
                    const uint j = (i & 1) ? (TN - 1 - jj) : jj;
                    simdgroup_multiply_accumulate(acc[i][j], afrag[i], bfrag[j], acc[i][j]);
                }
            }
        }
    }

    const uint out_row = row0 + sg_row;
    const uint out_col = col0 + sg_col;

    // Fast path: whole block in range and overwriting -> fragments go
    // straight to device memory.
    if (ALIGNED && !p.accumulate) {
#pragma clang loop unroll(full)
        for (uint i = 0; i < TM; ++i)
#pragma clang loop unroll(full)
            for (uint j = 0; j < TN; ++j)
                simdgroup_store(acc[i][j],
                                C + ulong(out_row + i * 8) * p.N + out_col + j * 8, p.N);
        return;
    }

    // Slow path: one 8x8 fragment at a time through simdgroup-local scratch
    // (no threadgroup barrier needed — only this simdgroup touches it), then
    // 2 predicated element writes per lane.
    threadgroup float* scratch = frag_scratch + sgid * 64;
    for (uint i = 0; i < TM; ++i) {
        for (uint j = 0; j < TN; ++j) {
            simdgroup_barrier(mem_flags::mem_threadgroup);
            simdgroup_store(acc[i][j], scratch, 8);
            simdgroup_barrier(mem_flags::mem_threadgroup);
            for (uint e = lane; e < 64; e += 32) {
                const uint gr = out_row + i * 8 + e / 8;
                const uint gc = out_col + j * 8 + e % 8;
                if (!ALIGNED && (gr >= p.M || gc >= p.N)) continue;
                const ulong o = ulong(gr) * p.N + gc;
                const float v = scratch[e];
                C[o] = p.accumulate ? (C[o] + v) : v;
            }
        }
    }
}