open-standards/jouleclaw/jouleclaw-rs/crates/jouleclaw-backend-apple/examples/mpp_gemm_bench.rs

//! M5 GPU Neural-Accelerator GEMM via Metal-4 Metal-Performance-
//! Primitives `mpp::tensor_ops::matmul2d`, using the canonical
//! recipe from liuliu/example_matmul_metal4. WORKING: bf16 in,
//! f32 accum, standard C[M,N]=A[M,K]·B[K,N] (no transpose).
//!
//! The unblock was K = `dynamic_length_v<int>` in the descriptor
//! — NOT `dynamic_extent`/0, which silently caps the K-sum at
//! one ~32-wide tile (the bug that defeated every earlier
//! attempt). Plain `slice()` (origin shift), single `run()`
//! that loops K internally, `execution_simdgroups<4>`,
//! descriptor tile (64, 32). Constraints: M%64, N%32 (N must be
//! ≥128 — N=64 mis-tiles), K%32.
//!
//! VERIFIED CORRECT (max_abs 0.0) at all shapes except N=64.
//! On M5 Max: 24.6 TFLOP/s at 2048×2048×4096 vs the bf16
//! simdgroup path's ~4.8 TFLOP/s — ~5×, both bit-correct.
#[cfg(target_os = "macos")]
fn main() {
    use std::time::Instant;
    use objc2::rc::Retained;
    use objc2::runtime::ProtocolObject;
    use objc2_foundation::NSString;
    use objc2_metal::{
        MTLBuffer, MTLCommandBuffer, MTLCommandEncoder, MTLCommandQueue,
        MTLComputeCommandEncoder, MTLComputePipelineState, MTLCreateSystemDefaultDevice,
        MTLDevice, MTLLibrary, MTLResourceOptions, MTLSize,
    };
    fn f32_to_bf16(v: f32) -> u16 { let b = v.to_bits(); ((b + (((b >> 16) & 1) + 0x7fff)) >> 16) as u16 }
    fn bf16_to_f32(h: u16) -> f32 { f32::from_bits((h as u32) << 16) }

    const KERNEL: &str = r#"
#include <metal_stdlib>
#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
using namespace metal;
using namespace mpp::tensor_ops;
kernel void gemm_bf16_mpp(device bfloat* A_buf [[buffer(0)]],   // [K,M]
                          device bfloat* B_buf [[buffer(1)]],   // [N,K]
                          device float*  C_buf [[buffer(2)]],   // [M,N]
                          constant uint& M [[buffer(3)]],
                          constant uint& N [[buffer(4)]],
                          constant uint& K [[buffer(5)]],
                          uint2 tgid [[threadgroup_position_in_grid]]) {
    auto A = tensor<device bfloat, dextents<int32_t,2>, tensor_inline>(A_buf, dextents<int32_t,2>(K, M));
    auto B = tensor<device bfloat, dextents<int32_t,2>, tensor_inline>(B_buf, dextents<int32_t,2>(N, K));
    auto C = tensor<device float,  dextents<int32_t,2>, tensor_inline>(C_buf, dextents<int32_t,2>(M, N));
    constexpr auto desc = matmul2d_descriptor(64, 32, dynamic_length_v<int>, false, false, false,
                                              matmul2d_descriptor::mode::multiply);
    matmul2d<desc, execution_simdgroups<4>> mm;
    auto mA = A.slice(0, tgid.y * 64);
    auto mB = B.slice(tgid.x * 32, 0);
    auto mC = C.slice(tgid.x * 32, tgid.y * 64);
    mm.run(mA, mB, mC);
}
"#;
    let device = MTLCreateSystemDefaultDevice().expect("no device");
    println!("device: {}", device.name());
    let queue = device.newCommandQueue().unwrap();
    let lib = match device.newLibraryWithSource_options_error(&NSString::from_str(KERNEL), None) {
        Ok(l) => l, Err(e) => { println!("COMPILE FAIL:\n{}", e.localizedDescription()); return; }
    };
    let func = lib.newFunctionWithName(&NSString::from_str("gemm_bf16_mpp")).unwrap();
    let pipe: Retained<ProtocolObject<dyn MTLComputePipelineState>> =
        device.newComputePipelineStateWithFunction_error(&func).unwrap();

    // M%64==0, N%32==0, K%32==0.
    let shapes = [(256usize, 256usize, 512usize), (512, 512, 2048), (1024, 1024, 4096), (2048, 2048, 4096)];
    for &(m, n, k) in &shapes {
        let a_bf: Vec<u16> = (0..k * m).map(|i| f32_to_bf16((((i % 17) as f32) - 8.0) * 0.05)).collect();
        let b_bf: Vec<u16> = (0..n * k).map(|i| f32_to_bf16((((i % 13) as f32) - 6.0) * 0.04)).collect();
        let mkb = |d: &[u16]| unsafe { device.newBufferWithBytes_length_options(
            std::ptr::NonNull::new(d.as_ptr() as *mut _).unwrap(), std::mem::size_of_val(d),
            MTLResourceOptions::StorageModeShared).unwrap() };
        let buf_a = mkb(&a_bf);
        let buf_b = mkb(&b_bf);
        let buf_c = device.newBufferWithLength_options(m * n * 4, MTLResourceOptions::StorageModeShared).unwrap();
        let cp = MTLBuffer::contents(&*buf_c).as_ptr() as *mut f32;
        let dispatch = || unsafe {
            std::ptr::write_bytes(cp, 0, m * n);
            let cb = queue.commandBuffer().unwrap();
            let enc = cb.computeCommandEncoder().unwrap();
            enc.setComputePipelineState(&pipe);
            enc.setBuffer_offset_atIndex(Some(&buf_a), 0, 0);
            enc.setBuffer_offset_atIndex(Some(&buf_b), 0, 1);
            enc.setBuffer_offset_atIndex(Some(&buf_c), 0, 2);
            for (s, &v) in [m as u32, n as u32, k as u32].iter().enumerate() {
                enc.setBytes_length_atIndex(std::ptr::NonNull::new(&v as *const u32 as *mut _).unwrap(), 4, 3 + s);
            }
            enc.dispatchThreadgroups_threadsPerThreadgroup(
                MTLSize { width: n / 32, height: m / 64, depth: 1 },
                MTLSize { width: 128, height: 1, depth: 1 });
            enc.endEncoding();
            cb.commit();
            cb.waitUntilCompleted();
        };
        dispatch();
        let c: Vec<f32> = (0..m * n).map(|i| unsafe { *cp.add(i) }).collect();
        let mut max_abs = 0f32; let mut max_mag = 0f32; let mut wrong = 0usize;
        // FULL-output check: every row, every column (not a corner).
        for mm_ in 0..m { for nn in 0..n {
            let mut acc = 0f32;
            for l in 0..k { acc += bf16_to_f32(a_bf[mm_ * k + l]) * bf16_to_f32(b_bf[l * n + nn]); }
            let d = (c[mm_ * n + nn] - acc).abs();
            max_abs = max_abs.max(d);
            max_mag = max_mag.max(acc.abs());
            if d > 1e-2 { wrong += 1; }
        }}
        let iters = 50;
        for _ in 0..5 { dispatch(); }
        let t0 = Instant::now();
        for _ in 0..iters { dispatch(); }
        let ms = t0.elapsed().as_secs_f64() * 1e3 / iters as f64;
        let gflop = 2.0 * m as f64 * n as f64 * k as f64 / 1e9;
        println!("  m={m:<4} n={n:<4} k={k:<4}  {ms:6.3} ms  {:8.1} GFLOP/s  max_abs={max_abs:.2e}  WRONG={wrong}/{} (mag≈{max_mag:.2e})",
                 gflop / (ms / 1e3), m*n);
    }
}
#[cfg(not(target_os = "macos"))]
fn main() {}