open-standards/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/tests/metal_kernels.rs

//! Kernels in `sources::PHASE27_OPS` and the matvec pair, executed on a GPU.
//!
//! Every shader in this crate is built at runtime by `compile_pipeline`, and
//! until this file existed nothing ever dispatched one: `cargo test --features
//! metal` ran exactly the same 46 tests as the default build. Compiling a
//! kernel says nothing about whether it computes the right thing, so each test
//! here checks a dispatch against a CPU reference, or against the kernel that
//! is supposed to invert it.
//!
//! These need a real Metal device. They do not skip when one is missing --
//! a test that passes by doing nothing is worse than no test.

#![cfg(feature = "metal")]
// The CPU reference loops walk an output by sample index while reading the
// delay line at a different index (n % delay_len), which is the whole point
// of the recurrence; iterator form cannot express it.
#![allow(clippy::needless_range_loop)]

use half::f16;
use jouleclaw_omni::hal::metal::shader::sources;
use jouleclaw_omni::hal::metal::{MetalCompute, MetalDevice};
use jouleclaw_omni::inference::gpu_ops;
use jouleclaw_omni::tensor::{DType, Shape, Tensor};
use std::sync::Arc;

fn compute() -> MetalCompute {
    let device = MetalDevice::new().expect("a Metal device is required by these tests");
    MetalCompute::new(Arc::new(device))
}

/// Deterministic pseudo-random values in roughly [-1, 1].
fn ramp(n: usize, seed: u32) -> Vec<f16> {
    (0..n)
        .map(|i| {
            let x = (i as u32).wrapping_mul(2654435761).wrapping_add(seed);
            let u = ((x >> 8) & 0xffff) as f32 / 32768.0 - 1.0;
            f16::from_f32(u)
        })
        .collect()
}

fn tensor(c: &MetalCompute, data: &[f16], shape: impl Into<Shape>) -> Tensor {
    Tensor::from_slice(data, shape, DType::F16, c.device().info().id).expect("tensor")
}

fn read(t: &Tensor) -> Vec<f32> {
    t.to_vec::<f16>()
        .expect("read back")
        .iter()
        .map(|v| v.to_f32())
        .collect()
}

#[track_caller]
fn assert_close(got: &[f32], want: &[f32], tol: f32, what: &str) {
    assert_eq!(got.len(), want.len(), "{what}: length");
    for (i, (g, w)) in got.iter().zip(want).enumerate() {
        let scale = w.abs().max(1.0);
        assert!(
            (g - w).abs() <= tol * scale,
            "{what}: index {i} got {g}, want {w} (tol {tol})"
        );
    }
}

// ─── shape plumbing: an inverse pair is its own reference ───────────────────

#[test]
fn window_partition_then_unpartition_is_identity() {
    let c = compute();
    let part = c
        .compile_pipeline("t_wp", sources::PHASE27_OPS, "window_partition_f16")
        .expect("window_partition_f16");
    let unpart = c
        .compile_pipeline("t_wu", sources::PHASE27_OPS, "window_unpartition_f16")
        .expect("window_unpartition_f16");

    let (h, w, dim, wh, ww) = (8usize, 12usize, 5usize, 4usize, 4usize);
    let src = ramp(h * w * dim, 11);
    let x = tensor(&c, &src, [h * w, dim]);

    let cb = c.new_command_buffer();
    let windows = gpu_ops::window_partition_on(&c, &part, cb.as_ref(), &x, h, w, dim, wh, ww);
    cb.commit();
    cb.wait_until_completed();

    assert_eq!(windows.shape().dims(), &[(h / wh) * (w / ww), wh * ww, dim]);

    let cb = c.new_command_buffer();
    let back = gpu_ops::window_unpartition_on(&c, &unpart, cb.as_ref(), &windows, h, w, dim, wh, ww);
    cb.commit();
    cb.wait_until_completed();

    let want: Vec<f32> = src.iter().map(|v| v.to_f32()).collect();
    assert_close(&read(&back), &want, 0.0, "window round trip");
}

#[test]
fn patchify_3d_then_unpatchify_is_identity() {
    let c = compute();
    let pat = c
        .compile_pipeline("t_p3", sources::PHASE27_OPS, "patchify_3d_f16")
        .expect("patchify_3d_f16");
    let unpat = c
        .compile_pipeline("t_u3", sources::PHASE27_OPS, "unpatchify_3d_f16")
        .expect("unpatchify_3d_f16");

    let (ch, t, h, w, pt, ph, pw) = (3usize, 2usize, 8usize, 6usize, 1usize, 2usize, 2usize);
    let src = ramp(ch * t * h * w, 23);
    let x = tensor(&c, &src, [ch, t, h, w]);

    let cb = c.new_command_buffer();
    let patches = gpu_ops::patchify_3d_on(&c, &pat, cb.as_ref(), &x, ch, t, h, w, pt, ph, pw);
    cb.commit();
    cb.wait_until_completed();

    let num_patches = (t / pt) * (h / ph) * (w / pw);
    assert_eq!(patches.shape().dims(), &[num_patches, ch * pt * ph * pw]);

    let cb = c.new_command_buffer();
    let back = gpu_ops::unpatchify_3d_on(&c, &unpat, cb.as_ref(), &patches, ch, t, h, w, pt, ph, pw);
    cb.commit();
    cb.wait_until_completed();

    let want: Vec<f32> = src.iter().map(|v| v.to_f32()).collect();
    assert_close(&read(&back), &want, 0.0, "patchify round trip");
}

#[test]
fn split_qkv_separates_the_three_projections() {
    let c = compute();
    let k = c
        .compile_pipeline("t_sq", sources::PHASE27_OPS, "split_qkv_f16")
        .expect("split_qkv_f16");

    let (seq, dim) = (7usize, 6usize);
    let src = ramp(seq * 3 * dim, 31);
    let qkv = tensor(&c, &src, [seq, 3 * dim]);

    let cb = c.new_command_buffer();
    let (q, kk, v) = gpu_ops::split_qkv_on(&c, &k, cb.as_ref(), &qkv, seq, dim).expect("split");
    cb.commit();
    cb.wait_until_completed();

    for (which, got) in [(0usize, &q), (1, &kk), (2, &v)] {
        let mut want = Vec::with_capacity(seq * dim);
        for s in 0..seq {
            for d in 0..dim {
                want.push(src[s * 3 * dim + which * dim + d].to_f32());
            }
        }
        assert_close(&read(got), &want, 0.0, &format!("split_qkv part {which}"));
    }
}

