//! bf16-native tiled GEMM via `simdgroup_matrix<bfloat,8,8>`
//! cooperative matmul — the Metal-4 / M5-Neural-Accelerator
//! path. Weights stay bf16 (no f16 conversion); activations
//! widen f32→bf16 once per call; accumulation is f32.
//!
//! Computes C[m,n] = A[m,k] · B[n,k]ᵀ (B is the weight in its
//! natural transposed `[n,k]` layout, same as gemm_bt). It
//! verifies against a CPU bf16-input reference and times the
//! kernel against the classic f16 `MPSMatrix` resident path.
//!
//! Run on an M5 (Metal 4) for the numbers that matter; it
//! also runs on M3 (simdgroup bf16 is Metal 3+) as a control.
#[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,
    };

    // f32 → bf16 bits, round-to-nearest-even.
    fn f32_to_bf16(v: f32) -> u16 {
        let b = v.to_bits();
        let round = ((b >> 16) & 1) + 0x7fff;
        ((b + round) >> 16) as u16
    }
    fn bf16_to_f32(h: u16) -> f32 {
        f32::from_bits((h as u32) << 16)
    }

    const KERNEL: &str = r#"
#include <metal_stdlib>
using namespace metal;
// One simdgroup (32 threads) per 8×8 output tile.
kernel void gemm_bf16_bt(
    device const bfloat* A [[buffer(0)]],   // [M, K] bf16 (row-major)
    device const bfloat* B [[buffer(1)]],   // [N, K] bf16 (weights, bt)
    device float*        C [[buffer(2)]],   // [M, N] f32
    constant uint&       M [[buffer(3)]],
    constant uint&       N [[buffer(4)]],
    constant uint&       K [[buffer(5)]],
    uint2 tg [[threadgroup_position_in_grid]])
{
    uint mi = tg.y * 8;
    uint nj = tg.x * 8;
    if (mi >= M || nj >= N) return;
    simdgroup_matrix<float, 8, 8> acc = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
    for (uint kt = 0; kt < K; kt += 8) {
        simdgroup_matrix<bfloat, 8, 8> a, b;
        simdgroup_load(a, A + mi * K + kt, K);
        // transpose=true: b[l,j] = B[nj+j, kt+l] → C = A·Bᵀ
        simdgroup_load(b, B + nj * K + kt, K, 0, true);
        simdgroup_multiply_accumulate(acc, a, b, acc);
    }
    simdgroup_store(acc, C + mi * N + nj, N);
}
"#;

    let device = MTLCreateSystemDefaultDevice().expect("no Metal device");
    println!("device: {}", device.name());
    let queue = device.newCommandQueue().expect("no queue");

    let src = NSString::from_str(KERNEL);
    let lib = device
        .newLibraryWithSource_options_error(&src, None)
        .expect("kernel compile failed");
    let func = lib
        .newFunctionWithName(&NSString::from_str("gemm_bf16_bt"))
        .expect("no fn");
    let pipe: Retained<ProtocolObject<dyn MTLComputePipelineState>> = device
        .newComputePipelineStateWithFunction_error(&func)
        .expect("no pipeline");

    // Prefill-ish shapes (multiples of 8). m = batch tokens.
    let shapes = [(64usize, 512usize, 512usize), (256, 2048, 2048), (512, 4096, 4096)];

    for &(m, n, k) in &shapes {
        // Synthetic A[m,k], B[n,k].
        let a_f: Vec<f32> =
            (0..m * k).map(|i| (((i % 17) as f32) - 8.0) * 0.05).collect();
        let b_f: Vec<f32> =
            (0..n * k).map(|i| (((i % 13) as f32) - 6.0) * 0.04).collect();
        let a_bf: Vec<u16> = a_f.iter().map(|&v| f32_to_bf16(v)).collect();
        let b_bf: Vec<u16> = b_f.iter().map(|&v| f32_to_bf16(v)).collect();

        let mk_buf = |data: &[u16]| unsafe {
            device
                .newBufferWithBytes_length_options(
                    std::ptr::NonNull::new(data.as_ptr() as *mut _).unwrap(),
                    std::mem::size_of_val(data),
                    MTLResourceOptions::StorageModeShared,
                )
                .unwrap()
        };
        let buf_a = mk_buf(&a_bf);
        let buf_b = mk_buf(&b_bf);
        let buf_c = device
            .newBufferWithLength_options(m * n * 4, MTLResourceOptions::StorageModeShared)
            .unwrap();

        let dispatch = || 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);
            for (slot, &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 + slot,
                );
            }
            enc.dispatchThreadgroups_threadsPerThreadgroup(
                MTLSize { width: n / 8, height: m / 8, depth: 1 },
                // one simdgroup = 32 threads
                MTLSize { width: 32, height: 1, depth: 1 },
            );
            enc.endEncoding();
            cb.commit();
            cb.waitUntilCompleted();
        };

        // Correctness: CPU reference over bf16-rounded inputs.
        dispatch();
        let c_ptr = MTLBuffer::contents(&*buf_c).as_ptr() as *const f32;
        let c_gpu: Vec<f32> = (0..m * n).map(|i| unsafe { *c_ptr.add(i) }).collect();
        let mut max_rel = 0f32;
        for i in 0..m.min(8) {
            for j in 0..n.min(8) {
                let mut acc = 0f32;
                for l in 0..k {
                    acc += bf16_to_f32(a_bf[i * k + l]) * bf16_to_f32(b_bf[j * k + l]);
                }
                let g = c_gpu[i * n + j];
                let rel = (g - acc).abs() / (acc.abs() + 1e-3);
                max_rel = max_rel.max(rel);
            }
        }

        // Time it.
        let iters = 50;
        for _ in 0..5 {
            dispatch();
        }
        let t0 = Instant::now();
        for _ in 0..iters {
            dispatch();
        }
        let bf16_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:<5} k={k:<5}  bf16-simdgroup {bf16_ms:6.3} ms  \
             {:7.1} GFLOP/s  max_rel={max_rel:.2e}",
            gflop / (bf16_ms / 1e3),
        );
    }
}

#[cfg(not(target_os = "macos"))]
fn main() {
    println!("macOS only");
}