// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// GEMM via Metal Performance Primitives cooperative tensors
// (mpp::tensor_ops::matmul2d). This is the Metal 4 path that targets the
// per-core neural accelerators on M5-class hardware, and the same one
// llama.cpp uses behind GGML_METAL_HAS_TENSOR. Whether it beats the
// hand-written simdgroup_matrix kernel is a hardware question, so it exists
// here to be benchmarked (tests/bench_precision.cpp) rather than assumed:
// the "Rigel" paper measured matmul2d on an M4 Max still executing on the
// shader cores, with a hand-fused GEMM winning.
//
// Tiling: 64x32 output tile per threadgroup, 4 simdgroups (128 threads), K
// as a dynamic extent so one pipeline serves every K. Shapes must be exact
// multiples of the tile; the caller falls back otherwise.
#include <metal_stdlib>
#include <metal_tensor>
#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>

using namespace metal;
using namespace mpp::tensor_ops;

enum : int { TILE_M = 64, TILE_N = 32 };

// Operands arrive as plain device pointers and become tensors in-kernel via
// the `tensor_inline` descriptor, so the host binds ordinary buffers and needs
// no MTLTensor plumbing (the default `tensor_handle` descriptor wraps an
// opaque handle that only a host-side MTLTensor can supply).
struct MPPParams { uint32_t M, N, K; };

// Transposes are template parameters because the descriptor is a constexpr
// template argument. Training needs all of nn (dX = dY.W), nt (fwd X.W^T) and
// tn (dW = dY^T.X), so all three are instantiated.
template <typename T, bool TA, bool TB>
kernel void matmul_mpp(device T* Ap [[buffer(0)]],
                       device T* Bp [[buffer(1)]],
                       device float*   Cp [[buffer(2)]],
                       constant MPPParams& p [[buffer(3)]],
                       uint2 tgid [[threadgroup_position_in_grid]]) {
    const int32_t M = int32_t(p.M), N = int32_t(p.N), K = int32_t(p.K);
    // extents are (columns, rows): A is MxK, B is KxN, C is MxN
    // Element type must be non-const: MPP static_asserts it is exactly one of
    // uint8_t/int8_t/uint4b/int4b/float/half/bfloat.
    tensor<device T, dextents<int32_t, 2>, tensor_inline> A(
        Ap, TA ? dextents<int32_t, 2>(M, K) : dextents<int32_t, 2>(K, M));
    tensor<device T, dextents<int32_t, 2>, tensor_inline> B(
        Bp, TB ? dextents<int32_t, 2>(K, N) : dextents<int32_t, 2>(N, K));
    tensor<device float, dextents<int32_t, 2>, tensor_inline> C(
        Cp, dextents<int32_t, 2>(N, M));

    // relaxed_precision=false keeps f32 accumulation semantics.
    constexpr auto desc = matmul2d_descriptor(
        TILE_M, TILE_N, static_cast<int>(dynamic_extent),
        /*transpose_left=*/TA, /*transpose_right=*/TB,
        /*relaxed_precision=*/false, matmul2d_descriptor::mode::multiply);

    matmul2d<desc, execution_simdgroups<4>> op;

    // slice() is (column, row); a transposed operand is indexed the other way.
    auto mA = TA ? A.slice(int(tgid.y) * TILE_M, 0) : A.slice(0, int(tgid.y) * TILE_M);
    auto mB = TB ? B.slice(0, int(tgid.x) * TILE_N) : B.slice(int(tgid.x) * TILE_N, 0);
    auto mC = C.slice(int(tgid.x) * TILE_N, int(tgid.y) * TILE_M);

    // The cooperative tensor's element->lane distribution is implementation
    // defined and not every slot a thread holds is valid, hence the
    // is_valid_element guard. (The header's own example still says get_mask,
    // which does not exist in this SDK.)
    auto cT = op.template get_destination_cooperative_tensor<decltype(mA), decltype(mB), float>();
#pragma clang loop unroll(full)
    for (uint16_t i = 0; i < cT.get_capacity(); ++i)
        if (cT.is_valid_element(i)) cT[i] = 0.0f;

    op.run(mA, mB, cT);
    cT.store(mC);
}

#define INST_MPP(SUF, T, TA, TB)                                          \
    template [[host_name("matmul_mpp_" #SUF)]] kernel void                \
    matmul_mpp<T, TA, TB>(device T*, device T*, device float*,            \
                          constant MPPParams&, uint2);

INST_MPP(f32, float, false, false)
INST_MPP(f32_nt, float, false, true)
INST_MPP(f32_tn, float, true, false)
INST_MPP(f16, half, false, false)
INST_MPP(f16_nt, half, false, true)
INST_MPP(f16_tn, half, true, false)