#[test]
fn kv_head_expand_repeats_each_kv_head() {
    let c = compute();
    let k = c
        .compile_pipeline("t_kv", sources::PHASE27_OPS, "kv_head_expand_f16")
        .expect("kv_head_expand_f16");

    let (seq, kv_heads, heads, hd) = (4usize, 2usize, 6usize, 3usize);
    let src = ramp(seq * kv_heads * hd, 41);
    let x = tensor(&c, &src, [seq, kv_heads, hd]);

    let cb = c.new_command_buffer();
    let out = gpu_ops::kv_head_expand_on(&c, &k, cb.as_ref(), &x, seq, kv_heads, heads, hd);
    cb.commit();
    cb.wait_until_completed();

    let group = heads / kv_heads;
    let mut want = Vec::with_capacity(seq * heads * hd);
    for s in 0..seq {
        for h in 0..heads {
            let kvh = h / group;
            for d in 0..hd {
                want.push(src[(s * kv_heads + kvh) * hd + d].to_f32());
            }
        }
    }
    assert_close(&read(&out), &want, 0.0, "kv_head_expand");
}

#[test]
fn max_pool2d_takes_the_window_maximum() {
    let c = compute();
    let k = c
        .compile_pipeline("t_mp", sources::PHASE27_OPS, "max_pool2d_f16")
        .expect("max_pool2d_f16");

    let (ch, h, w) = (2usize, 6usize, 8usize);
    let src = ramp(ch * h * w, 53);
    let x = tensor(&c, &src, [ch, h, w]);

    let cb = c.new_command_buffer();
    let out = gpu_ops::max_pool2d_on(&c, &k, cb.as_ref(), &x, ch, h, w);
    cb.commit();
    cb.wait_until_completed();

    let (oh, ow) = (h / 2, w / 2);
    let mut want = Vec::with_capacity(ch * oh * ow);
    for cc in 0..ch {
        for y in 0..oh {
            for x0 in 0..ow {
                let mut m = f32::NEG_INFINITY;
                for dy in 0..2 {
                    for dx in 0..2 {
                        m = m.max(src[cc * h * w + (y * 2 + dy) * w + (x0 * 2 + dx)].to_f32());
                    }
                }
                want.push(m);
            }
        }
    }
    assert_close(&read(&out), &want, 0.0, "max_pool2d");
}

#[test]
fn modulation_split_splits_six_ways() {
    let c = compute();
    let k = c
        .compile_pipeline("t_ms", sources::PHASE27_OPS, "modulation_split_f16")
        .expect("modulation_split_f16");

    let hidden = 9usize;
    let src = ramp(hidden * 6, 61);
    let params = tensor(&c, &src, [hidden * 6]);

    let cb = c.new_command_buffer();
    let parts = gpu_ops::modulation_split_on(&c, &k, cb.as_ref(), &params, hidden).expect("split");
    cb.commit();
    cb.wait_until_completed();

    for (i, part) in parts.iter().enumerate() {
        let want: Vec<f32> = (0..hidden).map(|j| src[i * hidden + j].to_f32()).collect();
        assert_close(&read(part), &want, 0.0, &format!("modulation part {i}"));
    }
}

// ─── arithmetic: checked against a CPU reference ────────────────────────────

#[test]
fn elu_matches_cpu() {
    let c = compute();
    let k = c
        .compile_pipeline("t_elu", sources::PHASE27_OPS, "elu_f16")
        .expect("elu_f16");

    let src = ramp(257, 71); // deliberately not a multiple of the threadgroup
    let x = tensor(&c, &src, [src.len()]);
    let alpha = 1.3f32;

    let cb = c.new_command_buffer();
    let out = gpu_ops::elu_on(&c, &k, cb.as_ref(), &x, alpha);
    cb.commit();
    cb.wait_until_completed();

    let want: Vec<f32> = src
        .iter()
        .map(|v| {
            let v = v.to_f32();
            if v >= 0.0 { v } else { alpha * (v.exp() - 1.0) }
        })
        .collect();
    assert_close(&read(&out), &want, 2e-3, "elu");
}

#[test]
fn euler_step_matches_cpu() {
    let c = compute();
    let k = c
        .compile_pipeline("t_eu", sources::PHASE27_OPS, "euler_step_f16")
        .expect("euler_step_f16");

    let xs = ramp(130, 83);
    let vs = ramp(130, 97);
    let dt = 0.25f32;
    let xt = tensor(&c, &xs, [xs.len()]);
    let vt = tensor(&c, &vs, [vs.len()]);

    let cb = c.new_command_buffer();
    let out = gpu_ops::euler_step_on(&c, &k, cb.as_ref(), &xt, &vt, dt);
    cb.commit();
    cb.wait_until_completed();

    let want: Vec<f32> = xs
        .iter()
        .zip(&vs)
        .map(|(a, b)| a.to_f32() + dt * b.to_f32())
        .collect();
    assert_close(&read(&out), &want, 2e-3, "euler_step");
}

#[test]
fn rope_batched_matches_cpu_rotation() {
    let c = compute();
    let k = c
        .compile_pipeline("t_rope", sources::PHASE27_OPS, "rope_batched_f16")
        .expect("rope_batched_f16");

    let (seq, heads, hd, start) = (5usize, 3usize, 8usize, 2usize);
    let theta = 10000.0f32;
    let src = ramp(seq * heads * hd, 101);
    let io = tensor(&c, &src, [seq, heads, hd]);

    let cb = c.new_command_buffer();
    gpu_ops::rope_batched_on(&c, &k, cb.as_ref(), &io, theta, seq, heads, hd, start);
    cb.commit();
    cb.wait_until_completed();

    let halfd = hd / 2;
    let mut want: Vec<f32> = src.iter().map(|v| v.to_f32()).collect();
    for s in 0..seq {
        for h in 0..heads {
            let base = (s * heads + h) * hd;
            for j in 0..halfd {
                let inv = theta.powf(-((2 * j) as f32) / hd as f32);
                let angle = (start + s) as f32 * inv;
                let (sn, cs) = angle.sin_cos();
                let x0 = src[base + j].to_f32();
                let x1 = src[base + halfd + j].to_f32();
                want[base + j] = x0 * cs - x1 * sn;
                want[base + halfd + j] = x0 * sn + x1 * cs;
            }
        }
    }
    assert_close(&read(&io), &want, 3e-3, "rope_batched");
}

