spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// Does simdgroup_matrix actually run faster with f16/bf16 operands than with4// f32 on this hardware? The literature disagrees: metal-benchmarks measures5// f16 and f32 FMA at the SAME rate on Apple GPUs (the win coming from6// registers/bandwidth), while Apple's M3 material claims up to 2x ALU for7// family 9 via FP16/FP32/INT co-issue. Since mixed precision is a large8// amount of work to plumb through a training framework, measure before9// committing.10//11// All three variants use the identical 64x64x16 tiling from matmul_simd and12// an f32 accumulator (required for training); only the staged-tile operand13// type and the MMA fragment type differ.14#include <Foundation/Foundation.hpp>15#include <Metal/Metal.hpp>1617#include "core/device.h"18#include "core/tensor.h"19#include "ops/metal/metal_ops.h"2021#include <cstdio>22#include <random>23#include <string>2425using namespace forge;2627namespace {2829struct Params { uint32_t M, N, K; };3031double run(const char* kernel, DType dt, int64_t M, int64_t N, int64_t K, int iters) {32    Tensor a = Tensor::empty({M, K}, dt);33    Tensor b = Tensor::empty({K, N}, dt);34    Tensor c = Tensor::empty({M, N}, DType::F32);35    std::mt19937 rng(3);36    std::uniform_real_distribution<float> dist(-1.0f, 1.0f);37    for (int64_t i = 0; i < a.numel(); ++i) a.set_item(i, dist(rng));38    for (int64_t i = 0; i < b.numel(); ++i) b.set_item(i, dist(rng));3940    Device& dev = Device::get();41    MTL::ComputePipelineState* pso = dev.pipeline(kernel);42    Params p{uint32_t(M), uint32_t(N), uint32_t(K)};4344    auto encode = [&]() {45        MTL::ComputeCommandEncoder* enc = metal::Stream::get().encoder();46        enc->setComputePipelineState(pso);47        enc->setBuffer(a.buffer(), a.buffer_offset(), 0);48        enc->setBuffer(b.buffer(), b.buffer_offset(), 1);49        enc->setBuffer(c.buffer(), c.buffer_offset(), 2);50        enc->setBytes(&p, sizeof(p), 3);51        enc->dispatchThreadgroups(52            MTL::Size(NS::UInteger(N) / 64, NS::UInteger(M) / 64, 1),53            MTL::Size(128, 1, 1));54    };5556    encode();57    metal::sync(); // warm up, build pipeline58    const double t0 = metal::Stream::get().gpu_seconds();59    for (int i = 0; i < iters; ++i) encode();60    metal::sync();61    const double dt_s = metal::Stream::get().gpu_seconds() - t0;62    return 2.0 * double(M) * double(N) * double(K) * iters / dt_s / 1e12;63}6465double run_mpp(const char* kernel, DType dt, int64_t M, int64_t N, int64_t K,66               int iters) {67    Tensor a = Tensor::empty({M, K}, dt);68    Tensor b = Tensor::empty({K, N}, dt);69    Tensor c = Tensor::empty({M, N}, DType::F32);70    std::mt19937 rng(3);71    std::uniform_real_distribution<float> dist(-1.0f, 1.0f);72    for (int64_t i = 0; i < a.numel(); ++i) a.set_item(i, dist(rng));73    for (int64_t i = 0; i < b.numel(); ++i) b.set_item(i, dist(rng));7475    Device& dev = Device::get();76    MTL::ComputePipelineState* pso = dev.pipeline(kernel);77    Params p{uint32_t(M), uint32_t(N), uint32_t(K)};7879    auto encode = [&]() {80        MTL::ComputeCommandEncoder* enc = metal::Stream::get().encoder();81        enc->setComputePipelineState(pso);82        enc->setBuffer(a.buffer(), a.buffer_offset(), 0);83        enc->setBuffer(b.buffer(), b.buffer_offset(), 1);84        enc->setBuffer(c.buffer(), c.buffer_offset(), 2);85        enc->setBytes(&p, sizeof(p), 3);86        enc->dispatchThreadgroups(87            MTL::Size(NS::UInteger(N) / 32, NS::UInteger(M) / 64, 1),88            MTL::Size(128, 1, 1));89    };9091    encode();92    metal::sync();93    const double t0 = metal::Stream::get().gpu_seconds();94    for (int i = 0; i < iters; ++i) encode();95    metal::sync();96    const double dt_s = metal::Stream::get().gpu_seconds() - t0;97    return 2.0 * double(M) * double(N) * double(K) * iters / dt_s / 1e12;98}99100} // namespace101102int main() {103    NS::AutoreleasePool* pool = NS::AutoreleasePool::alloc()->init();104    std::printf("device: %s\n\n", Device::get().name().c_str());105    std::printf("simdgroup_matrix throughput by operand precision "106                "(f32 accumulator throughout)\n");107    std::printf("%-16s %10s %10s %10s\n", "shape", "f32", "f16", "bf16");108109    struct Shape { int64_t M, N, K; };110    const Shape shapes[] = {{2048, 2048, 2048}, {4096, 4096, 4096},111                            {65536, 512, 1408}};112    for (const Shape& s : shapes) {113        const double gflop = 2.0 * double(s.M) * double(s.N) * double(s.K) / 1e9;114        const int iters = gflop > 50.0 ? 3 : 10;115        const double f32 = run("gemm_f32", DType::F32, s.M, s.N, s.K, iters);116        const double f16 = run("gemm_f16", DType::F16, s.M, s.N, s.K, iters);117        const double bf16 = run("gemm_bf16", DType::BF16, s.M, s.N, s.K, iters);118        char tag[32];119        std::snprintf(tag, sizeof(tag), "%lldx%lldx%lld", (long long)s.M,120                      (long long)s.N, (long long)s.K);121        std::printf("%-16s %9.2fT %9.2fT %9.2fT\n", tag, f32, f16, bf16);122        std::fflush(stdout);123    }124125    // Metal Performance Primitives cooperative-tensor matmul2d — the Metal 4126    // path that targets M5 neural accelerators. 64x32 tile, 4 simdgroups.127    std::printf("\nMPP matmul2d (cooperative tensors, Metal 4)\n");128    std::printf("%-16s %10s %10s\n", "shape", "f32", "f16");129    for (const Shape& s : shapes) {130        if (s.M % 64 != 0 || s.N % 32 != 0) continue;131        const double gflop = 2.0 * double(s.M) * double(s.N) * double(s.K) / 1e9;132        const int iters = gflop > 50.0 ? 3 : 10;133        const double f32 = run_mpp("matmul_mpp_f32", DType::F32, s.M, s.N, s.K, iters);134        const double f16 = run_mpp("matmul_mpp_f16", DType::F16, s.M, s.N, s.K, iters);135        char tag[32];136        std::snprintf(tag, sizeof(tag), "%lldx%lldx%lld", (long long)s.M,137                      (long long)s.N, (long long)s.K);138        std::printf("%-16s %9.2fT %9.2fT\n", tag, f32, f16);139        std::fflush(stdout);140    }141142    pool->drain();143    return 0;144}145