// Author: Simon-Pierre Boucher — contact@spboucher.ai
// Micro-benchmark: does simdgroup_matrix run faster with f16/bf16 operands
// than with f32 on this hardware? Same 64x64x16 tiling as the real GEMM;
// only the operand type of the staged tiles and the MMA changes. The
// accumulator stays f32 in every variant (required for training).
#include <metal_stdlib>
#include <metal_simdgroup_matrix>
using namespace metal;

enum : uint { BM = 64, BN = 64, BK = 16, WM = 2, WN = 2,
              NSG = WM*WN, THREADS = NSG*32, TM = BM/(8*WM), TN = BN/(8*WN),
              PAD = 8, LDA_S = BK+PAD, LDB_S = BN+PAD };

struct P { uint M, N, K; };

// T = operand precision, ACC fragment always f32.
template <typename T, typename FRAG>
kernel void gemm_prec(device const T* A [[buffer(0)]],
                      device const T* B [[buffer(1)]],
                      device float*   C [[buffer(2)]],
                      constant P& p [[buffer(3)]],
                      uint2 tgid [[threadgroup_position_in_grid]],
                      uint  tid  [[thread_index_in_threadgroup]],
                      uint  sgid [[simdgroup_index_in_threadgroup]]) {
    threadgroup T As[BM * LDA_S];
    threadgroup T Bs[BK * LDB_S];
    const uint row0 = tgid.y * BM, col0 = tgid.x * BN;
    const uint sr = (sgid / WN) * (TM * 8), sc = (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);
        for (uint e = tid; e < BM*BK; e += THREADS) {
            const uint i = e / BK, k = e % BK;
            As[i*LDA_S + k] = A[(row0+i)*p.K + k0+k];
        }
        for (uint e = tid; e < BK*BN; e += THREADS) {
            const uint k = e / BN, j = e % BN;
            Bs[k*LDB_S + j] = B[(k0+k)*p.N + col0+j];
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
#pragma clang loop unroll(full)
        for (uint kk = 0; kk < BK; kk += 8) {
            FRAG af[TM], bf[TN];
#pragma clang loop unroll(full)
            for (uint i = 0; i < TM; ++i) simdgroup_load(af[i], As + (sr+i*8)*LDA_S + kk, LDA_S);
#pragma clang loop unroll(full)
            for (uint j = 0; j < TN; ++j) simdgroup_load(bf[j], Bs + kk*LDB_S + sc + j*8, LDB_S);
#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_multiply_accumulate(acc[i][j], af[i], bf[j], acc[i][j]);
        }
    }
#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 + (row0+sr+i*8)*p.N + col0+sc+j*8, p.N);
}

template [[host_name("gemm_f32")]] kernel void
gemm_prec<float, simdgroup_float8x8>(device const float*, device const float*,
    device float*, constant P&, uint2, uint, uint);
template [[host_name("gemm_f16")]] kernel void
gemm_prec<half, simdgroup_half8x8>(device const half*, device const half*,
    device float*, constant P&, uint2, uint, uint);
template [[host_name("gemm_bf16")]] kernel void
gemm_prec<bfloat, simdgroup_bfloat8x8>(device const bfloat*, device const bfloat*,
    device float*, constant P&, uint2, uint, uint);