#[test]
fn matvec_matches_cpu() {
    let c = compute();
    let k = c
        .compile_pipeline("t_mv", sources::MATVEC, "matvec_f16")
        .expect("matvec_f16");

    let (n, kdim) = (12usize, 64usize);
    let xs = ramp(kdim, 113);
    let ws = ramp(n * kdim, 127);
    let x = tensor(&c, &xs, [1, kdim]);
    let w = tensor(&c, &ws, [n, kdim]);

    let device = c.device().raw();
    let out_buf = device.new_buffer((n * 2) as u64, metal_shared());
    let cb = c.new_command_buffer();
    let (nn, kk) = (n as u32, kdim as u32);
    c.dispatch(cb.as_ref(), &k, (n, 1, 1), (32, 1, 1), |e| {
        gpu_ops::set_tensor_buffer(e, 0, &x);
        gpu_ops::set_tensor_buffer(e, 1, &w);
        e.set_buffer(2, Some(&out_buf), 0);
        e.set_bytes(3, 4, &nn as *const u32 as *const _);
        e.set_bytes(4, 4, &kk as *const u32 as *const _);
    });
    cb.commit();
    cb.wait_until_completed();

    let got = unsafe { std::slice::from_raw_parts(out_buf.contents() as *const f16, n) };
    let got: Vec<f32> = got.iter().map(|v| v.to_f32()).collect();

    let want: Vec<f32> = (0..n)
        .map(|r| {
            (0..kdim)
                .map(|i| ws[r * kdim + i].to_f32() * xs[i].to_f32())
                .sum()
        })
        .collect();
    assert_close(&got, &want, 5e-3, "matvec");
}

#[test]
fn matmul_q8_0_matches_cpu_dequantized_matmul() {
    let c = compute();
    let k = c
        .compile_pipeline("t_q8", sources::MATMUL_Q8_0, "matmul_q8_0_f16")
        .expect("matmul_q8_0_f16");

    // Q8_0: 32 values per block, 34 bytes -- f16 scale then 32 int8.
    const QK: usize = 32;
    const BLK: usize = 34;
    let (n, kdim) = (6usize, 64usize);
    let blocks_per_row = kdim / QK;

    let xs = ramp(kdim, 131);
    let mut raw = vec![0u8; n * blocks_per_row * BLK];
    let mut dequant = vec![0f32; n * kdim];
    for r in 0..n {
        for b in 0..blocks_per_row {
            let scale = 0.01 + 0.003 * ((r * blocks_per_row + b) as f32);
            let off = (r * blocks_per_row + b) * BLK;
            raw[off..off + 2].copy_from_slice(&f16::from_f32(scale).to_le_bytes());
            for i in 0..QK {
                let q = (((r * 7 + b * 13 + i * 3) % 255) as i32 - 127) as i8;
                raw[off + 2 + i] = q as u8;
                dequant[r * kdim + b * QK + i] = q as f32 * f16::from_f32(scale).to_f32();
            }
        }
    }

    let x = tensor(&c, &xs, [1, kdim]);
    let device = c.device().raw();
    let w_buf = device.new_buffer_with_data(
        raw.as_ptr() as *const _,
        raw.len() as u64,
        metal_shared(),
    );
    let out_buf = device.new_buffer((n * 2) as u64, metal_shared());

    let cb = c.new_command_buffer();
    let (nn, kk) = (n as u32, kdim as u32);
    c.dispatch(cb.as_ref(), &k, (n, 1, 1), (32, 1, 1), |e| {
        gpu_ops::set_tensor_buffer(e, 0, &x);
        e.set_buffer(1, Some(&w_buf), 0);
        e.set_buffer(2, Some(&out_buf), 0);
        e.set_bytes(3, 4, &nn as *const u32 as *const _);
        e.set_bytes(4, 4, &kk as *const u32 as *const _);
    });
    cb.commit();
    cb.wait_until_completed();

    let got = unsafe { std::slice::from_raw_parts(out_buf.contents() as *const f16, n) };
    let got: Vec<f32> = got.iter().map(|v| v.to_f32()).collect();

    let want: Vec<f32> = (0..n)
        .map(|r| {
            (0..kdim)
                .map(|i| dequant[r * kdim + i] * xs[i].to_f32())
                .sum()
        })
        .collect();
    assert_close(&got, &want, 1e-2, "matmul_q8_0");
}

#[test]
fn flash_attention_matches_naive_softmax_attention() {
    let c = compute();
    let k = c
        .compile_pipeline("t_fa", sources::PHASE27_OPS, "flash_attention_f16")
        .expect("flash_attention_f16");

    // kv_seq deliberately not a multiple of the 32-wide tile, and q_seq
    // spans more than one query block.
    let (heads, q_seq, kv_seq, hd) = (2usize, 40usize, 37usize, 16usize);
    let scale = 1.0f32 / (hd as f32).sqrt();

    let qs = ramp(heads * q_seq * hd, 211);
    let ks = ramp(heads * kv_seq * hd, 223);
    let vs = ramp(heads * kv_seq * hd, 227);
    let q = tensor(&c, &qs, [heads, q_seq, hd]);
    let kt = tensor(&c, &ks, [heads, kv_seq, hd]);
    let vt = tensor(&c, &vs, [heads, kv_seq, hd]);

    let device = c.device().raw();
    let out_buf = device.new_buffer((heads * q_seq * hd * 2) as u64, metal_shared());

    let block = 32usize;
    let q_blocks = q_seq.div_ceil(block);
    let tg_mem = (3 * block * hd * 2) as u64;
    let (qs_u, kv_u, hd_u) = (q_seq as u32, kv_seq as u32, hd as u32);

    let cb = c.new_command_buffer();
    c.dispatch(cb.as_ref(), &k, (heads, q_blocks, 1), (block, 1, 1), |e| {
        gpu_ops::set_tensor_buffer(e, 0, &q);
        gpu_ops::set_tensor_buffer(e, 1, &kt);
        gpu_ops::set_tensor_buffer(e, 2, &vt);
        e.set_buffer(3, Some(&out_buf), 0);
        e.set_bytes(4, 4, &qs_u as *const u32 as *const _);
        e.set_bytes(5, 4, &kv_u as *const u32 as *const _);
        e.set_bytes(6, 4, &hd_u as *const u32 as *const _);
        e.set_bytes(7, 4, &scale as *const f32 as *const _);
        e.set_threadgroup_memory_length(0, tg_mem);
    });
    cb.commit();
    cb.wait_until_completed();

    let got = unsafe {
        std::slice::from_raw_parts(out_buf.contents() as *const f16, heads * q_seq * hd)
    };
    let got: Vec<f32> = got.iter().map(|v| v.to_f32()).collect();

    // Naive reference: full softmax over all keys, non-causal.
    let mut want = vec![0f32; heads * q_seq * hd];
    for h in 0..heads {
        for i in 0..q_seq {
            let qb = (h * q_seq + i) * hd;
            let mut logits = Vec::with_capacity(kv_seq);
            for j in 0..kv_seq {
                let kb = (h * kv_seq + j) * hd;
                let dot: f32 = (0..hd)
                    .map(|d| qs[qb + d].to_f32() * ks[kb + d].to_f32())
                    .sum();
                logits.push(dot * scale);
            }
            let m = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
            let exps: Vec<f32> = logits.iter().map(|l| (l - m).exp()).collect();
            let denom: f32 = exps.iter().sum();
            for d in 0..hd {
                let acc: f32 = (0..kv_seq)
                    .map(|j| exps[j] * vs[(h * kv_seq + j) * hd + d].to_f32())
                    .sum();
                want[(h * q_seq + i) * hd + d] = acc / denom;
            }
        }
    }
    assert_close(&got, &want, 1e-2, "flash_attention");
}

