// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// f32 GEMM, correctness-first (kernel roadmap steps 1–2; the simdgroup_matrix
// version is M5). Transpose variants are compiled via function constants
// TA/TB — four specialized pipelines from one source, no runtime branching
// on the hot path beyond what the compiler folds.
//
// matmul_naive_f32: one thread per output element. Grid (N, M), any
//   threadgroup shape.
// matmul_tiled_f32: 16x16 output tile per threadgroup, 16x16 threads,
//   K consumed in 16-wide steps through padded threadgroup tiles (+1 column
//   to break bank conflicts). Edge tiles zero-fill. Grid = ceil(N/16),
//   ceil(M/16) threadgroups of 16x16.
#include <metal_stdlib>
using namespace metal;

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

struct MatmulParams {
    uint M, N, K;
    uint lda, ldb;   // leading dims of A and B as stored
    uint accumulate; // 1: C += A·B (backward-pass gradients)
};

inline float load_a(device const float* A, uint i, uint k, uint lda) {
    return TA ? A[k * lda + i] : A[i * lda + k];
}
inline float load_b(device const float* B, uint k, uint j, uint ldb) {
    return TB ? B[j * ldb + k] : B[k * ldb + j];
}

kernel void matmul_naive_f32(device const float* A [[buffer(0)]],
                             device const float* B [[buffer(1)]],
                             device float*       C [[buffer(2)]],
                             constant MatmulParams& p [[buffer(3)]],
                             uint2 gid [[thread_position_in_grid]]) {
    if (gid.x >= p.N || gid.y >= p.M) return;
    float acc = 0.0f;
    for (uint k = 0; k < p.K; ++k)
        acc += load_a(A, gid.y, k, p.lda) * load_b(B, k, gid.x, p.ldb);
    const uint idx = gid.y * p.N + gid.x;
    C[idx] = p.accumulate ? (C[idx] + acc) : acc;
}

constant constexpr uint TILE = 16;

kernel void matmul_tiled_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]],
                             uint2 tid  [[thread_position_in_threadgroup]]) {
    threadgroup float As[TILE][TILE + 1];
    threadgroup float Bs[TILE][TILE + 1];

    const uint row = tgid.y * TILE + tid.y;
    const uint col = tgid.x * TILE + tid.x;

    float acc = 0.0f;
    for (uint kt = 0; kt < p.K; kt += TILE) {
        const uint ka = kt + tid.x;
        const uint kb = kt + tid.y;
        As[tid.y][tid.x] = (row < p.M && ka < p.K) ? load_a(A, row, ka, p.lda) : 0.0f;
        Bs[tid.y][tid.x] = (kb < p.K && col < p.N) ? load_b(B, kb, col, p.ldb) : 0.0f;
        threadgroup_barrier(mem_flags::mem_threadgroup);
        for (uint kk = 0; kk < TILE; ++kk)
            acc = fma(As[tid.y][kk], Bs[kk][tid.x], acc);
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    if (row < p.M && col < p.N) {
        const uint idx = row * p.N + col;
        C[idx] = p.accumulate ? (C[idx] + acc) : acc;
    }
}