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

//! matmul2d replicating the liuliu/example_matmul_metal4 recipe:
//! A[K,M], B[N,K], C[M,N], no transpose, K = dynamic_length_v<int>
//! (NOT dynamic_extent — that caps K at one tile), plain slice(),
//! single run(). Full-K test: A column m=3 all ones, B all ones
//! => C[3,n] = K for all n; rest 0.
#[cfg(target_os = "macos")]
fn main() {
    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 }

    const KERNEL: &str = r#"
#include <metal_stdlib>
#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
using namespace metal;
using namespace mpp::tensor_ops;
// A[K,M]=[256,128], B[N,K]=[64,256], C[M,N]=[128,64].
kernel void k(device bfloat* A_buf [[buffer(0)]],
              device bfloat* B_buf [[buffer(1)]],
              device float*  C_buf [[buffer(2)]],
              uint2 tgid [[threadgroup_position_in_grid]]) {
    auto A = tensor<device bfloat, dextents<int32_t,2>, tensor_inline>(A_buf, dextents<int32_t,2>(256, 128));
    auto B = tensor<device bfloat, dextents<int32_t,2>, tensor_inline>(B_buf, dextents<int32_t,2>(64, 256));
    auto C = tensor<device float,  dextents<int32_t,2>, tensor_inline>(C_buf, dextents<int32_t,2>(128, 64));
    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("k")).unwrap();
    let pipe: Retained<ProtocolObject<dyn MTLComputePipelineState>> =
        device.newComputePipelineStateWithFunction_error(&func).unwrap();
    let (m, n, k) = (128usize, 64usize, 256usize);
    // A[K,M]: a_bf[kk*M+mm]=1 iff mm==3 (column m=3 all ones).
    let a_bf: Vec<u16> = (0..k * m).map(|_| f32_to_bf16(1.0)).collect();
    let b_bf: Vec<u16> = (0..n * k).map(|i| f32_to_bf16(if i == 200 { 1.0 } else { 0.0 })).collect();
    let mk = |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 = mk(&a_bf);
    let buf_b = mk(&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;
    unsafe { std::ptr::write_bytes(cp, 0, m * n) };
    unsafe {
        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);
        // grid: N/32 x M/64 = 2 x 2
        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();
    }
    let c: Vec<f32> = (0..m * n).map(|i| unsafe { *cp.add(i) }).collect();
    println!("single-element B_buf[200]=1, A all ones. [N,K]:n=200/256,k=200%256; [K,N]:k=200/64,n=200%64");
    let nzv: Vec<(usize, i32)> = c.iter().enumerate().filter(|(_, v)| v.abs() > 0.5).map(|(i, v)| (i, *v as i32)).collect();
    println!("nonzeros: {}", nzv.len());
    for (i, v) in nzv.iter().take(8) {
        println!("  c[{i}] = {v}   ([M,N]: m={},n={})  ([N,M]: n={},m={})", i / n, i % n, i / m, i % m);
    }
}
#[cfg(not(target_os = "macos"))]
fn main() {}