/// `StorageModeShared`, without naming the `metal` crate as a dev-dependency.
fn metal_shared() -> metal::MTLResourceOptions {
    metal::MTLResourceOptions::StorageModeShared
}

// ─── the Mochi video-VAE decode stages ──────────────────────────────────────
//
// These six are dispatched directly from `architecture/mochi.rs` rather than
// through a `*_on` helper, so each test reproduces that call site's grid shape
// as well as its buffer layout. Note the grids are threadgroup counts, not
// thread counts, so they oversubscribe heavily -- that is the existing
// convention, and the bounds guards in the kernels are what make it safe.

#[test]
fn vae_conv1x1_matches_cpu() {
    let c = compute();
    let k = c
        .compile_pipeline("t_v1", sources::PHASE27_OPS, "vae_conv1x1_f16")
        .expect("vae_conv1x1_f16");

    let (in_ch, out_ch, spatial) = (5usize, 3usize, 7usize);
    let xs = ramp(in_ch * spatial, 301);
    let ws = ramp(out_ch * in_ch, 307);
    let bs = ramp(out_ch, 311);
    let x = tensor(&c, &xs, [in_ch, spatial]);
    let w = tensor(&c, &ws, [out_ch, in_ch]);
    let b = tensor(&c, &bs, [out_ch]);

    let out = c
        .device()
        .raw()
        .new_buffer((out_ch * spatial * 2) as u64, metal_shared());
    let (i_u, o_u, s_u) = (in_ch as u32, out_ch as u32, spatial as u32);
    let cb = c.new_command_buffer();
    c.dispatch(cb.as_ref(), &k, (spatial, out_ch, 1), (16, 16, 1), |e| {
        gpu_ops::set_tensor_buffer(e, 0, &x);
        gpu_ops::set_tensor_buffer(e, 1, &w);
        gpu_ops::set_tensor_buffer(e, 2, &b);
        e.set_buffer(3, Some(&out), 0);
        e.set_bytes(4, 4, &i_u as *const u32 as *const _);
        e.set_bytes(5, 4, &o_u as *const u32 as *const _);
        e.set_bytes(6, 4, &s_u as *const u32 as *const _);
    });
    cb.commit();
    cb.wait_until_completed();

    let mut want = vec![0f32; out_ch * spatial];
    for oc in 0..out_ch {
        for s in 0..spatial {
            let mut acc = bs[oc].to_f32();
            for ic in 0..in_ch {
                acc += ws[oc * in_ch + ic].to_f32() * xs[ic * spatial + s].to_f32();
            }
            want[oc * spatial + s] = acc;
        }
    }
    assert_close(&read_buf_f16(&out, out_ch * spatial), &want, 5e-3, "vae_conv1x1");
}

#[test]
fn vae_spatial_upsample_2x_repeats_each_pixel() {
    let c = compute();
    let k = c
        .compile_pipeline("t_vsu", sources::PHASE27_OPS, "vae_spatial_upsample_2x_f16")
        .expect("vae_spatial_upsample_2x_f16");

    let (ch, t, ih, iw) = (2usize, 3usize, 4usize, 5usize);
    let (oh, ow) = (ih * 2, iw * 2);
    let xs = ramp(ch * t * ih * iw, 317);
    let x = tensor(&c, &xs, [ch, t * ih * iw]);

    let out = c
        .device()
        .raw()
        .new_buffer((ch * t * oh * ow * 2) as u64, metal_shared());
    let (c_u, t_u, h_u, w_u) = (ch as u32, t as u32, ih as u32, iw as u32);
    let cb = c.new_command_buffer();
    c.dispatch(cb.as_ref(), &k, (ow, oh, ch * t), (16, 16, 1), |e| {
        gpu_ops::set_tensor_buffer(e, 0, &x);
        e.set_buffer(1, Some(&out), 0);
        e.set_bytes(2, 4, &c_u as *const u32 as *const _);
        e.set_bytes(3, 4, &t_u as *const u32 as *const _);
        e.set_bytes(4, 4, &h_u as *const u32 as *const _);
        e.set_bytes(5, 4, &w_u as *const u32 as *const _);
    });
    cb.commit();
    cb.wait_until_completed();

    let mut want = vec![0f32; ch * t * oh * ow];
    for cc in 0..ch {
        for tt in 0..t {
            for y in 0..oh {
                for x0 in 0..ow {
                    let src = xs[cc * t * ih * iw + tt * ih * iw + (y / 2) * iw + (x0 / 2)];
                    want[cc * t * oh * ow + tt * oh * ow + y * ow + x0] = src.to_f32();
                }
            }
        }
    }
    assert_close(
        &read_buf_f16(&out, ch * t * oh * ow),
        &want,
        0.0,
        "vae_spatial_upsample_2x",
    );
}

#[test]
fn vae_channel_reduce_averages_each_channel_group() {
    let c = compute();
    let k = c
        .compile_pipeline("t_vcr", sources::PHASE27_OPS, "vae_channel_reduce_f16")
        .expect("vae_channel_reduce_f16");

    let (in_ch, out_ch, spatial) = (8usize, 3usize, 6usize); // 8/3 = group of 2, remainder
    let xs = ramp(in_ch * spatial, 331);
    let x = tensor(&c, &xs, [in_ch, spatial]);

    let out = c
        .device()
        .raw()
        .new_buffer((out_ch * spatial * 2) as u64, metal_shared());
    let (i_u, o_u, s_u) = (in_ch as u32, out_ch as u32, spatial as u32);
    let cb = c.new_command_buffer();
    c.dispatch(cb.as_ref(), &k, (spatial, out_ch, 1), (16, 16, 1), |e| {
        gpu_ops::set_tensor_buffer(e, 0, &x);
        e.set_buffer(1, Some(&out), 0);
        e.set_bytes(2, 4, &i_u as *const u32 as *const _);
        e.set_bytes(3, 4, &o_u as *const u32 as *const _);
        e.set_bytes(4, 4, &s_u as *const u32 as *const _);
    });
    cb.commit();
    cb.wait_until_completed();

    let group = (in_ch / out_ch).max(1);
    let mut want = vec![0f32; out_ch * spatial];
    for oc in 0..out_ch {
        for s in 0..spatial {
            let mut acc = 0f32;
            let mut n = 0usize;
            for g in 0..group {
                let ic = oc * group + g;
                if ic >= in_ch {
                    break;
                }
                acc += xs[ic * spatial + s].to_f32();
                n += 1;
            }
            want[oc * spatial + s] = if n > 0 { acc / n as f32 } else { 0.0 };
        }
    }
    assert_close(&read_buf_f16(&out, out_ch * spatial), &want, 3e-3, "vae_channel_reduce");
}

