nuvec-gpu 0.0.1

//! The Metal runtime end to end. Every test skips when no Metal device is available.

#![cfg(target_os = "macos")]

use nuvec_gpu::{Arg, Buffer, Device, GpuError};

fn device() -> Option<Device> {
    let device = Device::system_default();
    if device.is_none() {
        eprintln!("skipped: no Metal device");
    }
    device
}

fn f32s(values: &[f32]) -> Vec<u8> {
    values.iter().flat_map(|v| v.to_le_bytes()).collect()
}

fn read_f32(buffer: &Buffer) -> Vec<f32> {
    buffer
        .contents()
        .unwrap()
        .as_chunks::<4>()
        .0
        .iter()
        .map(|c| f32::from_le_bytes(*c))
        .collect()
}

/// `out = a + b` over `n` floats.
fn add(
    cmd: &mut nuvec_gpu::CommandBuffer,
    device: &Device,
    a: &Buffer,
    b: &Buffer,
    out: &mut Buffer,
    n: u32,
) {
    let pipeline = device.pipeline("nuvec_add_f32").unwrap();
    cmd.dispatch(
        &pipeline,
        &[
            Arg::read(a),
            Arg::read(b),
            Arg::write(out),
            Arg::bytes(&n.to_le_bytes()),
        ],
        n as usize,
    )
    .unwrap();
}

#[test]
fn device_info() {
    let Some(device) = device() else { return };
    let info = device.info();
    assert!(!info.name.is_empty());
    assert!(
        info.architecture.starts_with("applegpu_g") || info.generation.is_none(),
        "{info:?}"
    );
    assert!(info.recommended_working_set > 0);
    assert!(info.max_threads_per_threadgroup >= 256);
}

#[test]
fn placeholder_kernel() {
    let Some(device) = device() else { return };
    let pipeline = device.pipeline("nuvec_placeholder").unwrap();
    let mut out = device.new_buffer(1000 * 4).unwrap();
    let mut cmd = device.command_buffer().unwrap();
    cmd.dispatch(&pipeline, &[Arg::write(&mut out)], 1000)
        .unwrap();
    cmd.commit().wait().unwrap();
    assert_eq!(
        read_f32(&out),
        (0..1000).map(|i| 2.0 * i as f32).collect::<Vec<_>>()
    );
}

#[test]
fn add_one_million() {
    let Some(device) = device() else { return };
    let n = 1_000_000u32;
    let a = device
        .buffer_from(&f32s(&(0..n).map(|i| i as f32).collect::<Vec<_>>()))
        .unwrap();
    let b = device.buffer_from(&f32s(&vec![0.5; n as usize])).unwrap();
    let mut out = device.new_buffer(n as usize * 4).unwrap();
    let mut cmd = device.command_buffer().unwrap();
    add(&mut cmd, &device, &a, &b, &mut out, n);
    cmd.commit().wait().unwrap();
    let got = read_f32(&out);
    assert!(got.iter().enumerate().all(|(i, &v)| v == i as f32 + 0.5));
}

#[test]
fn pipelines_are_cached() {
    let Some(device) = device() else { return };
    let first = device.pipeline("nuvec_add_f32").unwrap();
    let second = device.pipeline("nuvec_add_f32").unwrap();
    assert!(first.same_as(&second));
    assert_eq!(first.name(), "nuvec_add_f32");
    assert!(first.execution_width() >= 1);
}

#[test]
fn barriers_only_for_dependencies() {
    let Some(device) = device() else { return };
    let n = 4096u32;
    let one = device.buffer_from(&f32s(&vec![1.0; n as usize])).unwrap();

    // A chain x[i+1] = x[i] + 1: every dispatch after the first depends on its predecessor.
    let mut xs: Vec<Buffer> = (0..101)
        .map(|_| device.new_buffer(n as usize * 4).unwrap())
        .collect();
    let mut cmd = device.command_buffer().unwrap();
    for i in 0..100 {
        let (done, rest) = xs.split_at_mut(i + 1);
        add(&mut cmd, &device, &done[i], &one, &mut rest[0], n);
    }
    assert_eq!(cmd.barriers(), 99);
    cmd.commit().wait().unwrap();
    assert!(
        read_f32(&xs[100]).iter().all(|&v| v == 100.0),
        "a missing barrier would read stale data"
    );

    // Independent dispatches: no barriers.
    let mut outs: Vec<Buffer> = (0..100)
        .map(|_| device.new_buffer(n as usize * 4).unwrap())
        .collect();
    let mut cmd = device.command_buffer().unwrap();
    for out in &mut outs {
        add(&mut cmd, &device, &one, &one, out, n);
    }
    assert_eq!(cmd.barriers(), 0);
    cmd.commit().wait().unwrap();
    assert!(outs.iter().all(|o| read_f32(o).iter().all(|&v| v == 2.0)));
}

