spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2// Micro-benchmark: does simdgroup_matrix run faster with f16/bf16 operands3// than with f32 on this hardware? Same 64x64x16 tiling as the real GEMM;4// only the operand type of the staged tiles and the MMA changes. The5// accumulator stays f32 in every variant (required for training).6#include <metal_stdlib>7#include <metal_simdgroup_matrix>8using namespace metal;910enum : uint { BM = 64, BN = 64, BK = 16, WM = 2, WN = 2,11              NSG = WM*WN, THREADS = NSG*32, TM = BM/(8*WM), TN = BN/(8*WN),12              PAD = 8, LDA_S = BK+PAD, LDB_S = BN+PAD };1314struct P { uint M, N, K; };1516// T = operand precision, ACC fragment always f32.17template <typename T, typename FRAG>18kernel void gemm_prec(device const T* A [[buffer(0)]],19                      device const T* B [[buffer(1)]],20                      device float*   C [[buffer(2)]],21                      constant P& p [[buffer(3)]],22                      uint2 tgid [[threadgroup_position_in_grid]],23                      uint  tid  [[thread_index_in_threadgroup]],24                      uint  sgid [[simdgroup_index_in_threadgroup]]) {25    threadgroup T As[BM * LDA_S];26    threadgroup T Bs[BK * LDB_S];27    const uint row0 = tgid.y * BM, col0 = tgid.x * BN;28    const uint sr = (sgid / WN) * (TM * 8), sc = (sgid % WN) * (TN * 8);2930    simdgroup_float8x8 acc[TM][TN];31#pragma clang loop unroll(full)32    for (uint i = 0; i < TM; ++i)33#pragma clang loop unroll(full)34        for (uint j = 0; j < TN; ++j) acc[i][j] = make_filled_simdgroup_matrix<float,8,8>(0.0f);3536    for (uint k0 = 0; k0 < p.K; k0 += BK) {37        threadgroup_barrier(mem_flags::mem_threadgroup);38        for (uint e = tid; e < BM*BK; e += THREADS) {39            const uint i = e / BK, k = e % BK;40            As[i*LDA_S + k] = A[(row0+i)*p.K + k0+k];41        }42        for (uint e = tid; e < BK*BN; e += THREADS) {43            const uint k = e / BN, j = e % BN;44            Bs[k*LDB_S + j] = B[(k0+k)*p.N + col0+j];45        }46        threadgroup_barrier(mem_flags::mem_threadgroup);47#pragma clang loop unroll(full)48        for (uint kk = 0; kk < BK; kk += 8) {49            FRAG af[TM], bf[TN];50#pragma clang loop unroll(full)51            for (uint i = 0; i < TM; ++i) simdgroup_load(af[i], As + (sr+i*8)*LDA_S + kk, LDA_S);52#pragma clang loop unroll(full)53            for (uint j = 0; j < TN; ++j) simdgroup_load(bf[j], Bs + kk*LDB_S + sc + j*8, LDB_S);54#pragma clang loop unroll(full)55            for (uint i = 0; i < TM; ++i)56#pragma clang loop unroll(full)57                for (uint j = 0; j < TN; ++j)58                    simdgroup_multiply_accumulate(acc[i][j], af[i], bf[j], acc[i][j]);59        }60    }61#pragma clang loop unroll(full)62    for (uint i = 0; i < TM; ++i)63#pragma clang loop unroll(full)64        for (uint j = 0; j < TN; ++j)65            simdgroup_store(acc[i][j], C + (row0+sr+i*8)*p.N + col0+sc+j*8, p.N);66}6768template [[host_name("gemm_f32")]] kernel void69gemm_prec<float, simdgroup_float8x8>(device const float*, device const float*,70    device float*, constant P&, uint2, uint, uint);71template [[host_name("gemm_f16")]] kernel void72gemm_prec<half, simdgroup_half8x8>(device const half*, device const half*,73    device float*, constant P&, uint2, uint, uint);74template [[host_name("gemm_bf16")]] kernel void75gemm_prec<bfloat, simdgroup_bfloat8x8>(device const bfloat*, device const bfloat*,76    device float*, constant P&, uint2, uint, uint);77