#[test]
fn vae_temporal_upsample_holds_each_source_frame() {
    let c = compute();
    let k = c
        .compile_pipeline("t_vtu", sources::PHASE27_OPS, "vae_temporal_upsample_f16")
        .expect("vae_temporal_upsample_f16");

    let (ch, in_t, out_t, hw) = (2usize, 3usize, 7usize, 5usize); // 7/3 is not integral
    let xs = ramp(ch * in_t * hw, 337);
    let x = tensor(&c, &xs, [ch, in_t * hw]);

    let out = c
        .device()
        .raw()
        .new_buffer((ch * out_t * hw * 2) as u64, metal_shared());
    let (c_u, i_u, o_u, hw_u) = (ch as u32, in_t as u32, out_t as u32, hw as u32);
    let cb = c.new_command_buffer();
    c.dispatch(cb.as_ref(), &k, (hw, out_t, ch), (16, 16, 1), |e| {
        gpu_ops::set_tensor_buffer(e, 0, &x);
        e.set_buffer(1, Some(&out), 0);
        e.set_bytes(2, 4, &c_u as *const u32 as *const _);
        e.set_bytes(3, 4, &i_u as *const u32 as *const _);
        e.set_bytes(4, 4, &o_u as *const u32 as *const _);
        e.set_bytes(5, 4, &hw_u as *const u32 as *const _);
    });
    cb.commit();
    cb.wait_until_completed();

    let mut want = vec![0f32; ch * out_t * hw];
    for cc in 0..ch {
        for ot in 0..out_t {
            let it = ((ot * in_t) / out_t).min(in_t - 1);
            for p in 0..hw {
                want[cc * out_t * hw + ot * hw + p] = xs[cc * in_t * hw + it * hw + p].to_f32();
            }
        }
    }
    assert_close(
        &read_buf_f16(&out, ch * out_t * hw),
        &want,
        0.0,
        "vae_temporal_upsample",
    );
}

#[test]
fn vae_final_conv_sigmoid_maps_to_three_channels_in_zero_one() {
    let c = compute();
    let k = c
        .compile_pipeline("t_vfc", sources::PHASE27_OPS, "vae_final_conv_sigmoid_f16")
        .expect("vae_final_conv_sigmoid_f16");

    let (in_ch, spatial) = (9usize, 11usize);
    let xs = ramp(in_ch * spatial, 347);
    let x = tensor(&c, &xs, [in_ch, spatial]);

    let out = c
        .device()
        .raw()
        .new_buffer((3 * spatial * 2) as u64, metal_shared());
    let (i_u, s_u) = (in_ch as u32, spatial as u32);
    let cb = c.new_command_buffer();
    c.dispatch(cb.as_ref(), &k, (spatial, 3, 1), (256, 1, 1), |e| {
        gpu_ops::set_tensor_buffer(e, 0, &x);
        e.set_buffer(1, Some(&out), 0);
        e.set_bytes(2, 4, &i_u as *const u32 as *const _);
        e.set_bytes(3, 4, &s_u as *const u32 as *const _);
    });
    cb.commit();
    cb.wait_until_completed();

    let group = (in_ch / 3).max(1);
    let mut want = vec![0f32; 3 * spatial];
    for oc in 0..3usize {
        for s in 0..spatial {
            let mut acc = 0f32;
            let mut n = 0usize;
            for g in 0..group {
                let ic = oc * group + g;
                if ic >= in_ch {
                    break;
                }
                acc += xs[ic * spatial + s].to_f32();
                n += 1;
            }
            let mean = if n > 0 { acc / n as f32 } else { 0.0 };
            want[oc * spatial + s] = 1.0 / (1.0 + (-mean).exp());
        }
    }
    let got = read_buf_f16(&out, 3 * spatial);
    assert!(
        got.iter().all(|v| (0.0..=1.0).contains(v)),
        "a sigmoid must land in [0, 1]"
    );
    assert_close(&got, &want, 3e-3, "vae_final_conv_sigmoid");
}

#[test]
fn vae_extract_frame_interleaves_one_frame_as_f32() {
    let c = compute();
    let k = c
        .compile_pipeline("t_vef", sources::PHASE27_OPS, "vae_extract_frame_f16")
        .expect("vae_extract_frame_f16");

    let (t, h, w) = (4usize, 3usize, 5usize);
    let xs = ramp(3 * t * h * w, 353);
    let rgb = tensor(&c, &xs, [3, t * h * w]);

    let frame_idx = 2usize;
    // f32 destination: four bytes per element, as the Mochi call site expects.
    let out = c
        .device()
        .raw()
        .new_buffer((h * w * 3 * 4) as u64, metal_shared());
    let (t_u, h_u, w_u, f_u) = (t as u32, h as u32, w as u32, frame_idx as u32);
    let cb = c.new_command_buffer();
    c.dispatch(cb.as_ref(), &k, (w, h, 1), (16, 16, 1), |e| {
        gpu_ops::set_tensor_buffer(e, 0, &rgb);
        e.set_buffer(1, Some(&out), 0);
        e.set_bytes(2, 4, &t_u as *const u32 as *const _);
        e.set_bytes(3, 4, &h_u as *const u32 as *const _);
        e.set_bytes(4, 4, &w_u as *const u32 as *const _);
        e.set_bytes(5, 4, &f_u as *const u32 as *const _);
    });
    cb.commit();
    cb.wait_until_completed();

    let got = unsafe { std::slice::from_raw_parts(out.contents() as *const f32, h * w * 3) };
    let plane = t * h * w;
    let mut want = vec![0f32; h * w * 3];
    for y in 0..h {
        for x in 0..w {
            let off = frame_idx * h * w + y * w + x;
            for ch in 0..3usize {
                want[(y * w + x) * 3 + ch] = xs[ch * plane + off].to_f32();
            }
        }
    }
    assert_close(got, &want, 0.0, "vae_extract_frame");
}

// ─── the remaining three ────────────────────────────────────────────────────