#[test]
fn write_after_read_needs_a_barrier() {
    let Some(device) = device() else { return };
    let n = 1024u32;
    let mut a = device.buffer_from(&f32s(&vec![1.0; n as usize])).unwrap();
    let c = device.buffer_from(&f32s(&vec![5.0; n as usize])).unwrap();
    let mut b = device.new_buffer(n as usize * 4).unwrap();
    let mut cmd = device.command_buffer().unwrap();
    add(&mut cmd, &device, &a, &a, &mut b, n); // reads a
    add(&mut cmd, &device, &c, &c, &mut a, n); // then overwrites a
    assert_eq!(cmd.barriers(), 1);
    cmd.commit().wait().unwrap();
    assert!(
        read_f32(&b).iter().all(|&v| v == 2.0),
        "b must see the old a"
    );
    assert!(read_f32(&a).iter().all(|&v| v == 10.0));
}

#[test]
fn misuse_is_reported() {
    let Some(device) = device() else { return };
    assert_eq!(
        device.pipeline("no_such_kernel").unwrap_err(),
        GpuError::MissingKernel("no_such_kernel".into())
    );
    let pipeline = device.pipeline("nuvec_add_f32").unwrap();
    let small = device.new_buffer(16).unwrap();
    let mut cmd = device.command_buffer().unwrap();
    let past_end = cmd.dispatch(&pipeline, &[Arg::Read(&small, 17)], 4);
    assert!(matches!(past_end, Err(GpuError::InvalidArgument(_))));
    let bytes = [0u8; 4];
    let many: Vec<Arg<'_>> = (0..32).map(|_| Arg::bytes(&bytes)).collect();
    assert!(matches!(
        cmd.dispatch(&pipeline, &many, 4),
        Err(GpuError::InvalidArgument(_))
    ));
    let big = vec![0u8; 5000];
    assert!(matches!(
        cmd.dispatch(&pipeline, &[Arg::bytes(&big)], 4),
        Err(GpuError::InvalidArgument(_))
    ));
    assert_eq!(cmd.barriers(), 0, "rejected dispatches encode nothing");
}

#[test]
fn cpu_access_waits_for_the_gpu() {
    let Some(device) = device() else { return };
    let n = 1 << 20;
    let a = device.buffer_from(&f32s(&vec![1.0; n as usize])).unwrap();
    let mut out = device.new_buffer(n as usize * 4).unwrap();
    let mut cmd = device.command_buffer().unwrap();
    add(&mut cmd, &device, &a, &a, &mut out, n);
    let submitted = cmd.commit();
    assert!(out.in_use());
    assert_eq!(out.contents().unwrap_err(), GpuError::BufferInUse);
    assert_eq!(
        a.contents().unwrap_err(),
        GpuError::BufferInUse,
        "inputs are guarded too"
    );
    drop(submitted); // waits
    assert!(!out.in_use());
    assert_eq!(read_f32(&out)[0], 2.0);
}

#[test]
fn abandoned_command_buffer_releases_buffers() {
    let Some(device) = device() else { return };
    let a = device.buffer_from(&f32s(&[1.0; 4])).unwrap();
    let mut out = device.new_buffer(16).unwrap();
    let mut cmd = device.command_buffer().unwrap();
    add(&mut cmd, &device, &a, &a, &mut out, 4);
    drop(cmd); // never submitted
    assert!(!out.in_use() && !a.in_use());
    assert!(out.contents().is_ok());
}

#[test]
fn empty_buffers_and_zero_threads() {
    let Some(device) = device() else { return };
    let empty = device.new_buffer(0).unwrap();
    assert!(empty.is_empty() && empty.contents().unwrap().is_empty());
    let pipeline = device.pipeline("nuvec_add_f32").unwrap();
    let mut cmd = device.command_buffer().unwrap();
    cmd.dispatch(&pipeline, &[], 0).unwrap();
    cmd.commit().wait().unwrap();
}