spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// GEMM via Metal Performance Primitives cooperative tensors4// (mpp::tensor_ops::matmul2d). This is the Metal 4 path that targets the5// per-core neural accelerators on M5-class hardware, and the same one6// llama.cpp uses behind GGML_METAL_HAS_TENSOR. Whether it beats the7// hand-written simdgroup_matrix kernel is a hardware question, so it exists8// here to be benchmarked (tests/bench_precision.cpp) rather than assumed:9// the "Rigel" paper measured matmul2d on an M4 Max still executing on the10// shader cores, with a hand-fused GEMM winning.11//12// Tiling: 64x32 output tile per threadgroup, 4 simdgroups (128 threads), K13// as a dynamic extent so one pipeline serves every K. Shapes must be exact14// multiples of the tile; the caller falls back otherwise.15#include <metal_stdlib>16#include <metal_tensor>17#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>1819using namespace metal;20using namespace mpp::tensor_ops;2122enum : int { TILE_M = 64, TILE_N = 32 };2324// Operands arrive as plain device pointers and become tensors in-kernel via25// the `tensor_inline` descriptor, so the host binds ordinary buffers and needs26// no MTLTensor plumbing (the default `tensor_handle` descriptor wraps an27// opaque handle that only a host-side MTLTensor can supply).28struct MPPParams { uint32_t M, N, K; };2930// Transposes are template parameters because the descriptor is a constexpr31// template argument. Training needs all of nn (dX = dY.W), nt (fwd X.W^T) and32// tn (dW = dY^T.X), so all three are instantiated.33template <typename T, bool TA, bool TB>34kernel void matmul_mpp(device T* Ap [[buffer(0)]],35                       device T* Bp [[buffer(1)]],36                       device float*   Cp [[buffer(2)]],37                       constant MPPParams& p [[buffer(3)]],38                       uint2 tgid [[threadgroup_position_in_grid]]) {39    const int32_t M = int32_t(p.M), N = int32_t(p.N), K = int32_t(p.K);40    // extents are (columns, rows): A is MxK, B is KxN, C is MxN41    // Element type must be non-const: MPP static_asserts it is exactly one of42    // uint8_t/int8_t/uint4b/int4b/float/half/bfloat.43    tensor<device T, dextents<int32_t, 2>, tensor_inline> A(44        Ap, TA ? dextents<int32_t, 2>(M, K) : dextents<int32_t, 2>(K, M));45    tensor<device T, dextents<int32_t, 2>, tensor_inline> B(46        Bp, TB ? dextents<int32_t, 2>(K, N) : dextents<int32_t, 2>(N, K));47    tensor<device float, dextents<int32_t, 2>, tensor_inline> C(48        Cp, dextents<int32_t, 2>(N, M));4950    // relaxed_precision=false keeps f32 accumulation semantics.51    constexpr auto desc = matmul2d_descriptor(52        TILE_M, TILE_N, static_cast<int>(dynamic_extent),53        /*transpose_left=*/TA, /*transpose_right=*/TB,54        /*relaxed_precision=*/false, matmul2d_descriptor::mode::multiply);5556    matmul2d<desc, execution_simdgroups<4>> op;5758    // slice() is (column, row); a transposed operand is indexed the other way.59    auto mA = TA ? A.slice(int(tgid.y) * TILE_M, 0) : A.slice(0, int(tgid.y) * TILE_M);60    auto mB = TB ? B.slice(0, int(tgid.x) * TILE_N) : B.slice(int(tgid.x) * TILE_N, 0);61    auto mC = C.slice(int(tgid.x) * TILE_N, int(tgid.y) * TILE_M);6263    // The cooperative tensor's element->lane distribution is implementation64    // defined and not every slot a thread holds is valid, hence the65    // is_valid_element guard. (The header's own example still says get_mask,66    // which does not exist in this SDK.)67    auto cT = op.template get_destination_cooperative_tensor<decltype(mA), decltype(mB), float>();68#pragma clang loop unroll(full)69    for (uint16_t i = 0; i < cT.get_capacity(); ++i)70        if (cT.is_valid_element(i)) cT[i] = 0.0f;7172    op.run(mA, mB, cT);73    cT.store(mC);74}7576#define INST_MPP(SUF, T, TA, TB)                                          \77    template [[host_name("matmul_mpp_" #SUF)]] kernel void                \78    matmul_mpp<T, TA, TB>(device T*, device T*, device float*,            \79                          constant MPPParams&, uint2);8081INST_MPP(f32, float, false, false)82INST_MPP(f32_nt, float, false, true)83INST_MPP(f32_tn, float, true, false)84INST_MPP(f16, half, false, false)85INST_MPP(f16_nt, half, false, true)86INST_MPP(f16_tn, half, true, false)87