#[test]
fn im2col_patch_embed_matches_cpu_with_zero_padding() {
    let c = compute();
    let k = c
        .compile_pipeline("t_i2c", sources::PHASE27_OPS, "im2col_patch_embed_f16")
        .expect("im2col_patch_embed_f16");

    // Padding and a stride that does not divide the input: the taps that fall
    // outside must read as zero.
    let (ch, h, w, kh, kw, sh, sw, ph, pw) = (2usize, 7usize, 6usize, 3, 3, 2, 2, 1, 1);
    let xs = ramp(ch * h * w, 401);
    let x = tensor(&c, &xs, [ch, h, w]);

    let cb = c.new_command_buffer();
    let out = gpu_ops::im2col_patch_embed_on(
        &c, &k, cb.as_ref(), &x, ch, h, w, kh, kw, sh, sw, ph, pw,
    );
    cb.commit();
    cb.wait_until_completed();

    let out_h = (h + 2 * ph - kh) / sh + 1;
    let out_w = (w + 2 * pw - kw) / sw + 1;
    let patch_dim = ch * kh * kw;
    assert_eq!(out.shape().dims(), &[out_h * out_w, patch_dim]);

    let mut want = vec![0f32; out_h * out_w * patch_dim];
    let mut zeros = 0usize;
    for oy in 0..out_h {
        for ox in 0..out_w {
            for cc in 0..ch {
                for ky in 0..kh {
                    for kx in 0..kw {
                        let iy = (oy * sh + ky) as isize - ph as isize;
                        let ix = (ox * sw + kx) as isize - pw as isize;
                        let e = cc * kh * kw + ky * kw + kx;
                        let v = if iy >= 0 && iy < h as isize && ix >= 0 && ix < w as isize {
                            xs[(cc * h + iy as usize) * w + ix as usize].to_f32()
                        } else {
                            zeros += 1;
                            0.0
                        };
                        want[(oy * out_w + ox) * patch_dim + e] = v;
                    }
                }
            }
        }
    }
    assert!(zeros > 0, "this geometry must exercise the padding path");
    assert_close(&read(&out), &want, 0.0, "im2col_patch_embed");
}

#[test]
fn dequantize_rvq_sums_the_per_codebook_embeddings() {
    let c = compute();
    let k = c
        .compile_pipeline("t_rvq", sources::PHASE27_OPS, "dequantize_rvq_f16")
        .expect("dequantize_rvq_f16");

    let (n_cb, n_frames, cb_size, embed) = (3usize, 5usize, 7usize, 4usize);
    let embeds = ramp(n_cb * cb_size * embed, 409);
    let et = tensor(&c, &embeds, [n_cb, cb_size, embed]);

    // One code per (codebook, frame); the last is out of range on purpose and
    // must be skipped rather than read past the table.
    let mut codes: Vec<u32> = (0..n_cb * n_frames)
        .map(|i| ((i * 3 + 1) % cb_size) as u32)
        .collect();
    *codes.last_mut().unwrap() = cb_size as u32 + 2;

    let codes_bytes: &[u8] = bytemuck_cast(&codes);
    let codes_buf = c.device().raw().new_buffer_with_data(
        codes_bytes.as_ptr() as *const _,
        codes_bytes.len() as u64,
        metal_shared(),
    );

    let cb = c.new_command_buffer();
    let out = gpu_ops::dequantize_rvq_on(
        &c, &k, cb.as_ref(), &codes_buf, &et, n_cb, n_frames, cb_size, embed,
    );
    cb.commit();
    cb.wait_until_completed();

    let mut want = vec![0f32; n_frames * embed];
    for f in 0..n_frames {
        for d in 0..embed {
            let mut acc = 0f32;
            for q in 0..n_cb {
                let code = codes[q * n_frames + f] as usize;
                if code >= cb_size {
                    continue;
                }
                acc += embeds[(q * cb_size + code) * embed + d].to_f32();
            }
            want[f * embed + d] = acc;
        }
    }
    assert_close(&read(&out), &want, 3e-3, "dequantize_rvq");
}

#[test]
fn dilated_conv1d_matches_cpu_with_same_padding() {
    let c = compute();
    let k = c
        .compile_pipeline("t_dc1", sources::PHASE27_OPS, "dilated_conv1d_f16")
        .expect("dilated_conv1d_f16");

    let (in_ch, out_ch, len, ksize, dilation) = (3usize, 2usize, 9usize, 3usize, 2usize);
    let xs = ramp(in_ch * len, 419);
    let ws = ramp(out_ch * in_ch * ksize, 421);
    let bs = ramp(out_ch, 431);
    let x = tensor(&c, &xs, [in_ch, len]);
    let w = tensor(&c, &ws, [out_ch, in_ch, ksize]);
    let b = tensor(&c, &bs, [out_ch]);

    let out = c
        .device()
        .raw()
        .new_buffer((out_ch * len * 2) as u64, metal_shared());
    let dims: [u32; 5] = [
        in_ch as u32, out_ch as u32, len as u32, ksize as u32, dilation as u32,
    ];
    let cb = c.new_command_buffer();
    c.dispatch(cb.as_ref(), &k, (len, out_ch, 1), (16, 16, 1), |e| {
        gpu_ops::set_tensor_buffer(e, 0, &x);
        gpu_ops::set_tensor_buffer(e, 1, &w);
        gpu_ops::set_tensor_buffer(e, 2, &b);
        e.set_buffer(3, Some(&out), 0);
        for (i, d) in dims.iter().enumerate() {
            e.set_bytes(4 + i as u64, 4, d as *const u32 as *const _);
        }
    });
    cb.commit();
    cb.wait_until_completed();

    let pad = ((ksize - 1) * dilation) / 2;
    let mut want = vec![0f32; out_ch * len];
    for oc in 0..out_ch {
        for t in 0..len {
            let mut acc = bs[oc].to_f32();
            for ic in 0..in_ch {
                for kk in 0..ksize {
                    let s = t as isize + (kk * dilation) as isize - pad as isize;
                    if s < 0 || s >= len as isize {
                        continue;
                    }
                    acc += ws[(oc * in_ch + ic) * ksize + kk].to_f32()
                        * xs[ic * len + s as usize].to_f32();
                }
            }
            want[oc * len + t] = acc;
        }
    }
    assert_close(&read_buf_f16(&out, out_ch * len), &want, 5e-3, "dilated_conv1d");
}

/// Read `count` f16 values out of a shared Metal buffer.
fn read_buf_f16(buf: &metal::Buffer, count: usize) -> Vec<f32> {
    let s = unsafe { std::slice::from_raw_parts(buf.contents() as *const f16, count) };
    s.iter().map(|v| v.to_f32()).collect()
}

/// `u32` slice as bytes, without pulling in bytemuck as a dev-dependency.
fn bytemuck_cast(v: &[u32]) -> &[u8] {
    unsafe { std::slice::from_raw_parts(v.as_ptr() as *const u8, std::mem::size_of_val(v)) }
}

// ─── sources::PROCEDURAL ────────────────────────────────────────────────────
//
// The fusion pipeline's kernels. `FusionPipeline` is not constructed anywhere
// yet, so these dispatch the shaders the way `inference/fusion.rs` does and
// check them directly. All five work in f32 buffers, not tensors.

fn f32_buffer(c: &MetalCompute, v: &[f32]) -> metal::Buffer {
    let bytes =
        unsafe { std::slice::from_raw_parts(v.as_ptr() as *const u8, std::mem::size_of_val(v)) };
    c.device()
        .raw()
        .new_buffer_with_data(bytes.as_ptr() as *const _, bytes.len() as u64, metal_shared())
}

fn read_buf_f32(buf: &metal::Buffer, count: usize) -> Vec<f32> {
    unsafe { std::slice::from_raw_parts(buf.contents() as *const f32, count) }.to_vec()
}

#[test]
fn perlin_noise_is_deterministic_bounded_and_blends() {
    let c = compute();
    let k = c
        .compile_pipeline("t_pn", sources::PROCEDURAL, "perlin_noise_f32")
        .expect("perlin_noise_f32");

    let (w, h) = (16usize, 12usize);
    let amplitude = 2.0f32;

    // blend: 0 replace, 1 add. noise_type: 0 perlin, 1 simplex, 2 worley, 3 value.
    let run = |noise_type: u32, blend: u32, seed: u32, init: &[f32]| -> Vec<f32> {
        let buf = f32_buffer(&c, init);
        let (w_u, h_u) = (w as u32, h as u32);
        let (oct, lac, per, scale) = (3u32, 2.0f32, 0.5f32, 0.07f32);
        let cb = c.new_command_buffer();
        c.dispatch_2d(cb.as_ref(), &k, w, h, |e| {
            e.set_buffer(0, Some(&buf), 0);
            e.set_bytes(1, 4, &w_u as *const u32 as *const _);
            e.set_bytes(2, 4, &h_u as *const u32 as *const _);
            e.set_bytes(3, 4, &noise_type as *const u32 as *const _);
            e.set_bytes(4, 4, &oct as *const u32 as *const _);
            e.set_bytes(5, 4, &lac as *const f32 as *const _);
            e.set_bytes(6, 4, &per as *const f32 as *const _);
            e.set_bytes(7, 4, &scale as *const f32 as *const _);
            e.set_bytes(8, 4, &amplitude as *const f32 as *const _);
            e.set_bytes(9, 4, &seed as *const u32 as *const _);
            e.set_bytes(10, 4, &blend as *const u32 as *const _);
        });
        cb.commit();
        cb.wait_until_completed();
        read_buf_f32(&buf, w * h)
    };

    let zero = vec![0f32; w * h];
    for nt in 0..4u32 {
        let a = run(nt, 0, 7, &zero);
        let b = run(nt, 0, 7, &zero);
        assert_eq!(a, b, "noise type {nt} must be deterministic for a fixed seed");

        // fbm is a weighted mean of bases in [-1, 1], so |v| <= amplitude.
        assert!(
            a.iter().all(|v| v.abs() <= amplitude + 1e-4),
            "noise type {nt} exceeded its amplitude"
        );
        assert!(
            a.iter().any(|v| v.abs() > 1e-6),
            "noise type {nt} produced an all-zero field"
        );

        // A different seed must give a different field.
        let other = run(nt, 0, 8, &zero);
        assert_ne!(a, other, "noise type {nt} ignored its seed");

        // Additive blend onto an existing field adds to it.
        let added = run(nt, 1, 7, &a);
        let want: Vec<f32> = a.iter().map(|v| v + v).collect();
        assert_close(&added, &want, 1e-5, &format!("noise {nt} add blend"));
    }
}

#[test]
fn apply_palette_interpolates_between_stops() {
    let c = compute();
    let k = c
        .compile_pipeline("t_ap", sources::PROCEDURAL, "apply_palette_f32")
        .expect("apply_palette_f32");

    let (w, h) = (8usize, 4usize);
    // A field spanning [-1, 1], which the kernel remaps to [0, 1].
    let field: Vec<f32> = (0..w * h)
        .map(|i| -1.0 + 2.0 * (i as f32) / ((w * h - 1) as f32))
        .collect();
    // Three stops: black at 0, red at 0.5, white at 1.
    let palette: Vec<f32> = vec![
        0.0, 0.0, 0.0, 0.0, //
        0.5, 1.0, 0.0, 0.0, //
        1.0, 1.0, 1.0, 1.0,
    ];
    let stops = 3u32;

    let fb = f32_buffer(&c, &field);
    let pb = f32_buffer(&c, &palette);
    let out = c
        .device()
        .raw()
        .new_buffer((w * h * 3 * 4) as u64, metal_shared());
    let (w_u, h_u) = (w as u32, h as u32);
    let cb = c.new_command_buffer();
    c.dispatch_2d(cb.as_ref(), &k, w, h, |e| {
        e.set_buffer(0, Some(&fb), 0);
        e.set_buffer(1, Some(&pb), 0);
        e.set_buffer(2, Some(&out), 0);
        e.set_bytes(3, 4, &w_u as *const u32 as *const _);
        e.set_bytes(4, 4, &h_u as *const u32 as *const _);
        e.set_bytes(5, 4, &stops as *const u32 as *const _);
    });
    cb.commit();
    cb.wait_until_completed();

    let mut want = vec![0f32; w * h * 3];
    for (i, f) in field.iter().enumerate() {
        let t = (f * 0.5 + 0.5).clamp(0.0, 1.0);
        let (c0, c1, lo, hi) = if t <= 0.5 {
            ([0.0, 0.0, 0.0], [1.0, 0.0, 0.0], 0.0, 0.5)
        } else {
            ([1.0, 0.0, 0.0], [1.0, 1.0, 1.0], 0.5, 1.0)
        };
        let u = (t - lo) / (hi - lo);
        for ch in 0..3 {
            want[i * 3 + ch] = c0[ch] + u * (c1[ch] - c0[ch]);
        }
    }
    assert_close(&read_buf_f32(&out, w * h * 3), &want, 1e-5, "apply_palette");
}

#[test]
fn sdf_raymarch_hits_a_sphere_and_misses_past_its_edge() {
    let c = compute();
    let k = c
        .compile_pipeline("t_sdf", sources::PROCEDURAL, "sdf_raymarch_f32")
        .expect("sdf_raymarch_f32");

    let (w, h) = (32usize, 32usize);
    // One node: a sphere of radius 1 at the origin (op code 0).
    let mut scene = vec![0f32; 16];
    scene[0] = 0.0;
    scene[1] = 1.0;
    // eye(3), target(3), fov, max_dist, max_steps, background rgb(3), light(3)
    let camera: Vec<f32> = vec![
        0.0, 0.0, 4.0, //
        0.0, 0.0, 0.0, //
        0.6, 32.0, 128.0, //
        0.0, 0.25, 0.75, // a background nothing shaded can equal
        0.0, 0.0, 1.0,
    ];

    let sb = f32_buffer(&c, &scene);
    let camb = f32_buffer(&c, &camera);
    let out = c
        .device()
        .raw()
        .new_buffer((w * h * 3 * 4) as u64, metal_shared());
    let (w_u, h_u, n_u) = (w as u32, h as u32, 1u32);
    let cb = c.new_command_buffer();
    c.dispatch_2d(cb.as_ref(), &k, w, h, |e| {
        e.set_buffer(0, Some(&sb), 0);
        e.set_buffer(1, Some(&camb), 0);
        e.set_buffer(2, Some(&out), 0);
        e.set_bytes(3, 4, &w_u as *const u32 as *const _);
        e.set_bytes(4, 4, &h_u as *const u32 as *const _);
        e.set_bytes(5, 4, &n_u as *const u32 as *const _);
    });
    cb.commit();
    cb.wait_until_completed();

    let px = read_buf_f32(&out, w * h * 3);
    let bg = [0.0f32, 0.25, 0.75];
    let at = |y: usize, x: usize| {
        let i = (y * w + x) * 3;
        [px[i], px[i + 1], px[i + 2]]
    };
    let is_bg = |c: [f32; 3]| (0..3).all(|i| (c[i] - bg[i]).abs() < 1e-6);

    // The sphere subtends the middle of the frame; the corners cannot reach it.
    let centre = at(h / 2, w / 2);
    assert!(!is_bg(centre), "centre pixel should have hit the sphere, got {centre:?}");
    // Facing the light head-on, the centre is the brightest point: 0.1 + 0.9.
    assert!(
        (centre[0] - 1.0).abs() < 2e-2 && (centre[1] - 1.0).abs() < 2e-2,
        "centre should be fully lit, got {centre:?}"
    );
    for (y, x) in [(0, 0), (0, w - 1), (h - 1, 0), (h - 1, w - 1)] {
        let c = at(y, x);
        assert!(is_bg(c), "corner ({y},{x}) should be background, got {c:?}");
    }

    // Shading must vary across the sphere: a lit sphere is not a flat disc.
    let edge = at(h / 2, w / 2 + 9);
    assert!(
        !is_bg(edge) && edge[0] < centre[0] - 1e-2,
        "the sphere should fall off towards its limb: centre {centre:?}, edge {edge:?}"
    );
}

#[test]
fn ks_excite_fills_the_delay_line_within_the_velocity() {
    let c = compute();
    let k = c
        .compile_pipeline("t_kse", sources::PROCEDURAL, "ks_excite_f32")
        .expect("ks_excite_f32");

    let delay_len = 97usize; // not a multiple of any threadgroup width
    let velocity = 0.8f32;

    let excite = |etype: u32, position: f32, seed: u32| -> Vec<f32> {
        let buf = f32_buffer(&c, &vec![0f32; delay_len]);
        let dl = delay_len as u32;
        let cb = c.new_command_buffer();
        c.dispatch_1d(cb.as_ref(), &k, delay_len, |e| {
            e.set_buffer(0, Some(&buf), 0);
            e.set_bytes(1, 4, &dl as *const u32 as *const _);
            e.set_bytes(2, 4, &etype as *const u32 as *const _);
            e.set_bytes(3, 4, &position as *const f32 as *const _);
            e.set_bytes(4, 4, &velocity as *const f32 as *const _);
            e.set_bytes(5, 4, &seed as *const u32 as *const _);
        });
        cb.commit();
        cb.wait_until_completed();
        read_buf_f32(&buf, delay_len)
    };

    for etype in 0..3u32 {
        let v = excite(etype, 0.5, 11);
        assert_eq!(v.len(), delay_len);
        assert!(
            v.iter().all(|x| x.abs() <= velocity + 1e-5),
            "excitation {etype} exceeded the velocity envelope"
        );
        assert!(
            v.iter().any(|x| x.abs() > 1e-4),
            "excitation {etype} left the string silent"
        );
        assert_eq!(v, excite(etype, 0.5, 11), "excitation {etype} is not deterministic");
    }

    // A pluck is a triangle with its corner at the pluck point, so moving the
    // point moves where the maximum sits.
    let argmax = |v: &[f32]| {
        v.iter()
            .enumerate()
            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
            .unwrap()
            .0
    };
    let early = argmax(&excite(0, 0.2, 5));
    let late = argmax(&excite(0, 0.8, 5));
    assert!(
        early < late,
        "the pluck peak should follow the pluck position: {early} vs {late}"
    );
}

#[test]
fn ks_synthesize_matches_the_cpu_waveguide_loop() {
    let c = compute();
    let k = c
        .compile_pipeline("t_kss", sources::PROCEDURAL, "ks_synthesize_f32")
        .expect("ks_synthesize_f32");

    let (delay_len, num_samples) = (23usize, 200usize);
    let (damping, brightness) = (0.98f32, 0.35f32);
    // A deterministic starting delay line.
    let init: Vec<f32> = (0..delay_len)
        .map(|i| ((i as f32) / delay_len as f32) * 2.0 - 1.0)
        .collect();

    let delay = f32_buffer(&c, &init);
    let out = c
        .device()
        .raw()
        .new_buffer((num_samples * 4) as u64, metal_shared());
    let (dl, ns) = (delay_len as u32, num_samples as u32);
    let cb = c.new_command_buffer();
    c.dispatch_1d(cb.as_ref(), &k, 1, |e| {
        e.set_buffer(0, Some(&delay), 0);
        e.set_buffer(1, Some(&out), 0);
        e.set_bytes(2, 4, &dl as *const u32 as *const _);
        e.set_bytes(3, 4, &ns as *const u32 as *const _);
        e.set_bytes(4, 4, &damping as *const f32 as *const _);
        e.set_bytes(5, 4, &brightness as *const f32 as *const _);
    });
    cb.commit();
    cb.wait_until_completed();

    // The same serial recurrence on the CPU.
    let mut d = init.clone();
    let mut want = vec![0f32; num_samples];
    for n in 0..num_samples {
        let i = n % delay_len;
        let j = (i + 1) % delay_len;
        let cur = d[i];
        let nxt = d[j];
        want[n] = cur;
        let lowpass = 0.5 * (cur + nxt);
        d[i] = (lowpass + brightness * (cur - lowpass)) * damping;
    }

    let got = read_buf_f32(&out, num_samples);
    assert_close(&got, &want, 1e-4, "ks_synthesize");
    // The loop gain is below one, so the string must decay rather than ring on.
    let head: f32 = got[..delay_len].iter().map(|v| v.abs()).sum();
    let tail: f32 = got[num_samples - delay_len..].iter().map(|v| v.abs()).sum();
    assert!(tail < head, "a damped waveguide must decay: {head} -> {tail}");
}