fuel-dispatch 0.20.0

// SPDX-License-Identifier: MIT OR Apache-2.0
//! Live-Vulkan tests for the pipelined-executor binding-table
//! dispatch on Vulkan. Phase 7.6 step 9c Vulkan catch-up V.1.C —
//! the proof-of-life that one op (Add f32) flows end-to-end through
//! the new path on the Vulkan device.
//!
//! Requires a working Vulkan device (RTX 4070 on the dev machine
//! per the dev-environment memory). Tests are `#[ignore]` so the
//! default `cargo test` doesn't require a GPU; run with
//! `--include-ignored` for the GPU sweep.

#![cfg(feature = "vulkan")]

use std::sync::{Arc, Mutex, MutexGuard, OnceLock, RwLock};

use fuel_dispatch::{
    kernel::{KernelBindingTable, MatmulM, OpParams},
    vulkan_dispatch::register_vulkan_kernels,
};
use fuel_ir::{DType, Layout, Shape, dispatch::OpKind, probe::BackendId};
use fuel_memory::{BackendStorage, Storage, alloc_cpu_zeroed};
use fuel_vulkan_backend::VulkanBackend;

/// A serialized handle to the shared Vulkan backend for one live test.
///
/// Derefs to `Arc<VulkanBackend>`, so every existing call site — `&backend`,
/// `backend.method()`, `backend.clone()` — compiles unchanged. Holding it keeps
/// the process-wide `GPU_SERIAL` lock for the test's duration.
struct VkTest {
    _serial: MutexGuard<'static, ()>,
    backend: Arc<VulkanBackend>,
}

impl std::ops::Deref for VkTest {
    type Target = Arc<VulkanBackend>;
    fn deref(&self) -> &Self::Target {
        &self.backend
    }
}

/// Acquire the shared Vulkan backend, serialized against every other live test
/// in this binary. Returns `None` (skip) when no Vulkan device is present.
///
/// **Why shared + serialized.** The old body was `VulkanBackend::new()` — a fresh
/// VkInstance + VkDevice + VMA allocator + 512 MB recycler pool **per test**. Run
/// under libtest's default thread pool (~14 on this box), ~14 backends are alive at
/// once and their combined device memory OOMs the 8 GB RTX 4070 (observed:
/// `ERROR_OUT_OF_DEVICE_MEMORY` on `upload_bytes_handle`, scattered late in the
/// 204-test run; every test passed in isolation). Sharing ONE backend bounds live
/// device memory to a single working set regardless of parallelism.
///
/// Sharing then *requires* serialization: the backend's `one_shot` upload/download
/// path submits to its `VkQueue` without an internal lock, so concurrent use across
/// libtest threads would be a `vkQueueSubmit` data race (the queue must be
/// externally synchronized). `GPU_SERIAL` gives that external synchronization — one
/// test touches the GPU at a time — which is also the correct posture for live-GPU
/// tests on a shared card. Each test calls this exactly once, so the non-reentrant
/// lock never self-deadlocks.
fn backend_or_skip() -> Option<VkTest> {
    static GPU_SERIAL: OnceLock<Mutex<()>> = OnceLock::new();
    static SHARED: OnceLock<Option<Arc<VulkanBackend>>> = OnceLock::new();

    let serial = GPU_SERIAL
        .get_or_init(|| Mutex::new(()))
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner());
    let backend = SHARED
        .get_or_init(|| VulkanBackend::new().ok().map(Arc::new))
        .clone()?;
    Some(VkTest {
        _serial: serial,
        backend,
    })
}

fn upload_f32(backend: &Arc<VulkanBackend>, host: &[f32]) -> Storage {
    let bytes: &[u8] = bytemuck::cast_slice(host);
    let vk_bytes = backend
        .upload_bytes_handle(bytes)
        .expect("vulkan upload_bytes_handle");
    Storage::new(BackendStorage::Vulkan(vk_bytes), DType::F32)
}

fn download_f32(backend: &Arc<VulkanBackend>, s: &Storage) -> Vec<f32> {
    let bytes = match &s.inner {
        BackendStorage::Vulkan(v) => backend.download_bytes(v).expect("d2h"),
        _ => panic!("not on Vulkan"),
    };
    bytemuck::cast_slice::<u8, f32>(&bytes).to_vec()
}

/// Direct kernel-wrapper invocation. Skips the pipelined executor's
/// output-allocation arm and proves the dispatch wrapper + Slang
/// kernel produce correct bytes in isolation.
#[test]
#[ignore]
fn vulkan_dispatch_binary_add_f32_direct_wrapper() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let a_data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
    let b_data: Vec<f32> = vec![10.0, 20.0, 30.0, 40.0, 50.0, 60.0];
    let n = a_data.len();

    let a_storage = upload_f32(&backend, &a_data);
    let b_storage = upload_f32(&backend, &b_data);
    let out_bytes = backend
        .alloc_bytes_handle(n * 4)
        .expect("alloc_bytes_handle");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let alts = table.lookup_alternatives(
        OpKind::AddElementwise,
        &[DType::F32, DType::F32, DType::F32],
        BackendId::Vulkan,
    );
    assert!(!alts.is_empty(), "no Vulkan AddElementwise registration");
    let kernel = alts[0].kernel;

    // Per-input layouts: contiguous rank-1 [n].
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    let layouts = vec![layout.clone(), layout.clone(), layout];

    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::None,
    )
    .expect("kernel dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    let expected: Vec<f32> = a_data
        .iter()
        .zip(b_data.iter())
        .map(|(x, y)| x + y)
        .collect();
    assert_eq!(got, expected, "Vulkan Add f32 result mismatch");
}

/// Dispatch-table presence check (no GPU required). Confirms all
/// 6 binary-f32 ops register after V.2.A.
#[test]
fn vulkan_dispatch_binary_f32_registered() {
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);
    let key = [DType::F32, DType::F32, DType::F32];
    for op in [
        OpKind::AddElementwise,
        OpKind::SubElementwise,
        OpKind::MulElementwise,
        OpKind::DivElementwise,
        OpKind::MaximumElementwise,
        OpKind::MinimumElementwise,
    ] {
        let alts = table.lookup_alternatives(op, &key, BackendId::Vulkan);
        assert_eq!(
            alts.len(),
            1,
            "expected 1 Vulkan alternative for {op:?} f32 after register_vulkan_kernels, got {}",
            alts.len(),
        );
    }
}

/// Helper for V.2.A binary correctness tests — uploads `a` and `b`,
/// dispatches `op` against `(F32, F32, F32)`, downloads + returns.
fn run_binary_f32(
    backend: &Arc<VulkanBackend>,
    op: OpKind,
    a_data: &[f32],
    b_data: &[f32],
) -> Vec<f32> {
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);
    let n = a_data.len();
    let a_storage = upload_f32(backend, a_data);
    let b_storage = upload_f32(backend, b_data);
    let out_bytes = backend.alloc_bytes_handle(n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let alts =
        table.lookup_alternatives(op, &[DType::F32, DType::F32, DType::F32], BackendId::Vulkan);
    let kernel = alts[0].kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    let layouts = vec![layout.clone(), layout.clone(), layout];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::None,
    )
    .expect("kernel dispatch");
    download_f32(backend, &out_arc.read().unwrap())
}

#[test]
#[ignore]
fn vulkan_dispatch_binary_sub_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_binary_f32(
        &backend,
        OpKind::SubElementwise,
        &[10.0, 20.0, 30.0, 40.0],
        &[1.0, 2.0, 3.0, 4.0],
    );
    assert_eq!(got, vec![9.0, 18.0, 27.0, 36.0]);
}

#[test]
#[ignore]
fn vulkan_dispatch_binary_mul_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_binary_f32(
        &backend,
        OpKind::MulElementwise,
        &[2.0, 3.0, 4.0, 5.0],
        &[10.0, 10.0, 10.0, 10.0],
    );
    assert_eq!(got, vec![20.0, 30.0, 40.0, 50.0]);
}

#[test]
#[ignore]
fn vulkan_dispatch_binary_div_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_binary_f32(
        &backend,
        OpKind::DivElementwise,
        &[100.0, 80.0, 60.0, 40.0],
        &[2.0, 4.0, 5.0, 8.0],
    );
    assert_eq!(got, vec![50.0, 20.0, 12.0, 5.0]);
}

#[test]
#[ignore]
fn vulkan_dispatch_binary_maximum_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_binary_f32(
        &backend,
        OpKind::MaximumElementwise,
        &[1.0, 5.0, 3.0, 7.0],
        &[2.0, 4.0, 6.0, 1.0],
    );
    assert_eq!(got, vec![2.0, 5.0, 6.0, 7.0]);
}

#[test]
#[ignore]
fn vulkan_dispatch_binary_minimum_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_binary_f32(
        &backend,
        OpKind::MinimumElementwise,
        &[1.0, 5.0, 3.0, 7.0],
        &[2.0, 4.0, 6.0, 1.0],
    );
    assert_eq!(got, vec![1.0, 4.0, 3.0, 1.0]);
}

/// KISS §6.13 on Vulkan — the SHIPPABLE half: min/max PROPAGATE NaN a-first.
///
/// Vulkan `max`/`min` lower to SPIR-V FMax/FMin (spirv-dis confirmed, all four
/// dtypes), which on the RTX 4070 DROP NaN (measured: returned the non-NaN
/// operand). This fix adds an explicit `isnan(x)?x : isnan(y)?y : max(x,y)`
/// guard, matching CPU (#67) and CUDA. Born-red: on the pre-fix `.spv`,
/// `max(1,NaN)` returned `1.0` (NaN dropped) — this asserts it is NaN after.
/// The `isnan` guard is explicit control flow the driver cannot fuse, so the
/// propagation is reliable across devices (unlike the ±0 half — see below).
///
/// ⚠️ The ±0 a-bias half of §6.13 is NOT fixed here and is only MEASURED, never
/// asserted: the NVIDIA driver re-fuses any source-level select (a float select
/// AND a bit-cast select — BOTH measured) back into a canonicalizing FMNMX, so
/// ±0 stays the driver's canonical (max→+0, min→-0). `spirv-dis` shows a clean
/// IR regardless — it proves what we EMIT, never what EXECUTES; only the live
/// probe measures the latter. See the Vulkan ±0 gap.
///
/// Run:  pwsh scripts/gpu-run.ps1 -Project fuel-dispatch -- cargo test -p \
///   fuel-dispatch --features vulkan --test vulkan_dispatch_live \
///   vulkan_minmax_nan -- --ignored --nocapture
#[test]
#[ignore]
fn vulkan_minmax_nan_propagates_a_first() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let a_nan = f32::from_bits(0x7fc0_1234); // distinct payloads prove a-first
    let b_nan = f32::from_bits(0x7fc0_5678);
    let one = 1.0_f32;

    for op in [OpKind::MaximumElementwise, OpKind::MinimumElementwise] {
        // NaN in either position propagates (FMax would drop it to `one`).
        let r = run_binary_f32(&backend, op, &[a_nan, one], &[one, a_nan]);
        assert_eq!(
            r[0].to_bits(),
            a_nan.to_bits(),
            "{op:?}(NaN,1) must propagate a's NaN, got {:08x}",
            r[0].to_bits()
        );
        assert!(
            r[1].is_nan(),
            "{op:?}(1,NaN) must propagate NaN, not drop to 1 (got {:08x})",
            r[1].to_bits()
        );
        // a-first: with a NaN in BOTH operands, a's payload wins.
        let both = run_binary_f32(&backend, op, &[a_nan], &[b_nan]);
        assert_eq!(
            both[0].to_bits(),
            a_nan.to_bits(),
            "{op:?}(NaNa,NaNb) must keep a's payload (a-first)"
        );
    }

    // ±0 is the driver's canonical here and is NOT asserted (device-dependent —
    // see the Vulkan ±0 gap). Printed for diagnosis / re-measurement elsewhere.
    let z = run_binary_f32(
        &backend,
        OpKind::MaximumElementwise,
        &[0.0_f32, -0.0],
        &[-0.0, 0.0],
    );
    let s = |v: f32| if v.is_sign_negative() { '-' } else { '+' };
    eprintln!(
        "VULKAN ±0 NOTE (device-dependent, NOT asserted): max(+0,-0)={} max(-0,+0)={}",
        s(z[0]),
        s(z[1])
    );
}

// ---- V.2.B — unary f32 ---------------------------------------------------

/// Dispatch-table presence check for all 13 unary-f32 ops.
#[test]
fn vulkan_dispatch_unary_f32_registered() {
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);
    let key = [DType::F32, DType::F32];
    for op in [
        OpKind::NegElementwise,
        OpKind::SqrElementwise,
        OpKind::SqrtElementwise,
        OpKind::ExpElementwise,
        OpKind::LogElementwise,
        OpKind::SinElementwise,
        OpKind::CosElementwise,
        OpKind::TanhElementwise,
        OpKind::SigmoidElementwise,
        OpKind::SiluElementwise,
        OpKind::GeluTanhElementwise,
        OpKind::ReluElementwise,
        OpKind::StepElementwise,
    ] {
        let alts = table.lookup_alternatives(op, &key, BackendId::Vulkan);
        assert_eq!(
            alts.len(),
            1,
            "expected 1 Vulkan alternative for {op:?} f32, got {}",
            alts.len(),
        );
    }
}

fn run_unary_f32(backend: &Arc<VulkanBackend>, op: OpKind, a_data: &[f32]) -> Vec<f32> {
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);
    let n = a_data.len();
    let a_storage = upload_f32(backend, a_data);
    let out_bytes = backend.alloc_bytes_handle(n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let a_arc = Arc::new(RwLock::new(a_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let alts = table.lookup_alternatives(op, &[DType::F32, DType::F32], BackendId::Vulkan);
    let kernel = alts[0].kernel;
    // Rank-1 contiguous layout — the stride-aware unary kernel needs
    // the input layout to pack into Params; the contig flag in Params
    // takes the fast path.
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    let layouts = vec![layout.clone(), layout];
    kernel(
        &[Arc::clone(&a_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::None,
    )
    .expect("kernel dispatch");
    download_f32(backend, &out_arc.read().unwrap())
}

/// Tolerance helper for transcendental ops where bit-exact won't hold.
fn assert_close(got: &[f32], expected: &[f32], rel_tol: f32, abs_tol: f32) {
    assert_eq!(got.len(), expected.len(), "length mismatch");
    for (i, (a, b)) in got.iter().zip(expected.iter()).enumerate() {
        let diff = (a - b).abs();
        let rel = diff / a.abs().max(b.abs()).max(1e-12);
        assert!(
            diff < abs_tol || rel < rel_tol,
            "[{i}] got={a} want={b} diff={diff} rel={rel}",
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_neg_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_unary_f32(&backend, OpKind::NegElementwise, &[1.0, -2.0, 3.0, -4.0]);
    assert_eq!(got, vec![-1.0, 2.0, -3.0, 4.0]);
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_sqr_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_unary_f32(&backend, OpKind::SqrElementwise, &[1.0, 2.0, -3.0, 4.0]);
    assert_eq!(got, vec![1.0, 4.0, 9.0, 16.0]);
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_sqrt_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_unary_f32(&backend, OpKind::SqrtElementwise, &[1.0, 4.0, 9.0, 16.0]);
    assert_eq!(got, vec![1.0, 2.0, 3.0, 4.0]);
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_exp_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_unary_f32(&backend, OpKind::ExpElementwise, &[0.0, 1.0, 2.0]);
    let want = [
        1.0,
        std::f32::consts::E,
        std::f32::consts::E * std::f32::consts::E,
    ];
    assert_close(&got, &want, 1e-5, 1e-5);
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_log_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_unary_f32(
        &backend,
        OpKind::LogElementwise,
        &[
            1.0,
            std::f32::consts::E,
            std::f32::consts::E * std::f32::consts::E,
        ],
    );
    assert_close(&got, &[0.0, 1.0, 2.0], 1e-5, 1e-5);
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_sin_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_unary_f32(
        &backend,
        OpKind::SinElementwise,
        &[0.0, std::f32::consts::PI / 2.0, std::f32::consts::PI],
    );
    assert_close(&got, &[0.0, 1.0, 0.0], 1e-5, 1e-5);
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_cos_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_unary_f32(
        &backend,
        OpKind::CosElementwise,
        &[0.0, std::f32::consts::PI / 2.0, std::f32::consts::PI],
    );
    assert_close(&got, &[1.0, 0.0, -1.0], 1e-5, 1e-5);
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_tanh_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_unary_f32(&backend, OpKind::TanhElementwise, &[0.0, 1.0, -1.0]);
    assert_close(&got, &[0.0, 0.7615942, -0.7615942], 1e-5, 1e-5);
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_sigmoid_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_unary_f32(&backend, OpKind::SigmoidElementwise, &[0.0, 1.0, -1.0]);
    assert_close(&got, &[0.5, 0.7310586, 0.26894143], 1e-5, 1e-5);
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_silu_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    // silu(x) = x * sigmoid(x): silu(0) = 0, silu(1) = 0.7310586
    let got = run_unary_f32(&backend, OpKind::SiluElementwise, &[0.0, 1.0, -1.0]);
    assert_close(&got, &[0.0, 0.7310586, -0.26894143], 1e-5, 1e-5);
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_gelu_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    // gelu_tanh approximation: gelu_tanh(0)=0, gelu_tanh(1)≈0.8413, gelu_tanh(-1)≈-0.1587
    let got = run_unary_f32(&backend, OpKind::GeluTanhElementwise, &[0.0, 1.0, -1.0]);
    assert_close(&got, &[0.0, 0.8413, -0.1587], 1e-3, 1e-3);
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_relu_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_unary_f32(
        &backend,
        OpKind::ReluElementwise,
        &[-2.0, -1.0, 0.0, 1.0, 2.0],
    );
    assert_eq!(got, vec![0.0, 0.0, 0.0, 1.0, 2.0]);
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_step_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    // step(x) = 1 if x > 0 else 0
    let got = run_unary_f32(
        &backend,
        OpKind::StepElementwise,
        &[-2.0, -0.5, 0.0, 0.5, 2.0],
    );
    // step at 0 — unary.slang's exact convention varies; check the
    // unambiguous values + relax 0.
    assert_eq!(got[0], 0.0);
    assert_eq!(got[1], 0.0);
    assert_eq!(got[3], 1.0);
    assert_eq!(got[4], 1.0);
}

/// Rank-2 contiguous Add: proves the kernel handles multi-dim
/// shapes through the dispatch wrapper. Strided-input correctness
/// is exercised by V.2's downstream tests once more ops + view ops
/// land; this V.1.C smoke-test sticks to contiguous shapes.
#[test]
#[ignore]
fn vulkan_dispatch_binary_add_f32_rank2_contig() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let a_data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
    let b_data: Vec<f32> = vec![10.0, 20.0, 30.0, 40.0, 50.0, 60.0];

    let a_storage = upload_f32(&backend, &a_data);
    let b_storage = upload_f32(&backend, &b_data);
    let out_bytes = backend.alloc_bytes_handle(6 * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::AddElementwise,
        &[DType::F32, DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;

    let layout = Layout::contiguous(Shape::from_dims(&[2, 3]));
    let layouts = vec![layout.clone(), layout.clone(), layout];

    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::None,
    )
    .expect("kernel dispatch (rank-2 contig)");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    assert_eq!(got, vec![11.0, 22.0, 33.0, 44.0, 55.0, 66.0]);
}

// ===========================================================================
// V.2.C — Softmax + RmsNorm last-dim (Fused ops, f32)
// ===========================================================================

/// Presence check: SoftmaxLastDim + RmsNormLastDim register on
/// `[F32, F32]` (unary shape); Rope registers on the canonical
/// `[F32, F32, F32, F32]` 4-dtype key that matches CPU's
/// `register_cpu_kernels` registration (x, cos, sin, out).
#[test]
fn vulkan_dispatch_softmax_norm_registered() {
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);
    let unary = [DType::F32, DType::F32];
    for op in [OpKind::SoftmaxLastDim, OpKind::RmsNormLastDim] {
        let alts = table.lookup_alternatives(op, &unary, BackendId::Vulkan);
        assert_eq!(
            alts.len(),
            1,
            "expected 1 Vulkan alternative for {op:?} f32 after register_vulkan_kernels, got {}",
            alts.len(),
        );
    }
    let rope_key = [DType::F32, DType::F32, DType::F32, DType::F32];
    let alts = table.lookup_alternatives(OpKind::Rope, &rope_key, BackendId::Vulkan);
    assert_eq!(
        alts.len(),
        1,
        "expected 1 Vulkan alternative for Rope [x,cos,sin,out]=f32 after register_vulkan_kernels, got {}",
        alts.len(),
    );
}

/// Softmax row-wise: 2 rows × 4 cols. Each row should sum to ~1.0.
#[test]
#[ignore]
fn vulkan_dispatch_softmax_last_dim_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let n = outer * last;
    let host: Vec<f32> = vec![
        // row 0
        1.0, 2.0, 3.0, 4.0, // row 1 — shift-invariant
        -1.0, 0.0, 1.0, 2.0,
    ];

    let in_storage = upload_f32(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::SoftmaxLastDim,
        &[DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    let layouts = vec![layout.clone(), layout];
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::SoftmaxLastDim {
            outer_count: outer,
            last_dim: last,
        },
    )
    .expect("softmax dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    // Row sums to ~1
    for row in 0..outer {
        let s: f32 = got[row * last..(row + 1) * last].iter().sum();
        assert!((s - 1.0).abs() < 1e-5, "softmax row {row} sum {s} != 1.0");
    }
    // Both rows have the same shape (shift-invariant), so the two
    // rows should agree element-wise.
    for c in 0..last {
        let a = got[c];
        let b = got[last + c];
        assert!(
            (a - b).abs() < 1e-5,
            "softmax shift-invariance broken at col {c}: row0={a} row1={b}"
        );
    }
}

// ===========================================================================
// V.3.E — f16 unary ops (native float16_t)
// ===========================================================================

fn upload_f16(backend: &Arc<VulkanBackend>, host: &[half::f16]) -> Storage {
    let bytes: &[u8] = bytemuck::cast_slice(host);
    let vk_bytes = backend
        .upload_bytes_handle(bytes)
        .expect("vulkan upload_bytes_handle f16");
    Storage::new(BackendStorage::Vulkan(vk_bytes), DType::F16)
}

fn download_f16(backend: &Arc<VulkanBackend>, s: &Storage) -> Vec<half::f16> {
    let bytes = match &s.inner {
        BackendStorage::Vulkan(v) => backend.download_bytes(v).expect("d2h"),
        _ => panic!("not on Vulkan"),
    };
    bytemuck::cast_slice::<u8, half::f16>(&bytes).to_vec()
}

fn run_unary_f16(backend: &Arc<VulkanBackend>, op: OpKind, data: &[half::f16]) -> Vec<half::f16> {
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);
    let n = data.len();
    let in_storage = upload_f16(backend, data);
    let out_bytes = backend.alloc_bytes_handle(n * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let kernel =
        table.lookup_alternatives(op, &[DType::F16, DType::F16], BackendId::Vulkan)[0].kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    let layouts = vec![layout.clone(), layout];
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::None,
    )
    .expect("kernel dispatch");
    download_f16(backend, &out_arc.read().unwrap())
}

fn run_binary_f16(
    backend: &Arc<VulkanBackend>,
    op: OpKind,
    a_data: &[half::f16],
    b_data: &[half::f16],
) -> Vec<half::f16> {
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);
    let n = a_data.len();
    let a_storage = upload_f16(backend, a_data);
    let b_storage = upload_f16(backend, b_data);
    let out_bytes = backend.alloc_bytes_handle(n * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let kernel =
        table.lookup_alternatives(op, &[DType::F16, DType::F16, DType::F16], BackendId::Vulkan)[0]
            .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    let layouts = vec![layout.clone(), layout.clone(), layout];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::None,
    )
    .expect("kernel dispatch");
    download_f16(backend, &out_arc.read().unwrap())
}

#[test]
fn vulkan_dispatch_binary_f16_registered() {
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);
    let key = [DType::F16, DType::F16, DType::F16];
    for op in [
        OpKind::AddElementwise,
        OpKind::SubElementwise,
        OpKind::MulElementwise,
        OpKind::DivElementwise,
        OpKind::MaximumElementwise,
        OpKind::MinimumElementwise,
    ] {
        let alts = table.lookup_alternatives(op, &key, BackendId::Vulkan);
        assert_eq!(
            alts.len(),
            1,
            "expected 1 Vulkan alt for {op:?} f16, got {}",
            alts.len()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_binary_add_f16() {
    use half::f16;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let a: Vec<f16> = vec![1.0, 2.0, 3.0, 4.0]
        .into_iter()
        .map(f16::from_f32)
        .collect();
    let b: Vec<f16> = vec![10.0, 20.0, 30.0, 40.0]
        .into_iter()
        .map(f16::from_f32)
        .collect();
    let got = run_binary_f16(&backend, OpKind::AddElementwise, &a, &b);
    let expected: Vec<f32> = vec![11.0, 22.0, 33.0, 44.0];
    for (g, e) in got.iter().zip(expected.iter()) {
        assert!(
            (g.to_f32() - e).abs() < 0.01,
            "add f16: got={}, expected={e}",
            g.to_f32()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_binary_mul_f16() {
    use half::f16;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let a: Vec<f16> = vec![1.5, -2.0, 0.5]
        .into_iter()
        .map(f16::from_f32)
        .collect();
    let b: Vec<f16> = vec![2.0, 3.0, 4.0].into_iter().map(f16::from_f32).collect();
    let got = run_binary_f16(&backend, OpKind::MulElementwise, &a, &b);
    let expected: Vec<f32> = vec![3.0, -6.0, 2.0];
    for (g, e) in got.iter().zip(expected.iter()) {
        assert!(
            (g.to_f32() - e).abs() < 0.01,
            "mul f16: got={}, expected={e}",
            g.to_f32()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_binary_maximum_f16() {
    use half::f16;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let a: Vec<f16> = vec![1.0, -2.0, 3.0]
        .into_iter()
        .map(f16::from_f32)
        .collect();
    let b: Vec<f16> = vec![-1.0, 2.0, 0.5]
        .into_iter()
        .map(f16::from_f32)
        .collect();
    let got = run_binary_f16(&backend, OpKind::MaximumElementwise, &a, &b);
    let expected: Vec<f32> = vec![1.0, 2.0, 3.0];
    for (g, e) in got.iter().zip(expected.iter()) {
        assert!(
            (g.to_f32() - e).abs() < 0.01,
            "max f16: got={}, expected={e}",
            g.to_f32()
        );
    }
}

#[test]
fn vulkan_dispatch_unary_f16_registered() {
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);
    let key = [DType::F16, DType::F16];
    for op in [
        OpKind::NegElementwise,
        OpKind::SqrElementwise,
        OpKind::ReluElementwise,
        OpKind::TanhElementwise,
        OpKind::SiluElementwise,
        OpKind::GeluTanhElementwise,
    ] {
        let alts = table.lookup_alternatives(op, &key, BackendId::Vulkan);
        assert_eq!(
            alts.len(),
            1,
            "expected 1 Vulkan alt for {op:?} f16, got {}",
            alts.len()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_neg_f16() {
    use half::f16;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let host: Vec<f16> = vec![1.0, -2.5, 0.0, 3.0]
        .into_iter()
        .map(f16::from_f32)
        .collect();
    let got = run_unary_f16(&backend, OpKind::NegElementwise, &host);
    for (g, src) in got.iter().zip(host.iter()) {
        assert!(
            (g.to_f32() - (-src.to_f32())).abs() < 1e-3,
            "neg f16: src={}, got={}",
            src.to_f32(),
            g.to_f32()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_relu_f16() {
    use half::f16;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let host: Vec<f16> = vec![-1.0, 0.0, 1.0, -2.5, 2.5]
        .into_iter()
        .map(f16::from_f32)
        .collect();
    let got = run_unary_f16(&backend, OpKind::ReluElementwise, &host);
    let expected = [0.0, 0.0, 1.0, 0.0, 2.5];
    for (g, e) in got.iter().zip(expected.iter()) {
        assert!(
            (g.to_f32() - e).abs() < 1e-3,
            "relu f16: got={}, expected={e}",
            g.to_f32()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_tanh_f16() {
    use half::f16;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let host: Vec<f16> = vec![-1.0, 0.0, 1.0, 2.0]
        .into_iter()
        .map(f16::from_f32)
        .collect();
    let got = run_unary_f16(&backend, OpKind::TanhElementwise, &host);
    for (g, src) in got.iter().zip(host.iter()) {
        // f16 tanh has visible rounding error vs f32; widen tolerance.
        let expected = src.to_f32().tanh();
        assert!(
            (g.to_f32() - expected).abs() < 0.01,
            "tanh f16: src={}, got={}, expected={expected}",
            src.to_f32(),
            g.to_f32()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_sigmoid_f16() {
    use half::f16;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let host: Vec<f16> = vec![-2.0, 0.0, 2.0, 5.0]
        .into_iter()
        .map(f16::from_f32)
        .collect();
    let got = run_unary_f16(&backend, OpKind::SigmoidElementwise, &host);
    for (g, src) in got.iter().zip(host.iter()) {
        let expected = 1.0 / (1.0 + (-src.to_f32()).exp());
        assert!(
            (g.to_f32() - expected).abs() < 0.01,
            "sigmoid f16: src={}, got={}, expected={expected}",
            src.to_f32(),
            g.to_f32()
        );
    }
}

// ===========================================================================
// V.3.E.5 — f64 unary + binary (native double)
// ===========================================================================

fn upload_f64(backend: &Arc<VulkanBackend>, host: &[f64]) -> Storage {
    let bytes: &[u8] = bytemuck::cast_slice(host);
    let vk_bytes = backend
        .upload_bytes_handle(bytes)
        .expect("vulkan upload_bytes_handle f64");
    Storage::new(BackendStorage::Vulkan(vk_bytes), DType::F64)
}

fn download_f64(backend: &Arc<VulkanBackend>, s: &Storage) -> Vec<f64> {
    let bytes = match &s.inner {
        BackendStorage::Vulkan(v) => backend.download_bytes(v).expect("d2h"),
        _ => panic!("not on Vulkan"),
    };
    bytemuck::cast_slice::<u8, f64>(&bytes).to_vec()
}

#[test]
fn vulkan_dispatch_f64_registered() {
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);
    for op in [
        OpKind::NegElementwise,
        OpKind::SqrElementwise,
        OpKind::ReluElementwise,
    ] {
        let alts = table.lookup_alternatives(op, &[DType::F64, DType::F64], BackendId::Vulkan);
        assert_eq!(alts.len(), 1, "expected 1 Vulkan alt for {op:?} f64");
    }
    for op in [OpKind::AddElementwise, OpKind::MulElementwise] {
        let alts =
            table.lookup_alternatives(op, &[DType::F64, DType::F64, DType::F64], BackendId::Vulkan);
        assert_eq!(alts.len(), 1, "expected 1 Vulkan alt for {op:?} f64 binary");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_neg_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host: Vec<f64> = vec![1.5, -2.25, 0.0, 3.125];
    let n = host.len();
    let in_storage = upload_f64(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(n * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let kernel = table.lookup_alternatives(
        OpKind::NegElementwise,
        &[DType::F64, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::None,
    )
    .expect("neg f64 dispatch");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    let expected: Vec<f64> = host.iter().map(|x| -x).collect();
    for (g, e) in got.iter().zip(expected.iter()) {
        assert!((g - e).abs() < 1e-12, "neg f64: got={g}, expected={e}");
    }
}

// ===========================================================================
// V.3.E.5 transcendentals — polynomial approximations in Slang
// ===========================================================================
//
// Target precision: 1e-12 relative error (matches the kernel's
// design target; far below libm's ULP-correct standard but
// adequate for inference / training workloads at f64). For
// composites (sigmoid / silu / gelu_tanh) we expect ~1e-11 because
// errors accumulate across the composed exp / tanh calls.

fn run_unary_f64(backend: &Arc<VulkanBackend>, op: OpKind, data: &[f64]) -> Vec<f64> {
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);
    let n = data.len();
    let in_storage = upload_f64(backend, data);
    let out_bytes = backend.alloc_bytes_handle(n * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let kernel =
        table.lookup_alternatives(op, &[DType::F64, DType::F64], BackendId::Vulkan)[0].kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::None,
    )
    .expect("kernel dispatch");
    download_f64(backend, &out_arc.read().unwrap())
}

/// Assert error ≤ `tol`. For "ordinary" magnitudes (≥ 1e-14) the
/// check is relative; for near-zero expected values it switches to
/// absolute (relative-error against a value of magnitude 1e-16 would
/// blow up at any actual difference). NaN and ±inf require bit
/// equality; ±0 of either sign treats as zero.
fn check_f64(label: &str, got: &[f64], expected: &[f64], tol: f64) -> f64 {
    assert_eq!(got.len(), expected.len(), "{label}: length mismatch");
    let mut worst: f64 = 0.0;
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        if e.is_nan() {
            assert!(g.is_nan(), "{label}[{i}]: expected NaN, got {g}");
            continue;
        }
        if e.is_infinite() {
            assert_eq!(
                g.is_sign_positive(),
                e.is_sign_positive(),
                "{label}[{i}]: sign mismatch (got {g}, expected {e})"
            );
            assert!(g.is_infinite(), "{label}[{i}]: expected {e}, got {g}");
            continue;
        }
        let abs_err = (g - e).abs();
        // When the expected magnitude is at f64-noise level (~1e-14),
        // the kernel returning an exact zero or a slightly different
        // tiny value is more accurate than libm — relative-error
        // comparison would falsely fail. Switch to absolute error.
        if e.abs() < 1e-14 {
            assert!(
                abs_err < 1e-14,
                "{label}[{i}]: near-zero expected {e}, got {g}, abs err {abs_err:e}"
            );
            continue;
        }
        let rel = abs_err / e.abs();
        if rel > worst {
            worst = rel;
        }
        assert!(
            rel <= tol,
            "{label}[{i}]: got {g}, expected {e}, rel err {rel:e} > tol {tol:e}"
        );
    }
    worst
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_exp_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    // Cover the full natural range: tiny, small, 1, larger, near-overflow,
    // negative, very negative (underflow neighborhood).
    // `2.71828` is a DELIBERATE decimal approximation sitting alongside the
    // exact `E` on the line above -- two distinct inputs, chosen so the sweep
    // covers both an exactly-representable constant and a nearby value that is
    // not. Clippy suggests replacing it WITH `E`, which would collapse the two
    // into a duplicate entry and silently DROP a test input. The lint is right
    // about what the literal approximates and wrong about what to do here.
    #[allow(clippy::approx_constant)]
    let host = vec![
        0.0,
        1e-10,
        0.5,
        1.0,
        std::f64::consts::E,
        2.71828,
        10.0,
        100.0,
        500.0,
        700.0,
        -1e-10,
        -0.5,
        -1.0,
        -10.0,
        -100.0,
        -700.0,
    ];
    let got = run_unary_f64(&backend, OpKind::ExpElementwise, &host);
    let expected: Vec<f64> = host.iter().map(|x| x.exp()).collect();
    // DD Horner + DD reconstruction + full-precision LN2_LO (48
    // mantissa bits, not the prior 21) bring exp to sub-1-ULP.
    // 5e-16 tolerance ≈ 2.3 ULP — well above observed 0.8 ULP at
    // exp(700), the worst case before the DD upgrade was 75 ULP.
    let worst = check_f64("exp f64", &got, &expected, 5e-16);
    eprintln!("exp f64: worst rel err = {worst:e}");
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_log_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let host = vec![
        f64::MIN_POSITIVE,
        1e-15,
        1e-10,
        0.001,
        0.5,
        1.0,
        std::f64::consts::E,
        10.0,
        100.0,
        1e10,
        1e100,
        1e300,
    ];
    let got = run_unary_f64(&backend, OpKind::LogElementwise, &host);
    let expected: Vec<f64> = host.iter().map(|x| x.ln()).collect();
    // Allow slightly looser tolerance for the very small / very large
    // inputs where range reduction is the bottleneck.
    let worst = check_f64("log f64", &got, &expected, 5e-12);
    eprintln!("log f64: worst rel err = {worst:e}");
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_sin_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    use std::f64::consts::{FRAC_PI_2, FRAC_PI_4, FRAC_PI_6, PI};

    // Payne-Hanek reduction in DD precision: sin/cos are ULP-correct
    // throughout the finite f64 range, including the prior failure
    // regime |x| > 6.6e6 that Cody-Waite couldn't handle. Tolerance
    // 5e-16 (~2.3 ULP) — observed worst case is ~0.9 ULP at the 5e6
    // boundary, where small true |sin(x)| amplifies the residual
    // polynomial error.
    let small = vec![
        0.0,
        1e-10,
        FRAC_PI_6,
        FRAC_PI_4,
        FRAC_PI_2,
        PI,
        1.5 * PI,
        2.0 * PI,
        -FRAC_PI_4,
        -PI,
        10.0,
        100.0,
        -100.0,
    ];
    let got = run_unary_f64(&backend, OpKind::SinElementwise, &small);
    let expected: Vec<f64> = small.iter().map(|x| x.sin()).collect();
    let worst_small = check_f64("sin f64 (|x| <= 100)", &got, &expected, 5e-16);
    eprintln!("sin f64 (|x| <= 100): worst rel err = {worst_small:e}");

    let moderate = vec![300.0, 500.0, -700.0, 1000.0];
    let got = run_unary_f64(&backend, OpKind::SinElementwise, &moderate);
    let expected: Vec<f64> = moderate.iter().map(|x| x.sin()).collect();
    let worst_mod = check_f64("sin f64 (|x| <= 1000)", &got, &expected, 5e-16);
    eprintln!("sin f64 (|x| <= 1000): worst rel err = {worst_mod:e}");

    let large = vec![1e5, 5e5, 1e6, 5e6, -1e6];
    let got = run_unary_f64(&backend, OpKind::SinElementwise, &large);
    let expected: Vec<f64> = large.iter().map(|x| x.sin()).collect();
    let worst_lg = check_f64("sin f64 (|x| <= 5e6)", &got, &expected, 5e-16);
    eprintln!("sin f64 (|x| <= 5e6): worst rel err = {worst_lg:e}");
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_sin_f64_huge() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    // The Payne-Hanek regime: inputs that the prior three-term
    // Cody-Waite reduction failed on (|x| > 6.6e6) but full PH
    // handles correctly up to ~2^53.
    //
    // Adversarial inputs:
    //   1e10, 1e12, 1e15  — magnitudes well past the Cody-Waite cutoff
    //   2^53 - 1          — at the PH precision boundary (where the
    //                       106-bit DD pair for x*(2/π) just barely
    //                       captures the integer + fractional split)
    //   2^53 / 3          — fractional bits near a half-integer of
    //                       x*(2/π); known to expose buggy PH reductions
    //   huge * π          — multiples of π where true sin ≈ 0;
    //                       catastrophic-cancellation case
    let host = vec![
        1.0e10_f64,
        1.0e12,
        1.0e15,
        (1u64 << 53) as f64 - 1.0,
        ((1u64 << 53) as f64) / 3.0,
        1.0e10 * std::f64::consts::PI,
        1.0e12 * std::f64::consts::PI,
    ];
    let got = run_unary_f64(&backend, OpKind::SinElementwise, &host);
    let expected: Vec<f64> = host.iter().map(|x| x.sin()).collect();
    // Wide tolerance: at these magnitudes f64 input precision itself
    // limits how much of the math signal survives — libm and our PH
    // both reduce the SAME f64 value of the input (which has lost
    // information about the true real x), so they should agree to
    // a few ULP. 5e-14 ≈ ~250 ULP gives margin for the polynomial's
    // ~3 ULP plus the reduction's ~1 ULP plus any environmental jitter.
    let worst = check_f64("sin f64 (huge)", &got, &expected, 5e-14);
    eprintln!("sin f64 (huge): worst rel err = {worst:e}");
    for (x, (g, e)) in host.iter().zip(got.iter().zip(expected.iter())) {
        eprintln!("  sin({x:e}) = {g:e}  (libm: {e:e})");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_cos_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    use std::f64::consts::{FRAC_PI_2, FRAC_PI_4, PI};

    let small = vec![
        0.0,
        FRAC_PI_4,
        FRAC_PI_2,
        PI,
        2.0 * PI,
        -FRAC_PI_4,
        -PI,
        1.0,
        5.0,
        100.0,
    ];
    let got = run_unary_f64(&backend, OpKind::CosElementwise, &small);
    let expected: Vec<f64> = small.iter().map(|x| x.cos()).collect();
    let worst_small = check_f64("cos f64 (|x| <= 100)", &got, &expected, 5e-16);
    eprintln!("cos f64 (|x| <= 100): worst rel err = {worst_small:e}");

    let moderate = vec![300.0, 1000.0, -500.0];
    let got = run_unary_f64(&backend, OpKind::CosElementwise, &moderate);
    let expected: Vec<f64> = moderate.iter().map(|x| x.cos()).collect();
    let worst_mod = check_f64("cos f64 (|x| <= 1000)", &got, &expected, 5e-16);
    eprintln!("cos f64 (|x| <= 1000): worst rel err = {worst_mod:e}");
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_tanh_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let host = vec![
        0.0, 1e-10, 0.5, 1.0, 2.0, 5.0, 10.0, 18.0, 50.0, -0.5, -1.0, -5.0, -18.0,
    ];
    let got = run_unary_f64(&backend, OpKind::TanhElementwise, &host);
    let expected: Vec<f64> = host.iter().map(|x| x.tanh()).collect();
    // tanh composes from exp so errors compound slightly. Tolerance
    // 5e-12 relative; at |x|>=18 the saturation case lands at ±1 exactly.
    let worst = check_f64("tanh f64", &got, &expected, 5e-12);
    eprintln!("tanh f64: worst rel err = {worst:e}");
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_sigmoid_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let host = vec![-10.0, -3.0, -1.0, -0.5, 0.0, 0.5, 1.0, 3.0, 10.0];
    let got = run_unary_f64(&backend, OpKind::SigmoidElementwise, &host);
    let expected: Vec<f64> = host.iter().map(|x| 1.0 / (1.0 + (-x).exp())).collect();
    let worst = check_f64("sigmoid f64", &got, &expected, 5e-12);
    eprintln!("sigmoid f64: worst rel err = {worst:e}");
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_silu_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let host = vec![-5.0, -1.0, -0.1, 0.0, 0.1, 1.0, 5.0];
    let got = run_unary_f64(&backend, OpKind::SiluElementwise, &host);
    let expected: Vec<f64> = host.iter().map(|x| x / (1.0 + (-x).exp())).collect();
    let worst = check_f64("silu f64", &got, &expected, 5e-12);
    eprintln!("silu f64: worst rel err = {worst:e}");
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_gelu_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let host = vec![-3.0, -1.0, -0.5, 0.0, 0.5, 1.0, 3.0];
    let got = run_unary_f64(&backend, OpKind::GeluTanhElementwise, &host);
    // Reference matches the kernel's tanh-approx form exactly.
    let expected: Vec<f64> = host
        .iter()
        .map(|x| {
            let inner = 0.7978845608028654 * (x + 0.044715 * x * x * x);
            0.5 * x * (1.0 + inner.tanh())
        })
        .collect();
    // Gelu composes tanh which composes exp — error budget around 1e-11.
    let worst = check_f64("gelu_tanh f64", &got, &expected, 1e-11);
    eprintln!("gelu_tanh f64: worst rel err = {worst:e}");
}

#[test]
#[ignore]
fn vulkan_dispatch_binary_add_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let a = vec![1.0_f64, 2.0, 3.0, 4.0];
    let b = vec![10.0_f64, 20.0, 30.0, 40.0];
    let n = a.len();
    let a_storage = upload_f64(&backend, &a);
    let b_storage = upload_f64(&backend, &b);
    let out_bytes = backend.alloc_bytes_handle(n * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::AddElementwise,
        &[DType::F64, DType::F64, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout.clone(), layout],
        &OpParams::None,
    )
    .expect("add f64 dispatch");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    assert_eq!(got, vec![11.0, 22.0, 33.0, 44.0]);
}

// ===========================================================================
// V.3.J — WriteSlice (in-place slab assign)
// ===========================================================================

/// WriteSlice: 1D slab. dst shape [8] init [0,0,0,0,0,0,0,0]; src [3]
/// = [10, 20, 30] written at range [2..5]. Expected dst:
/// [0, 0, 10, 20, 30, 0, 0, 0].
#[test]
#[ignore]
fn vulkan_dispatch_write_slice_b4_1d_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // Pre-fill dst with zeros (alloc_bytes_handle does not zero — use
    // upload_bytes_handle to set bytes explicitly).
    let dst_init = vec![0.0_f32; 8];
    let dst_bytes: &[u8] = bytemuck::cast_slice(&dst_init);
    let dst_vk = backend.upload_bytes_handle(dst_bytes).expect("upload dst");
    let dst_storage = Storage::new(BackendStorage::Vulkan(dst_vk), DType::F32);

    let src = vec![10.0_f32, 20.0, 30.0];
    let src_vk = backend
        .upload_bytes_handle(bytemuck::cast_slice(&src))
        .expect("upload src");
    let src_storage = Storage::new(BackendStorage::Vulkan(src_vk), DType::F32);

    let src_arc = Arc::new(RwLock::new(src_storage));
    let dst_arc = Arc::new(RwLock::new(dst_storage));

    let kernel = table.lookup_alternatives(
        OpKind::WriteSlice,
        &[DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[3])),
        Layout::contiguous(Shape::from_dims(&[8])),
    ];
    kernel(
        &[Arc::clone(&src_arc)],
        &mut [Arc::clone(&dst_arc)],
        &layouts,
        &OpParams::WriteSlice {
            dest_shape: vec![8],
            ranges: vec![(2, 5)],
            deferred_dyn_offset: None,
        },
    )
    .expect("write_slice dispatch");

    let got = download_f32(&backend, &dst_arc.read().unwrap());
    assert_eq!(got, vec![0.0, 0.0, 10.0, 20.0, 30.0, 0.0, 0.0, 0.0]);
}

/// WriteSlice: 2D slab — write a 2×2 block at offset (1, 1) inside
/// a 3×4 destination.
#[test]
#[ignore]
fn vulkan_dispatch_write_slice_b4_2d_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // dst 3×4 init = matrix of 0.5 (so we can see what survives).
    let dst_init = vec![0.5_f32; 12];
    let dst_vk = backend
        .upload_bytes_handle(bytemuck::cast_slice(&dst_init))
        .expect("upload dst");
    let dst_storage = Storage::new(BackendStorage::Vulkan(dst_vk), DType::F32);

    // src 2×2 = [[10, 20], [30, 40]] (row-major)
    let src = vec![10.0_f32, 20.0, 30.0, 40.0];
    let src_vk = backend
        .upload_bytes_handle(bytemuck::cast_slice(&src))
        .expect("upload src");
    let src_storage = Storage::new(BackendStorage::Vulkan(src_vk), DType::F32);

    let src_arc = Arc::new(RwLock::new(src_storage));
    let dst_arc = Arc::new(RwLock::new(dst_storage));

    let kernel = table.lookup_alternatives(
        OpKind::WriteSlice,
        &[DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[2, 2])),
        Layout::contiguous(Shape::from_dims(&[3, 4])),
    ];
    kernel(
        &[Arc::clone(&src_arc)],
        &mut [Arc::clone(&dst_arc)],
        &layouts,
        &OpParams::WriteSlice {
            dest_shape: vec![3, 4],
            ranges: vec![(1, 3), (1, 3)],
            deferred_dyn_offset: None,
        },
    )
    .expect("write_slice 2d dispatch");

    let got = download_f32(&backend, &dst_arc.read().unwrap());
    // Layout (row-major, 3×4):
    //   row 0: 0.5 0.5 0.5 0.5
    //   row 1: 0.5 10  20  0.5
    //   row 2: 0.5 30  40  0.5
    assert_eq!(
        got,
        vec![
            0.5, 0.5, 0.5, 0.5, 0.5, 10.0, 20.0, 0.5, 0.5, 30.0, 40.0, 0.5,
        ]
    );
}

/// WriteSlice b2: f16 KV-cache slab write. 1×8 slab into a 4×8 dst
/// at row 2. last-dim slab is the full 8 cols (range_start=0,
/// src_shape=8) so even-alignment holds.
#[test]
#[ignore]
fn vulkan_dispatch_write_slice_b2_f16_kv_cache() {
    use half::f16;

    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let seq = 4usize;
    let head_dim = 8usize;

    let dst_init: Vec<f16> = vec![f16::from_f32(-1.0); seq * head_dim];
    let dst_vk = backend
        .upload_bytes_handle(bytemuck::cast_slice(&dst_init))
        .expect("upload");
    let dst_storage = Storage::new(BackendStorage::Vulkan(dst_vk), DType::F16);

    let src: Vec<f16> = (0..head_dim)
        .map(|i| f16::from_f32(i as f32 + 100.0))
        .collect();
    let src_vk = backend
        .upload_bytes_handle(bytemuck::cast_slice(&src))
        .expect("upload");
    let src_storage = Storage::new(BackendStorage::Vulkan(src_vk), DType::F16);

    let src_arc = Arc::new(RwLock::new(src_storage));
    let dst_arc = Arc::new(RwLock::new(dst_storage));

    let kernel = table.lookup_alternatives(
        OpKind::WriteSlice,
        &[DType::F16, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[1, head_dim])),
        Layout::contiguous(Shape::from_dims(&[seq, head_dim])),
    ];
    kernel(
        &[Arc::clone(&src_arc)],
        &mut [Arc::clone(&dst_arc)],
        &layouts,
        &OpParams::WriteSlice {
            dest_shape: vec![seq, head_dim],
            ranges: vec![(2, 3), (0, head_dim)],
            deferred_dyn_offset: None,
        },
    )
    .expect("write_slice b2 dispatch");

    let raw = match &dst_arc.read().unwrap().inner {
        BackendStorage::Vulkan(v) => backend.download_bytes(v).expect("d2h"),
        _ => panic!(),
    };
    let got: Vec<f16> = bytemuck::cast_slice::<u8, f16>(&raw).to_vec();
    // `0 * head_dim + j` is the ROW-MAJOR INDEX IDIOM in a matched series with
    // rows 1..3 below. The multiply is alignment, not arithmetic: writing bare
    // `j` for row 0 makes the row index invisible on exactly one line of four
    // and breaks the visual correspondence a reader checks the rows by.
    // Same disposition as the three sites in `fuel-core`. `0 *` trips
    // `erasing_op` and `1 *` trips `identity_op` -- one idiom, two lint names.
    #[allow(clippy::erasing_op, clippy::identity_op)]
    for j in 0..head_dim {
        let row0 = got[0 * head_dim + j];
        let row1 = got[1 * head_dim + j];
        let row2 = got[2 * head_dim + j];
        let row3 = got[3 * head_dim + j];
        assert_eq!(row0.to_f32(), -1.0);
        assert_eq!(row1.to_f32(), -1.0);
        assert_eq!(row2.to_f32(), j as f32 + 100.0);
        assert_eq!(row3.to_f32(), -1.0);
    }
}

/// WriteSlice b8: f64 1D slab into a [6] dst at range [1..4].
#[test]
#[ignore]
fn vulkan_dispatch_write_slice_b8_f64_1d() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let dst_init: Vec<f64> = vec![0.0; 6];
    let dst_vk = backend
        .upload_bytes_handle(bytemuck::cast_slice(&dst_init))
        .expect("upload");
    let dst_storage = Storage::new(BackendStorage::Vulkan(dst_vk), DType::F64);

    let src: Vec<f64> = vec![100.0, 200.0, 300.0];
    let src_vk = backend
        .upload_bytes_handle(bytemuck::cast_slice(&src))
        .expect("upload");
    let src_storage = Storage::new(BackendStorage::Vulkan(src_vk), DType::F64);

    let src_arc = Arc::new(RwLock::new(src_storage));
    let dst_arc = Arc::new(RwLock::new(dst_storage));

    let kernel = table.lookup_alternatives(
        OpKind::WriteSlice,
        &[DType::F64, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[3])),
        Layout::contiguous(Shape::from_dims(&[6])),
    ];
    kernel(
        &[Arc::clone(&src_arc)],
        &mut [Arc::clone(&dst_arc)],
        &layouts,
        &OpParams::WriteSlice {
            dest_shape: vec![6],
            ranges: vec![(1, 4)],
            deferred_dyn_offset: None,
        },
    )
    .expect("write_slice b8 dispatch");

    let raw = match &dst_arc.read().unwrap().inner {
        BackendStorage::Vulkan(v) => backend.download_bytes(v).expect("d2h"),
        _ => panic!(),
    };
    let got: Vec<f64> = bytemuck::cast_slice::<u8, f64>(&raw).to_vec();
    assert_eq!(got, vec![0.0, 100.0, 200.0, 300.0, 0.0, 0.0]);
}

/// WriteSlice: KV-cache shape — write a single token's K vector into
/// position 2 of a [seq=4, head_dim=8] cache. Models the inference
/// hot path that backs Op::WriteSlice in the lazy graph.
#[test]
#[ignore]
fn vulkan_dispatch_write_slice_b4_kv_cache_shape() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let seq = 4usize;
    let head_dim = 8usize;
    let dst_init = vec![-1.0_f32; seq * head_dim];
    let dst_vk = backend
        .upload_bytes_handle(bytemuck::cast_slice(&dst_init))
        .expect("upload dst");
    let dst_storage = Storage::new(BackendStorage::Vulkan(dst_vk), DType::F32);

    let src: Vec<f32> = (0..head_dim).map(|i| i as f32 + 100.0).collect();
    let src_vk = backend
        .upload_bytes_handle(bytemuck::cast_slice(&src))
        .expect("upload src");
    let src_storage = Storage::new(BackendStorage::Vulkan(src_vk), DType::F32);

    let src_arc = Arc::new(RwLock::new(src_storage));
    let dst_arc = Arc::new(RwLock::new(dst_storage));

    let kernel = table.lookup_alternatives(
        OpKind::WriteSlice,
        &[DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[1, head_dim])),
        Layout::contiguous(Shape::from_dims(&[seq, head_dim])),
    ];
    kernel(
        &[Arc::clone(&src_arc)],
        &mut [Arc::clone(&dst_arc)],
        &layouts,
        &OpParams::WriteSlice {
            dest_shape: vec![seq, head_dim],
            ranges: vec![(2, 3), (0, head_dim)],
            deferred_dyn_offset: None,
        },
    )
    .expect("write_slice kv dispatch");

    let got = download_f32(&backend, &dst_arc.read().unwrap());
    // Rows 0, 1 stay -1; row 2 = src; row 3 stays -1.
    // Row-major index idiom in a matched series -- see the note above.
    #[allow(clippy::erasing_op, clippy::identity_op)]
    for j in 0..head_dim {
        assert_eq!(got[0 * head_dim + j], -1.0);
        assert_eq!(got[1 * head_dim + j], -1.0);
        assert_eq!(got[2 * head_dim + j], (j as f32) + 100.0);
        assert_eq!(got[3 * head_dim + j], -1.0);
    }
}

// ===========================================================================
// V.3.B — Cast (f32 ↔ f16, f32 ↔ bf16)
// ===========================================================================

#[test]
#[ignore]
fn vulkan_dispatch_cast_f32_to_bf16() {
    use half::bf16;

    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host_f32: Vec<f32> = vec![1.0, -2.5, 3.25, 0.0, 1e-3, 1e3];
    let n = host_f32.len();

    let in_storage = upload_f32(&backend, &host_f32);
    let out_bytes_h = backend.alloc_bytes_handle(n * 2).expect("alloc bf16 out");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::BF16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Cast, &[DType::F32, DType::BF16], BackendId::Vulkan)[0]
            .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[n])),
        Layout::contiguous(Shape::from_dims(&[n])),
    ];
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::None,
    )
    .expect("cast f32→bf16 dispatch");

    // Download bf16 bytes and reinterpret.
    let raw = match &out_arc.read().unwrap().inner {
        BackendStorage::Vulkan(v) => backend.download_bytes(v).expect("d2h"),
        _ => panic!("not on Vulkan"),
    };
    let got: Vec<bf16> = bytemuck::cast_slice::<u8, bf16>(&raw).to_vec();
    for (i, (g, src)) in got.iter().zip(host_f32.iter()).enumerate() {
        let expected = bf16::from_f32(*src);
        // The kernel uses truncation (bits >> 16), not round-to-nearest;
        // bf16::from_f32 also truncates. Should be bit-identical.
        assert_eq!(
            g.to_bits(),
            expected.to_bits(),
            "cast[{i}]: f32={src}, got bf16 bits={:#x}, expected {:#x}",
            g.to_bits(),
            expected.to_bits()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_cast_bf16_to_f32() {
    use half::bf16;

    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host_bf16: Vec<bf16> = vec![
        bf16::from_f32(1.0),
        bf16::from_f32(-2.5),
        bf16::from_f32(3.25),
        bf16::from_f32(0.0),
        bf16::from_f32(1e-3),
        bf16::from_f32(1e3),
    ];
    let n = host_bf16.len();

    let bf16_bytes: &[u8] = bytemuck::cast_slice(&host_bf16);
    let in_vk = backend
        .upload_bytes_handle(bf16_bytes)
        .expect("upload bf16");
    let in_storage = Storage::new(BackendStorage::Vulkan(in_vk), DType::BF16);
    let out_bytes_h = backend.alloc_bytes_handle(n * 4).expect("alloc f32 out");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Cast, &[DType::BF16, DType::F32], BackendId::Vulkan)[0]
            .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[n])),
        Layout::contiguous(Shape::from_dims(&[n])),
    ];
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::None,
    )
    .expect("cast bf16→f32 dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for (i, (g, src)) in got.iter().zip(host_bf16.iter()).enumerate() {
        let expected = src.to_f32();
        assert!(
            (g - expected).abs() < 1e-3,
            "cast[{i}]: bf16={}, got f32={g}, expected {expected}",
            src.to_f32()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_cast_f32_to_f16() {
    use half::f16;

    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host_f32: Vec<f32> = vec![1.0, -2.5, 3.25, 0.0, 1e-3, 1e3];
    let n = host_f32.len();

    let in_storage = upload_f32(&backend, &host_f32);
    let out_bytes_h = backend.alloc_bytes_handle(n * 2).expect("alloc f16 out");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::F16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Cast, &[DType::F32, DType::F16], BackendId::Vulkan)[0]
            .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[n])),
        Layout::contiguous(Shape::from_dims(&[n])),
    ];
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::None,
    )
    .expect("cast f32→f16 dispatch");

    let raw = match &out_arc.read().unwrap().inner {
        BackendStorage::Vulkan(v) => backend.download_bytes(v).expect("d2h"),
        _ => panic!("not on Vulkan"),
    };
    let got: Vec<f16> = bytemuck::cast_slice::<u8, f16>(&raw).to_vec();
    for (i, (g, src)) in got.iter().zip(host_f32.iter()).enumerate() {
        let expected = f16::from_f32(*src);
        // f16 round-to-nearest-even; allow 1 ULP slack.
        let g_bits = g.to_bits() as i32;
        let e_bits = expected.to_bits() as i32;
        assert!(
            (g_bits - e_bits).abs() <= 1,
            "cast[{i}]: f32={src}, got f16 bits={g_bits:#x}, expected {e_bits:#x}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_cast_f16_to_f32() {
    use half::f16;

    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host_f16: Vec<f16> = vec![
        f16::from_f32(1.0),
        f16::from_f32(-2.5),
        f16::from_f32(3.25),
        f16::from_f32(0.0),
        f16::from_f32(1e-3),
        f16::from_f32(1e3),
    ];
    let n = host_f16.len();

    let f16_bytes: &[u8] = bytemuck::cast_slice(&host_f16);
    let in_vk = backend.upload_bytes_handle(f16_bytes).expect("upload f16");
    let in_storage = Storage::new(BackendStorage::Vulkan(in_vk), DType::F16);
    let out_bytes_h = backend.alloc_bytes_handle(n * 4).expect("alloc f32 out");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Cast, &[DType::F16, DType::F32], BackendId::Vulkan)[0]
            .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[n])),
        Layout::contiguous(Shape::from_dims(&[n])),
    ];
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::None,
    )
    .expect("cast f16→f32 dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for (i, (g, src)) in got.iter().zip(host_f16.iter()).enumerate() {
        let expected = src.to_f32();
        // f16 → f32 is exact (no rounding); allow tight tolerance.
        assert!(
            (g - expected).abs() < 1e-3,
            "cast[{i}]: f16={}, got f32={g}, expected {expected}",
            src.to_f32()
        );
    }
}

/// PowI y = x^3 across mixed-sign inputs.
#[test]
#[ignore]
fn vulkan_dispatch_powi_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host: Vec<f32> = vec![-2.0, -1.0, 0.0, 1.0, 2.0, 3.0];
    let n = host.len();

    let in_storage = upload_f32(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::PowIElementwise,
        &[DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    let layouts = vec![layout.clone(), layout];
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::PowI { exp: 3 },
    )
    .expect("powi dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    let expected: Vec<f32> = host.iter().map(|x| x.powi(3)).collect();
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert!((g - e).abs() < 1e-5, "powi at {i}: got {g}, expected {e}");
    }
}

/// PowI y = x^2 across the same inputs (hits the x*x fast path).
#[test]
#[ignore]
fn vulkan_dispatch_powi_squared_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host: Vec<f32> = vec![-2.5, -1.0, 0.0, 0.5, 1.5, 3.0];
    let n = host.len();

    let in_storage = upload_f32(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::PowIElementwise,
        &[DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    let layouts = vec![layout.clone(), layout];
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::PowI { exp: 2 },
    )
    .expect("powi dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    let expected: Vec<f32> = host.iter().map(|x| x * x).collect();
    assert_eq!(got, expected);
}

/// Clamp y = clamp(x, -1, 1) across 6 elements spanning the bounds.
#[test]
#[ignore]
fn vulkan_dispatch_clamp_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host: Vec<f32> = vec![-3.0, -1.5, -0.5, 0.5, 1.5, 3.0];
    let n = host.len();

    let in_storage = upload_f32(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::ClampElementwise,
        &[DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    let layouts = vec![layout.clone(), layout];
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Clamp {
            min: -1.0,
            max: 1.0,
        },
    )
    .expect("clamp dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    assert_eq!(got, vec![-1.0, -1.0, -0.5, 0.5, 1.0, 1.0]);
}

/// Affine y = 2*x + 3 across 6 elements.
#[test]
#[ignore]
fn vulkan_dispatch_affine_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
    let n = host.len();

    let in_storage = upload_f32(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Affine, &[DType::F32, DType::F32], BackendId::Vulkan)[0]
            .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    let layouts = vec![layout.clone(), layout];
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Affine { mul: 2.0, add: 3.0 },
    )
    .expect("affine dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    let expected: Vec<f32> = host.iter().map(|x| 2.0 * x + 3.0).collect();
    assert_eq!(got, expected);
}

/// Affine y = 2*x + 3, f64.
#[test]
#[ignore]
fn vulkan_dispatch_affine_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
    let n = host.len();

    let in_storage = upload_f64(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(n * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Affine, &[DType::F64, DType::F64], BackendId::Vulkan)[0]
            .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    let layouts = vec![layout.clone(), layout];
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Affine { mul: 2.0, add: 3.0 },
    )
    .expect("affine f64 dispatch");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    let expected: Vec<f64> = host.iter().map(|x| 2.0 * x + 3.0).collect();
    assert_eq!(got, expected);
}

/// Affine y = 2*x + 3, f16.
#[test]
#[ignore]
fn vulkan_dispatch_affine_f16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host_f32: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
    let host: Vec<half::f16> = host_f32.iter().map(|&x| half::f16::from_f32(x)).collect();
    let n = host.len();

    let in_storage = upload_f16(&backend, &host);
    // n*2 may not be u32-multiple; alloc rounded.
    let out_bytes = backend
        .alloc_bytes_handle(((n * 2 + 3) & !3) as usize)
        .expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Affine, &[DType::F16, DType::F16], BackendId::Vulkan)[0]
            .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    let layouts = vec![layout.clone(), layout];
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Affine { mul: 2.0, add: 3.0 },
    )
    .expect("affine f16 dispatch");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    for (i, (g, &x)) in got.iter().take(n).zip(host_f32.iter()).enumerate() {
        let expected = 2.0_f32 * x + 3.0;
        assert_eq!(
            g.to_f32(),
            expected,
            "affine f16[{i}]: got {}, expected {expected}",
            g.to_f32()
        );
    }
}

/// Affine y = 2*x + 3, bf16.
#[test]
#[ignore]
fn vulkan_dispatch_affine_bf16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host_f32: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
    let host: Vec<half::bf16> = host_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();
    let n = host.len();

    let in_storage = upload_bf16(&backend, &host);
    let out_bytes = backend
        .alloc_bytes_handle(((n * 2 + 3) & !3) as usize)
        .expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::Affine,
        &[DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    let layouts = vec![layout.clone(), layout];
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Affine { mul: 2.0, add: 3.0 },
    )
    .expect("affine bf16 dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    for (i, (g, &x)) in got.iter().take(n).zip(host_f32.iter()).enumerate() {
        let expected = 2.0_f32 * x + 3.0;
        assert_eq!(
            g.to_f32(),
            expected,
            "affine bf16[{i}]: got {}, expected {expected}",
            g.to_f32()
        );
    }
}

/// Helper for matmul: builds storages, dispatches, returns output.
fn run_matmul_f32(
    backend: &Arc<VulkanBackend>,
    lhs: &[f32],
    lhs_batch: &[usize],
    rhs: &[f32],
    rhs_batch: &[usize],
    m: usize,
    n: usize,
    k: usize,
) -> Vec<f32> {
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);
    let lhs_storage = upload_f32(backend, lhs);
    let rhs_storage = upload_f32(backend, rhs);
    let lhs_total: usize = lhs_batch.iter().product::<usize>().max(1);
    let out_n = lhs_total * m * n;
    let out_bytes = backend.alloc_bytes_handle(out_n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let lhs_arc = Arc::new(RwLock::new(lhs_storage));
    let rhs_arc = Arc::new(RwLock::new(rhs_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let kernel = table.lookup_alternatives(
        OpKind::MatMul,
        &[DType::F32, DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let mut lhs_dims = lhs_batch.to_vec();
    lhs_dims.extend_from_slice(&[m, k]);
    let mut rhs_dims = rhs_batch.to_vec();
    rhs_dims.extend_from_slice(&[k, n]);
    let mut out_dims = lhs_batch.to_vec();
    out_dims.extend_from_slice(&[m, n]);
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&lhs_dims)),
        Layout::contiguous(Shape::from_dims(&rhs_dims)),
        Layout::contiguous(Shape::from_dims(&out_dims)),
    ];
    kernel(
        &[Arc::clone(&lhs_arc), Arc::clone(&rhs_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Matmul {
            lhs_batch_dims: lhs_batch.to_vec(),
            rhs_batch_dims: rhs_batch.to_vec(),
            m,
            n,
            k,
            m_compute: MatmulM::All,
        },
    )
    .expect("matmul dispatch");
    download_f32(backend, &out_arc.read().unwrap())
}

/// Mixed-precision matmul f32 × bf16 → f32 (V.3.D). Small-M path.
#[test]
#[ignore]
fn vulkan_dispatch_matmul_f32_bf16_b_small_m() {
    use half::bf16;

    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // A = [[1,2,3], [4,5,6]] (2×3 f32); B = [[1,2], [3,4], [5,6]] (3×2 bf16).
    // A @ B = [[22, 28], [49, 64]] — but bf16 has limited precision so
    //         results are approximate.
    let a_f32: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
    let b_bf16: Vec<bf16> = vec![
        bf16::from_f32(1.0),
        bf16::from_f32(2.0),
        bf16::from_f32(3.0),
        bf16::from_f32(4.0),
        bf16::from_f32(5.0),
        bf16::from_f32(6.0),
    ];

    let a_bytes: &[u8] = bytemuck::cast_slice(&a_f32);
    let b_bytes: &[u8] = bytemuck::cast_slice(&b_bf16);
    let a_vk = backend.upload_bytes_handle(a_bytes).expect("a upload");
    let b_vk = backend.upload_bytes_handle(b_bytes).expect("b upload");
    let a_storage = Storage::new(BackendStorage::Vulkan(a_vk), DType::F32);
    let b_storage = Storage::new(BackendStorage::Vulkan(b_vk), DType::BF16);
    let out_n = 2 * 2;
    let out_bytes_h = backend.alloc_bytes_handle(out_n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::F32);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MatMul,
        &[DType::F32, DType::BF16, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[2, 3])),
        Layout::contiguous(Shape::from_dims(&[3, 2])),
        Layout::contiguous(Shape::from_dims(&[2, 2])),
    ];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Matmul {
            lhs_batch_dims: vec![],
            rhs_batch_dims: vec![],
            m: 2,
            n: 2,
            k: 3,
            m_compute: MatmulM::All,
        },
    )
    .expect("mixed-bf16 matmul dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    let expected = [22.0_f32, 28.0, 49.0, 64.0];
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        // bf16 has ~7-bit mantissa; small integer values are exact but
        // the multiply-accumulate can accumulate small drift. Tolerance
        // of 0.5 covers worst-case bf16 rounding for these magnitudes.
        assert!(
            (g - e).abs() < 0.5,
            "mixed-bf16 matmul [{i}]: got {g}, expected {e}"
        );
    }
}

/// Mixed-precision matmul f32 × bf16 → f32 with m=16, n=16, k=16 to
/// hit the cooperative-matrix (tensor-core) path when available.
/// Falls through to matmul_tiled_bf16_b otherwise — either way the
/// math must match.
#[test]
#[ignore]
fn vulkan_dispatch_matmul_f32_bf16_b_coop_size() {
    use half::bf16;

    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let m = 16usize;
    let n = 16usize;
    let k = 16usize;

    // A[i, j] = 1 for all (deterministic; simplifies expected).
    // B[i, j] = j (column index). Then (A @ B)[i, j] = sum_k(1 * j) = k * j.
    let a_f32: Vec<f32> = vec![1.0; m * k];
    let mut b_bf16: Vec<bf16> = Vec::with_capacity(k * n);
    for _i in 0..k {
        for j in 0..n {
            b_bf16.push(bf16::from_f32(j as f32));
        }
    }

    let a_bytes: &[u8] = bytemuck::cast_slice(&a_f32);
    let b_bytes: &[u8] = bytemuck::cast_slice(&b_bf16);
    let a_vk = backend.upload_bytes_handle(a_bytes).expect("a upload");
    let b_vk = backend.upload_bytes_handle(b_bytes).expect("b upload");
    let a_storage = Storage::new(BackendStorage::Vulkan(a_vk), DType::F32);
    let b_storage = Storage::new(BackendStorage::Vulkan(b_vk), DType::BF16);
    let out_n = m * n;
    let out_bytes_h = backend.alloc_bytes_handle(out_n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::F32);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MatMul,
        &[DType::F32, DType::BF16, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[m, k])),
        Layout::contiguous(Shape::from_dims(&[k, n])),
        Layout::contiguous(Shape::from_dims(&[m, n])),
    ];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Matmul {
            lhs_batch_dims: vec![],
            rhs_batch_dims: vec![],
            m,
            n,
            k,
            m_compute: MatmulM::All,
        },
    )
    .expect("mixed-bf16 coop-size matmul dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for i in 0..m {
        for j in 0..n {
            let expected = (k * j) as f32;
            let g = got[i * n + j];
            assert!(
                (g - expected).abs() < 0.5,
                "coop-size mixed-bf16 [{i}, {j}]: got {g}, expected {expected}"
            );
        }
    }
}

/// Pure-bf16 matmul bf16 × bf16 → f32 via cooperative-matrix tile.
/// 16×16×16 = smallest coop-eligible shape (one workgroup).
/// Reference: A is ones, B[:, j] = j → out[i, j] = K * j.
#[test]
#[ignore]
fn vulkan_dispatch_matmul_bf16_bf16_f32_coop_size() {
    use half::bf16;

    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let m = 16usize;
    let n = 16usize;
    let k = 16usize;

    let a_bf16: Vec<bf16> = vec![bf16::from_f32(1.0); m * k];
    let mut b_bf16: Vec<bf16> = Vec::with_capacity(k * n);
    for _i in 0..k {
        for j in 0..n {
            b_bf16.push(bf16::from_f32(j as f32));
        }
    }

    let a_bytes: &[u8] = bytemuck::cast_slice(&a_bf16);
    let b_bytes: &[u8] = bytemuck::cast_slice(&b_bf16);
    let a_vk = backend.upload_bytes_handle(a_bytes).expect("a upload");
    let b_vk = backend.upload_bytes_handle(b_bytes).expect("b upload");
    let a_storage = Storage::new(BackendStorage::Vulkan(a_vk), DType::BF16);
    let b_storage = Storage::new(BackendStorage::Vulkan(b_vk), DType::BF16);
    let out_n = m * n;
    let out_bytes_h = backend.alloc_bytes_handle(out_n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::F32);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MatMul,
        &[DType::BF16, DType::BF16, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[m, k])),
        Layout::contiguous(Shape::from_dims(&[k, n])),
        Layout::contiguous(Shape::from_dims(&[m, n])),
    ];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Matmul {
            lhs_batch_dims: vec![],
            rhs_batch_dims: vec![],
            m,
            n,
            k,
            m_compute: MatmulM::All,
        },
    )
    .expect("bf16 coop-size matmul dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for i in 0..m {
        for j in 0..n {
            let expected = (k * j) as f32;
            let g = got[i * n + j];
            assert!(
                (g - expected).abs() < 0.5,
                "bf16 coop-size [{i}, {j}]: got {g}, expected {expected}"
            );
        }
    }
}

/// Pure-bf16 matmul bf16 × bf16 → f32, larger shape (multiple workgroups).
/// 32×64×32 — exercises k_tiles=2, gx=1, gy=2.
#[test]
#[ignore]
fn vulkan_dispatch_matmul_bf16_bf16_f32_multi_tile() {
    use half::bf16;

    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let m = 32usize;
    let n = 64usize;
    let k = 32usize;

    // A[i, j] = 1; B[i, j] = j. Same pattern: out[i, j] = K * j.
    let a_bf16: Vec<bf16> = vec![bf16::from_f32(1.0); m * k];
    let mut b_bf16: Vec<bf16> = Vec::with_capacity(k * n);
    for _i in 0..k {
        for j in 0..n {
            b_bf16.push(bf16::from_f32(j as f32));
        }
    }

    let a_bytes: &[u8] = bytemuck::cast_slice(&a_bf16);
    let b_bytes: &[u8] = bytemuck::cast_slice(&b_bf16);
    let a_vk = backend.upload_bytes_handle(a_bytes).expect("a upload");
    let b_vk = backend.upload_bytes_handle(b_bytes).expect("b upload");
    let a_storage = Storage::new(BackendStorage::Vulkan(a_vk), DType::BF16);
    let b_storage = Storage::new(BackendStorage::Vulkan(b_vk), DType::BF16);
    let out_bytes_h = backend.alloc_bytes_handle(m * n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::F32);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MatMul,
        &[DType::BF16, DType::BF16, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[m, k])),
        Layout::contiguous(Shape::from_dims(&[k, n])),
        Layout::contiguous(Shape::from_dims(&[m, n])),
    ];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Matmul {
            lhs_batch_dims: vec![],
            rhs_batch_dims: vec![],
            m,
            n,
            k,
            m_compute: MatmulM::All,
        },
    )
    .expect("bf16 multi-tile matmul dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for i in 0..m {
        for j in 0..n {
            let expected = (k * j) as f32;
            let g = got[i * n + j];
            // Wider tolerance: K=32 means sum can drift more under bf16→f16 downcast.
            assert!(
                (g - expected).abs() < 2.0,
                "bf16 multi-tile [{i}, {j}]: got {g}, expected {expected}"
            );
        }
    }
}

/// Pure-bf16 matmul bf16 × bf16 → bf16 (downcast store).
/// Same shape as the f32-output test; expected matches but tolerance
/// reflects the bf16 truncation on the output store.
#[test]
#[ignore]
fn vulkan_dispatch_matmul_bf16_bf16_bf16_coop_size() {
    use half::bf16;

    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let m = 16usize;
    let n = 16usize;
    let k = 16usize;

    let a_bf16: Vec<bf16> = vec![bf16::from_f32(1.0); m * k];
    let mut b_bf16: Vec<bf16> = Vec::with_capacity(k * n);
    for _i in 0..k {
        for j in 0..n {
            b_bf16.push(bf16::from_f32(j as f32));
        }
    }

    let a_bytes: &[u8] = bytemuck::cast_slice(&a_bf16);
    let b_bytes: &[u8] = bytemuck::cast_slice(&b_bf16);
    let a_vk = backend.upload_bytes_handle(a_bytes).expect("a upload");
    let b_vk = backend.upload_bytes_handle(b_bytes).expect("b upload");
    let a_storage = Storage::new(BackendStorage::Vulkan(a_vk), DType::BF16);
    let b_storage = Storage::new(BackendStorage::Vulkan(b_vk), DType::BF16);
    let out_bytes_h = backend.alloc_bytes_handle(m * n * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::BF16);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MatMul,
        &[DType::BF16, DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[m, k])),
        Layout::contiguous(Shape::from_dims(&[k, n])),
        Layout::contiguous(Shape::from_dims(&[m, n])),
    ];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Matmul {
            lhs_batch_dims: vec![],
            rhs_batch_dims: vec![],
            m,
            n,
            k,
            m_compute: MatmulM::All,
        },
    )
    .expect("bf16→bf16 coop-size matmul dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    // Expected: out[i, j] = K * j. With K=16 and j ∈ [0, 16), all
    // outputs are integers in [0, 240]. bf16 represents integers
    // ≤ 256 exactly, so the post-downcast values round-trip clean.
    for i in 0..m {
        for j in 0..n {
            let expected = (k * j) as f32;
            let g = got[i * n + j].to_f32();
            assert_eq!(
                g, expected,
                "bf16→bf16 coop-size [{i}, {j}]: got {g}, expected {expected}"
            );
        }
    }
}

/// Pure-bf16 matmul bf16 × bf16 → bf16, larger shape (multi-tile).
#[test]
#[ignore]
fn vulkan_dispatch_matmul_bf16_bf16_bf16_multi_tile() {
    use half::bf16;

    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let m = 32usize;
    let n = 64usize;
    let k = 32usize;

    let a_bf16: Vec<bf16> = vec![bf16::from_f32(1.0); m * k];
    let mut b_bf16: Vec<bf16> = Vec::with_capacity(k * n);
    for _i in 0..k {
        for j in 0..n {
            b_bf16.push(bf16::from_f32(j as f32));
        }
    }

    let a_bytes: &[u8] = bytemuck::cast_slice(&a_bf16);
    let b_bytes: &[u8] = bytemuck::cast_slice(&b_bf16);
    let a_vk = backend.upload_bytes_handle(a_bytes).expect("a upload");
    let b_vk = backend.upload_bytes_handle(b_bytes).expect("b upload");
    let a_storage = Storage::new(BackendStorage::Vulkan(a_vk), DType::BF16);
    let b_storage = Storage::new(BackendStorage::Vulkan(b_vk), DType::BF16);
    let out_bytes_h = backend.alloc_bytes_handle(m * n * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::BF16);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MatMul,
        &[DType::BF16, DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[m, k])),
        Layout::contiguous(Shape::from_dims(&[k, n])),
        Layout::contiguous(Shape::from_dims(&[m, n])),
    ];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Matmul {
            lhs_batch_dims: vec![],
            rhs_batch_dims: vec![],
            m,
            n,
            k,
            m_compute: MatmulM::All,
        },
    )
    .expect("bf16→bf16 multi-tile matmul dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    // out[i, j] = K * j; K=32, j ∈ [0, 64), so max = 32 * 63 = 2016.
    // bf16's mantissa-7 can represent integer multiples of 8 above 1024
    // (when 2016 actually rounds to 2016 in bf16). Loose tolerance to
    // ±16 (one bf16 ULP near 2048) to be safe.
    for i in 0..m {
        for j in 0..n {
            let expected = (k * j) as f32;
            let g = got[i * n + j].to_f32();
            assert!(
                (g - expected).abs() <= 16.0,
                "bf16→bf16 multi-tile [{i}, {j}]: got {g}, expected {expected}"
            );
        }
    }
}

/// Pure-bf16 matmul bf16 × bf16 → f32, small-shape fallback path.
/// M=4, N=8, K=20 — fails the coop tile (M%16!=0 && N%16!=0) so the
/// wrapper routes to matmul_small_bf16_bf16_f32. Reference: matmul
/// is associative-respecting since K is small and inputs are exact
/// in bf16.
#[test]
#[ignore]
fn vulkan_dispatch_matmul_bf16_bf16_f32_small() {
    use half::bf16;

    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let m = 4usize;
    let n = 8usize;
    let k = 20usize;

    // A[i, j] = 1; B[i, j] = j. Out[i, j] = K * j.
    let a_bf16: Vec<bf16> = vec![bf16::from_f32(1.0); m * k];
    let mut b_bf16: Vec<bf16> = Vec::with_capacity(k * n);
    for _i in 0..k {
        for j in 0..n {
            b_bf16.push(bf16::from_f32(j as f32));
        }
    }

    let a_bytes: &[u8] = bytemuck::cast_slice(&a_bf16);
    let b_bytes: &[u8] = bytemuck::cast_slice(&b_bf16);
    let a_vk = backend.upload_bytes_handle(a_bytes).expect("a upload");
    let b_vk = backend.upload_bytes_handle(b_bytes).expect("b upload");
    let a_storage = Storage::new(BackendStorage::Vulkan(a_vk), DType::BF16);
    let b_storage = Storage::new(BackendStorage::Vulkan(b_vk), DType::BF16);
    let out_bytes_h = backend.alloc_bytes_handle(m * n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::F32);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MatMul,
        &[DType::BF16, DType::BF16, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[m, k])),
        Layout::contiguous(Shape::from_dims(&[k, n])),
        Layout::contiguous(Shape::from_dims(&[m, n])),
    ];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Matmul {
            lhs_batch_dims: vec![],
            rhs_batch_dims: vec![],
            m,
            n,
            k,
            m_compute: MatmulM::All,
        },
    )
    .expect("bf16 small-shape matmul dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for i in 0..m {
        for j in 0..n {
            let expected = (k * j) as f32;
            let g = got[i * n + j];
            // bf16 reduction; K=20, j ≤ 7 → max 140 (exact in bf16).
            assert!(
                (g - expected).abs() < 0.5,
                "bf16 small [{i}, {j}]: got {g}, expected {expected}"
            );
        }
    }
}

/// Pure-bf16 matmul bf16 × bf16 → bf16, small-shape fallback.
#[test]
#[ignore]
fn vulkan_dispatch_matmul_bf16_bf16_bf16_small() {
    use half::bf16;

    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let m = 4usize;
    let n = 8usize;
    let k = 20usize;

    let a_bf16: Vec<bf16> = vec![bf16::from_f32(1.0); m * k];
    let mut b_bf16: Vec<bf16> = Vec::with_capacity(k * n);
    for _i in 0..k {
        for j in 0..n {
            b_bf16.push(bf16::from_f32(j as f32));
        }
    }

    let a_bytes: &[u8] = bytemuck::cast_slice(&a_bf16);
    let b_bytes: &[u8] = bytemuck::cast_slice(&b_bf16);
    let a_vk = backend.upload_bytes_handle(a_bytes).expect("a upload");
    let b_vk = backend.upload_bytes_handle(b_bytes).expect("b upload");
    let a_storage = Storage::new(BackendStorage::Vulkan(a_vk), DType::BF16);
    let b_storage = Storage::new(BackendStorage::Vulkan(b_vk), DType::BF16);
    let out_bytes_h = backend
        .alloc_bytes_handle((m * n * 2 + 3) & !3)
        .expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::BF16);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MatMul,
        &[DType::BF16, DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[m, k])),
        Layout::contiguous(Shape::from_dims(&[k, n])),
        Layout::contiguous(Shape::from_dims(&[m, n])),
    ];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Matmul {
            lhs_batch_dims: vec![],
            rhs_batch_dims: vec![],
            m,
            n,
            k,
            m_compute: MatmulM::All,
        },
    )
    .expect("bf16→bf16 small-shape matmul dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    for i in 0..m {
        for j in 0..n {
            let expected = (k * j) as f32;
            let g = got[i * n + j].to_f32();
            assert_eq!(
                g, expected,
                "bf16→bf16 small [{i}, {j}]: got {g}, expected {expected}"
            );
        }
    }
}

/// Pure-f16 matmul f16 × f16 → f32 via cooperative-matrix tile.
/// 16×16×16 single-tile sanity check.
#[test]
#[ignore]
fn vulkan_dispatch_matmul_f16_f16_f32_coop_size() {
    use half::f16;

    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let m = 16usize;
    let n = 16usize;
    let k = 16usize;

    let a_f16: Vec<f16> = vec![f16::from_f32(1.0); m * k];
    let mut b_f16: Vec<f16> = Vec::with_capacity(k * n);
    for _i in 0..k {
        for j in 0..n {
            b_f16.push(f16::from_f32(j as f32));
        }
    }

    let a_bytes: &[u8] = bytemuck::cast_slice(&a_f16);
    let b_bytes: &[u8] = bytemuck::cast_slice(&b_f16);
    let a_vk = backend.upload_bytes_handle(a_bytes).expect("a upload");
    let b_vk = backend.upload_bytes_handle(b_bytes).expect("b upload");
    let a_storage = Storage::new(BackendStorage::Vulkan(a_vk), DType::F16);
    let b_storage = Storage::new(BackendStorage::Vulkan(b_vk), DType::F16);
    let out_bytes_h = backend.alloc_bytes_handle(m * n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::F32);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MatMul,
        &[DType::F16, DType::F16, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[m, k])),
        Layout::contiguous(Shape::from_dims(&[k, n])),
        Layout::contiguous(Shape::from_dims(&[m, n])),
    ];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Matmul {
            lhs_batch_dims: vec![],
            rhs_batch_dims: vec![],
            m,
            n,
            k,
            m_compute: MatmulM::All,
        },
    )
    .expect("f16 coop-size matmul dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for i in 0..m {
        for j in 0..n {
            let expected = (k * j) as f32;
            let g = got[i * n + j];
            // f16 has 10 mantissa bits — integers ≤ 1024 are exact.
            assert!(
                (g - expected).abs() < 0.5,
                "f16 coop-size [{i}, {j}]: got {g}, expected {expected}"
            );
        }
    }
}

/// Pure-f16 matmul f16 × f16 → f32, larger shape.
#[test]
#[ignore]
fn vulkan_dispatch_matmul_f16_f16_f32_multi_tile() {
    use half::f16;

    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let m = 32usize;
    let n = 64usize;
    let k = 32usize;

    let a_f16: Vec<f16> = vec![f16::from_f32(1.0); m * k];
    let mut b_f16: Vec<f16> = Vec::with_capacity(k * n);
    for _i in 0..k {
        for j in 0..n {
            b_f16.push(f16::from_f32(j as f32));
        }
    }

    let a_bytes: &[u8] = bytemuck::cast_slice(&a_f16);
    let b_bytes: &[u8] = bytemuck::cast_slice(&b_f16);
    let a_vk = backend.upload_bytes_handle(a_bytes).expect("a upload");
    let b_vk = backend.upload_bytes_handle(b_bytes).expect("b upload");
    let a_storage = Storage::new(BackendStorage::Vulkan(a_vk), DType::F16);
    let b_storage = Storage::new(BackendStorage::Vulkan(b_vk), DType::F16);
    let out_bytes_h = backend.alloc_bytes_handle(m * n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::F32);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MatMul,
        &[DType::F16, DType::F16, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[m, k])),
        Layout::contiguous(Shape::from_dims(&[k, n])),
        Layout::contiguous(Shape::from_dims(&[m, n])),
    ];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Matmul {
            lhs_batch_dims: vec![],
            rhs_batch_dims: vec![],
            m,
            n,
            k,
            m_compute: MatmulM::All,
        },
    )
    .expect("f16 multi-tile matmul dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for i in 0..m {
        for j in 0..n {
            let expected = (k * j) as f32;
            let g = got[i * n + j];
            assert!(
                (g - expected).abs() < 2.0,
                "f16 multi-tile [{i}, {j}]: got {g}, expected {expected}"
            );
        }
    }
}

/// Pure-f16 matmul f16 × f16 → f16 (downcast store).
#[test]
#[ignore]
fn vulkan_dispatch_matmul_f16_f16_f16_coop_size() {
    use half::f16;

    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let m = 16usize;
    let n = 16usize;
    let k = 16usize;

    let a_f16: Vec<f16> = vec![f16::from_f32(1.0); m * k];
    let mut b_f16: Vec<f16> = Vec::with_capacity(k * n);
    for _i in 0..k {
        for j in 0..n {
            b_f16.push(f16::from_f32(j as f32));
        }
    }

    let a_bytes: &[u8] = bytemuck::cast_slice(&a_f16);
    let b_bytes: &[u8] = bytemuck::cast_slice(&b_f16);
    let a_vk = backend.upload_bytes_handle(a_bytes).expect("a upload");
    let b_vk = backend.upload_bytes_handle(b_bytes).expect("b upload");
    let a_storage = Storage::new(BackendStorage::Vulkan(a_vk), DType::F16);
    let b_storage = Storage::new(BackendStorage::Vulkan(b_vk), DType::F16);
    let out_bytes_h = backend.alloc_bytes_handle(m * n * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::F16);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MatMul,
        &[DType::F16, DType::F16, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[m, k])),
        Layout::contiguous(Shape::from_dims(&[k, n])),
        Layout::contiguous(Shape::from_dims(&[m, n])),
    ];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Matmul {
            lhs_batch_dims: vec![],
            rhs_batch_dims: vec![],
            m,
            n,
            k,
            m_compute: MatmulM::All,
        },
    )
    .expect("f16→f16 coop-size matmul dispatch");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    // out[i, j] = K * j; max = 16 * 15 = 240. f16 represents integers
    // ≤ 1024 exactly, so values round-trip clean.
    for i in 0..m {
        for j in 0..n {
            let expected = (k * j) as f32;
            let g = got[i * n + j].to_f32();
            assert_eq!(
                g, expected,
                "f16→f16 coop-size [{i}, {j}]: got {g}, expected {expected}"
            );
        }
    }
}

/// Pure-f16 matmul f16 × f16 → f16, larger shape (multi-tile).
#[test]
#[ignore]
fn vulkan_dispatch_matmul_f16_f16_f16_multi_tile() {
    use half::f16;

    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let m = 32usize;
    let n = 64usize;
    let k = 32usize;

    let a_f16: Vec<f16> = vec![f16::from_f32(1.0); m * k];
    let mut b_f16: Vec<f16> = Vec::with_capacity(k * n);
    for _i in 0..k {
        for j in 0..n {
            b_f16.push(f16::from_f32(j as f32));
        }
    }

    let a_bytes: &[u8] = bytemuck::cast_slice(&a_f16);
    let b_bytes: &[u8] = bytemuck::cast_slice(&b_f16);
    let a_vk = backend.upload_bytes_handle(a_bytes).expect("a upload");
    let b_vk = backend.upload_bytes_handle(b_bytes).expect("b upload");
    let a_storage = Storage::new(BackendStorage::Vulkan(a_vk), DType::F16);
    let b_storage = Storage::new(BackendStorage::Vulkan(b_vk), DType::F16);
    let out_bytes_h = backend.alloc_bytes_handle(m * n * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::F16);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MatMul,
        &[DType::F16, DType::F16, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[m, k])),
        Layout::contiguous(Shape::from_dims(&[k, n])),
        Layout::contiguous(Shape::from_dims(&[m, n])),
    ];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Matmul {
            lhs_batch_dims: vec![],
            rhs_batch_dims: vec![],
            m,
            n,
            k,
            m_compute: MatmulM::All,
        },
    )
    .expect("f16→f16 multi-tile matmul dispatch");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    // out[i, j] = K * j; K=32, j ∈ [0, 64), so max = 32 * 63 = 2016.
    // f16 has 10-bit mantissa — integers ≤ 2048 exact, so all values
    // round-trip exactly.
    for i in 0..m {
        for j in 0..n {
            let expected = (k * j) as f32;
            let g = got[i * n + j].to_f32();
            assert_eq!(
                g, expected,
                "f16→f16 multi-tile [{i}, {j}]: got {g}, expected {expected}"
            );
        }
    }
}

/// Pure-f16 matmul f16 × f16 → f32, small-shape fallback path.
#[test]
#[ignore]
fn vulkan_dispatch_matmul_f16_f16_f32_small() {
    use half::f16;

    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let m = 4usize;
    let n = 8usize;
    let k = 20usize;

    let a_f16: Vec<f16> = vec![f16::from_f32(1.0); m * k];
    let mut b_f16: Vec<f16> = Vec::with_capacity(k * n);
    for _i in 0..k {
        for j in 0..n {
            b_f16.push(f16::from_f32(j as f32));
        }
    }

    let a_bytes: &[u8] = bytemuck::cast_slice(&a_f16);
    let b_bytes: &[u8] = bytemuck::cast_slice(&b_f16);
    let a_vk = backend.upload_bytes_handle(a_bytes).expect("a upload");
    let b_vk = backend.upload_bytes_handle(b_bytes).expect("b upload");
    let a_storage = Storage::new(BackendStorage::Vulkan(a_vk), DType::F16);
    let b_storage = Storage::new(BackendStorage::Vulkan(b_vk), DType::F16);
    let out_bytes_h = backend.alloc_bytes_handle(m * n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::F32);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MatMul,
        &[DType::F16, DType::F16, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[m, k])),
        Layout::contiguous(Shape::from_dims(&[k, n])),
        Layout::contiguous(Shape::from_dims(&[m, n])),
    ];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Matmul {
            lhs_batch_dims: vec![],
            rhs_batch_dims: vec![],
            m,
            n,
            k,
            m_compute: MatmulM::All,
        },
    )
    .expect("f16 small-shape matmul dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for i in 0..m {
        for j in 0..n {
            let expected = (k * j) as f32;
            let g = got[i * n + j];
            assert!(
                (g - expected).abs() < 0.5,
                "f16 small [{i}, {j}]: got {g}, expected {expected}"
            );
        }
    }
}

/// Pure-f16 matmul f16 × f16 → f16, small-shape fallback.
#[test]
#[ignore]
fn vulkan_dispatch_matmul_f16_f16_f16_small() {
    use half::f16;

    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let m = 4usize;
    let n = 8usize;
    let k = 20usize;

    let a_f16: Vec<f16> = vec![f16::from_f32(1.0); m * k];
    let mut b_f16: Vec<f16> = Vec::with_capacity(k * n);
    for _i in 0..k {
        for j in 0..n {
            b_f16.push(f16::from_f32(j as f32));
        }
    }

    let a_bytes: &[u8] = bytemuck::cast_slice(&a_f16);
    let b_bytes: &[u8] = bytemuck::cast_slice(&b_f16);
    let a_vk = backend.upload_bytes_handle(a_bytes).expect("a upload");
    let b_vk = backend.upload_bytes_handle(b_bytes).expect("b upload");
    let a_storage = Storage::new(BackendStorage::Vulkan(a_vk), DType::F16);
    let b_storage = Storage::new(BackendStorage::Vulkan(b_vk), DType::F16);
    let out_bytes_h = backend
        .alloc_bytes_handle((m * n * 2 + 3) & !3)
        .expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::F16);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MatMul,
        &[DType::F16, DType::F16, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[m, k])),
        Layout::contiguous(Shape::from_dims(&[k, n])),
        Layout::contiguous(Shape::from_dims(&[m, n])),
    ];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Matmul {
            lhs_batch_dims: vec![],
            rhs_batch_dims: vec![],
            m,
            n,
            k,
            m_compute: MatmulM::All,
        },
    )
    .expect("f16→f16 small-shape matmul dispatch");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    for i in 0..m {
        for j in 0..n {
            let expected = (k * j) as f32;
            let g = got[i * n + j].to_f32();
            assert_eq!(
                g, expected,
                "f16→f16 small [{i}, {j}]: got {g}, expected {expected}"
            );
        }
    }
}

/// Mixed-precision matmul f32 × bf16 → f32, m == 1 (matvec_bf16_b path).
#[test]
#[ignore]
fn vulkan_dispatch_matmul_f32_bf16_b_matvec() {
    use half::bf16;

    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let a_f32: Vec<f32> = vec![1.0, 2.0, 3.0]; // 1×3
    let b_bf16: Vec<bf16> = vec![
        bf16::from_f32(1.0),
        bf16::from_f32(2.0),
        bf16::from_f32(3.0),
        bf16::from_f32(4.0),
        bf16::from_f32(5.0),
        bf16::from_f32(6.0),
    ]; // 3×2

    let a_bytes: &[u8] = bytemuck::cast_slice(&a_f32);
    let b_bytes: &[u8] = bytemuck::cast_slice(&b_bf16);
    let a_vk = backend.upload_bytes_handle(a_bytes).expect("a upload");
    let b_vk = backend.upload_bytes_handle(b_bytes).expect("b upload");
    let a_storage = Storage::new(BackendStorage::Vulkan(a_vk), DType::F32);
    let b_storage = Storage::new(BackendStorage::Vulkan(b_vk), DType::BF16);
    // `1 * 2` is the output SHAPE [1, 2] written as its dimensions, not an
    // arithmetic identity. Folding it to `2` deletes the rank.
    #[allow(clippy::identity_op)]
    let out_n = 1 * 2;
    let out_bytes_h = backend.alloc_bytes_handle(out_n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::F32);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MatMul,
        &[DType::F32, DType::BF16, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[1, 3])),
        Layout::contiguous(Shape::from_dims(&[3, 2])),
        Layout::contiguous(Shape::from_dims(&[1, 2])),
    ];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Matmul {
            lhs_batch_dims: vec![],
            rhs_batch_dims: vec![],
            m: 1,
            n: 2,
            k: 3,
            m_compute: MatmulM::All,
        },
    )
    .expect("mixed-bf16 matvec dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    let expected = [22.0_f32, 28.0];
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert!(
            (g - e).abs() < 0.5,
            "mixed-bf16 matvec [{i}]: got {g}, expected {e}"
        );
    }
}

/// Matvec path: m == 1. 1×3 @ 3×2 → 1×2.
#[test]
#[ignore]
fn vulkan_dispatch_matmul_matvec_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    // A = [[1, 2, 3]] (1×3); B = [[1, 2], [3, 4], [5, 6]] (3×2).
    // A @ B = [[1+6+15, 2+8+18]] = [[22, 28]].
    let got = run_matmul_f32(
        &backend,
        &[1.0, 2.0, 3.0],
        &[],
        &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
        &[],
        1,
        2,
        3,
    );
    assert_eq!(got, vec![22.0, 28.0]);
}

/// Small-M path: 2 ≤ m < 32. 2×3 @ 3×2 → 2×2.
#[test]
#[ignore]
fn vulkan_dispatch_matmul_small_m_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    // A = [[1, 2, 3], [4, 5, 6]]; B = [[1, 2], [3, 4], [5, 6]].
    // Row 0: [22, 28]; Row 1: [4+15+30, 8+20+36] = [49, 64].
    let got = run_matmul_f32(
        &backend,
        &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
        &[],
        &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
        &[],
        2,
        2,
        3,
    );
    assert_eq!(got, vec![22.0, 28.0, 49.0, 64.0]);
}

/// Tiled-M path: m >= 32. Build a 32×4 @ 4×3 with deterministic values.
#[test]
#[ignore]
fn vulkan_dispatch_matmul_tiled_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let m = 32usize;
    let n = 3usize;
    let k = 4usize;
    // A[i, j] = i; B[i, j] = j+1. Row sum across k is i * sum(j+1 for j..k=4)
    //                                          = i * (1+1+1+1) for each col? No.
    // Actually A[i,j] = i. So row i is [i,i,i,i]. B[i,j] = j+1.
    // (A @ B)[i, j] = sum_k(A[i,k] * B[k,j]) = sum_k(i * (j+1)) = i * k * (j+1).
    // Wait B[k,j] = j+1 (depends only on j). So sum_k(i * (j+1)) = i * k * (j+1) = i * 4 * (j+1).
    let mut a = Vec::with_capacity(m * k);
    for i in 0..m {
        for _j in 0..k {
            a.push(i as f32);
        }
    }
    let mut b = Vec::with_capacity(k * n);
    for _i in 0..k {
        for j in 0..n {
            b.push((j + 1) as f32);
        }
    }
    let got = run_matmul_f32(&backend, &a, &[], &b, &[], m, n, k);
    for i in 0..m {
        for j in 0..n {
            let expected = (i * k * (j + 1)) as f32;
            let g = got[i * n + j];
            assert!(
                (g - expected).abs() < 1e-3,
                "tiled matmul [{i}, {j}]: got {g}, expected {expected}"
            );
        }
    }
}

/// Batched matmul: B=2 batch heads. lhs [2, 2, 3] @ rhs [2, 3, 2] → [2, 2, 2].
#[test]
#[ignore]
fn vulkan_dispatch_matmul_batched_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    // Batch 0: same as small_m_f32 → [22, 28, 49, 64]
    // Batch 1: A=[[1,1,1],[2,2,2]], B=[[1,1],[1,1],[1,1]] → [[3,3],[6,6]]
    let lhs: Vec<f32> = vec![
        1.0, 2.0, 3.0, 4.0, 5.0, 6.0, // batch 0
        1.0, 1.0, 1.0, 2.0, 2.0, 2.0, // batch 1
    ];
    let rhs: Vec<f32> = vec![
        1.0, 2.0, 3.0, 4.0, 5.0, 6.0, // batch 0
        1.0, 1.0, 1.0, 1.0, 1.0, 1.0, // batch 1
    ];
    let got = run_matmul_f32(&backend, &lhs, &[2], &rhs, &[2], 2, 2, 3);
    assert_eq!(
        got,
        vec![
            22.0, 28.0, 49.0, 64.0, // batch 0
            3.0, 3.0, 6.0, 6.0, // batch 1
        ]
    );
}

/// GQA: lhs has 2× the batch heads of rhs. lhs[4,1,3] @ rhs[2,3,2] → [4,1,2].
#[test]
#[ignore]
fn vulkan_dispatch_matmul_gqa_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    // 4 query heads share 2 key heads (GQA factor = 2).
    // Q[head_q] @ K[head_q / 2].
    let lhs: Vec<f32> = vec![
        1.0, 2.0, 3.0, // q0
        4.0, 5.0, 6.0, // q1
        1.0, 0.0, 0.0, // q2
        0.0, 1.0, 0.0, // q3
    ];
    let rhs: Vec<f32> = vec![
        // k0 (used by q0, q1): same B as small_m
        1.0, 2.0, 3.0, 4.0, 5.0, 6.0, // k1 (used by q2, q3): identity-like 3x2
        1.0, 0.0, 0.0, 1.0, 1.0, 1.0,
    ];
    let got = run_matmul_f32(&backend, &lhs, &[4], &rhs, &[2], 1, 2, 3);
    // q0 @ k0 = [22, 28] (same as 1x3 @ 3x2)
    // q1 @ k0 = [4*1+5*3+6*5, 4*2+5*4+6*6] = [49, 64]
    // q2 @ k1 = [1*1+0*0+0*1, 1*0+0*1+0*1] = [1, 0]
    // q3 @ k1 = [0*1+1*0+0*1, 0*0+1*1+0*1] = [0, 1]
    assert_eq!(got, vec![22.0, 28.0, 49.0, 64.0, 1.0, 0.0, 0.0, 1.0]);
}

/// Concat N=3 along last dim: chains via one intermediate allocation.
#[test]
#[ignore]
fn vulkan_dispatch_concat_n3_along_last_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // 3 inputs, each [2 rows]:
    //   a (2×2): [[1, 2], [3, 4]]
    //   b (2×3): [[10, 20, 30], [40, 50, 60]]
    //   c (2×1): [[100], [200]]
    // Concat along dim=1 → 2×6: [[1, 2, 10, 20, 30, 100],
    //                            [3, 4, 40, 50, 60, 200]]
    let a_data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0];
    let b_data: Vec<f32> = vec![10.0, 20.0, 30.0, 40.0, 50.0, 60.0];
    let c_data: Vec<f32> = vec![100.0, 200.0];

    let a_storage = upload_f32(&backend, &a_data);
    let b_storage = upload_f32(&backend, &b_data);
    let c_storage = upload_f32(&backend, &c_data);
    let out_bytes = backend.alloc_bytes_handle(12 * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let c_arc = Arc::new(RwLock::new(c_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Concat, &[DType::F32, DType::F32], BackendId::Vulkan)[0]
            .kernel;
    let layouts = vec![
        Layout::contiguous(Shape::from_dims(&[2, 2])),
        Layout::contiguous(Shape::from_dims(&[2, 3])),
        Layout::contiguous(Shape::from_dims(&[2, 1])),
        Layout::contiguous(Shape::from_dims(&[2, 6])),
    ];
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc), Arc::clone(&c_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Concat {
            outer_count: 2,
            input_dim_sizes: vec![2, 3, 1],
            inner_count: 1,
            axis: 1,
        },
    )
    .expect("N=3 concat dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    assert_eq!(
        got,
        vec![
            1.0, 2.0, 10.0, 20.0, 30.0, 100.0, 3.0, 4.0, 40.0, 50.0, 60.0, 200.0,
        ]
    );
}

/// Concat N=4 along the leading dim: tests the chain ping-pong.
#[test]
#[ignore]
fn vulkan_dispatch_concat_n4_along_leading_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // 4 inputs, each shape [?, 2]:
    //   a (1×2): [[1, 2]]
    //   b (2×2): [[3, 4], [5, 6]]
    //   c (1×2): [[7, 8]]
    //   d (1×2): [[9, 10]]
    // Concat dim=0 → 5×2.
    let a = vec![1.0_f32, 2.0];
    let b = vec![3.0_f32, 4.0, 5.0, 6.0];
    let c = vec![7.0_f32, 8.0];
    let d = vec![9.0_f32, 10.0];

    let arc_for = |v: &[f32], dim0: usize| {
        let st = upload_f32(&backend, v);
        (
            Arc::new(RwLock::new(st)),
            Layout::contiguous(Shape::from_dims(&[dim0, 2])),
        )
    };
    let (a_arc, a_l) = arc_for(&a, 1);
    let (b_arc, b_l) = arc_for(&b, 2);
    let (c_arc, c_l) = arc_for(&c, 1);
    let (d_arc, d_l) = arc_for(&d, 1);

    let out_n = 5 * 2;
    let out_bytes = backend.alloc_bytes_handle(out_n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Concat, &[DType::F32, DType::F32], BackendId::Vulkan)[0]
            .kernel;
    let layouts = vec![
        a_l,
        b_l,
        c_l,
        d_l,
        Layout::contiguous(Shape::from_dims(&[5, 2])),
    ];
    kernel(
        &[
            Arc::clone(&a_arc),
            Arc::clone(&b_arc),
            Arc::clone(&c_arc),
            Arc::clone(&d_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Concat {
            outer_count: 1,
            input_dim_sizes: vec![1, 2, 1, 1],
            inner_count: 2,
            axis: 0,
        },
    )
    .expect("N=4 concat dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    assert_eq!(
        got,
        vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0,]
    );
}

/// Concat binary along last dim: [[1,2,3], [4,5,6]] + [[7,8], [9,10]]
/// → [[1,2,3,7,8], [4,5,6,9,10]].
#[test]
#[ignore]
fn vulkan_dispatch_concat_along_last_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let a_data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; // 2×3
    let b_data: Vec<f32> = vec![7.0, 8.0, 9.0, 10.0]; // 2×2

    let a_storage = upload_f32(&backend, &a_data);
    let b_storage = upload_f32(&backend, &b_data);
    let out_n = 2 * 5;
    let out_bytes = backend.alloc_bytes_handle(out_n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);

    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Concat, &[DType::F32, DType::F32], BackendId::Vulkan)[0]
            .kernel;
    let a_layout = Layout::contiguous(Shape::from_dims(&[2, 3]));
    let b_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[2, 5]));
    let layouts = vec![a_layout, b_layout, out_layout];
    // concat dim = 1 (last). outer_count = prod(dims[..1]) = 2; inner_count = 1.
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Concat {
            outer_count: 2,
            input_dim_sizes: vec![3, 2],
            inner_count: 1,
            axis: 1,
        },
    )
    .expect("concat dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    assert_eq!(got, vec![1.0, 2.0, 3.0, 7.0, 8.0, 4.0, 5.0, 6.0, 9.0, 10.0]);
}

/// Helper for V.2.D reduce tests — uploads `data` of shape `dims`,
/// dispatches `op` with the supplied reduce dims + keepdim flag, and
/// returns the downloaded result. Output buffer is sized to fit the
/// largest legal output (full-reduce = 1 elem; last-dim = n_rows).
fn run_reduce_f32(
    backend: &Arc<VulkanBackend>,
    op: OpKind,
    data: &[f32],
    dims: &[usize],
    reduce_dims: Vec<usize>,
) -> Vec<f32> {
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);
    let in_storage = upload_f32(backend, data);
    // Output element count: 1 for full reduce; n_rows for last-dim
    // reduce (n_rows = product of dims[..rank-1]).
    let rank = dims.len();
    let out_elems = if reduce_dims.is_empty() || reduce_dims.len() == rank {
        1
    } else if reduce_dims.len() == 1 && reduce_dims[0] == rank - 1 {
        dims[..rank - 1].iter().product::<usize>().max(1)
    } else {
        panic!("test helper only supports full + last-dim reduce")
    };
    let out_bytes = backend.alloc_bytes_handle(out_elems * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let kernel =
        table.lookup_alternatives(op, &[DType::F32, DType::F32], BackendId::Vulkan)[0].kernel;
    let layout = Layout::contiguous(Shape::from_dims(dims));
    let out_layout = if reduce_dims.is_empty() || reduce_dims.len() == rank {
        Layout::contiguous(Shape::from_dims(&[1]))
    } else {
        Layout::contiguous(Shape::from_dims(&dims[..rank - 1]))
    };
    let layouts = vec![layout, out_layout];
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Reduce {
            dims: reduce_dims,
            keepdim: false,
        },
    )
    .expect("reduce dispatch");
    download_f32(backend, &out_arc.read().unwrap())
}

#[test]
#[ignore]
fn vulkan_dispatch_reduce_sum_full_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_reduce_f32(
        &backend,
        OpKind::SumReduce,
        &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
        &[2, 3],
        vec![0, 1],
    );
    assert_eq!(got.len(), 1);
    assert!((got[0] - 21.0).abs() < 1e-5, "sum got {got:?}");
}

#[test]
#[ignore]
fn vulkan_dispatch_reduce_max_full_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_reduce_f32(
        &backend,
        OpKind::MaxReduce,
        &[1.0, 7.0, 3.0, 4.0, 2.0, 6.0],
        &[2, 3],
        vec![0, 1],
    );
    assert_eq!(got, vec![7.0]);
}

#[test]
#[ignore]
fn vulkan_dispatch_reduce_min_full_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_reduce_f32(
        &backend,
        OpKind::MinReduce,
        &[5.0, 2.0, 3.0, -1.0, 4.0, 0.0],
        &[2, 3],
        vec![0, 1],
    );
    assert_eq!(got, vec![-1.0]);
}

#[test]
#[ignore]
fn vulkan_dispatch_reduce_mean_full_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_reduce_f32(
        &backend,
        OpKind::MeanReduce,
        &[2.0, 4.0, 6.0, 8.0],
        &[2, 2],
        vec![0, 1],
    );
    assert_eq!(got.len(), 1);
    assert!((got[0] - 5.0).abs() < 1e-5, "mean got {got:?}");
}

#[test]
#[ignore]
fn vulkan_dispatch_reduce_mean_last_dim_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    // 2 rows × 3 cols → per-row mean: [2, 5]
    let got = run_reduce_f32(
        &backend,
        OpKind::MeanReduce,
        &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
        &[2, 3],
        vec![1],
    );
    assert_eq!(got, vec![2.0, 5.0]);
}

#[test]
#[ignore]
fn vulkan_dispatch_reduce_sum_last_dim_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    // 2 rows × 3 cols → per-row sum: [6, 15]
    let got = run_reduce_f32(
        &backend,
        OpKind::SumReduce,
        &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
        &[2, 3],
        vec![1],
    );
    assert_eq!(got, vec![6.0, 15.0]);
}

/// IndexSelect along dim 0: pick rows 0, 2, 1 from a 4×3 matrix.
#[test]
#[ignore]
fn vulkan_dispatch_index_select_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // src shape [4, 3] — rows: [1..3], [4..6], [7..9], [10..12].
    let src_data: Vec<f32> = vec![
        1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0,
    ];
    let ids_data: Vec<u32> = vec![0, 2, 1];

    let src_storage = upload_f32(&backend, &src_data);
    let ids_bytes: &[u8] = bytemuck::cast_slice(&ids_data);
    let ids_vk = backend.upload_bytes_handle(ids_bytes).expect("ids upload");
    let ids_storage = Storage::new(BackendStorage::Vulkan(ids_vk), DType::U32);

    let outer_count = 1usize; // dims before axis 0
    let source_dim_size = 4usize; // src.dims[0]
    let n_indices = 3usize; // ids.len()
    let inner_count = 3usize; // dims after axis 0
    let out_n = outer_count * n_indices * inner_count;
    let out_bytes = backend.alloc_bytes_handle(out_n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);

    let src_arc = Arc::new(RwLock::new(src_storage));
    let ids_arc = Arc::new(RwLock::new(ids_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::IndexSelect,
        &[DType::F32, DType::U32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let src_layout = Layout::contiguous(Shape::from_dims(&[4, 3]));
    let ids_layout = Layout::contiguous(Shape::from_dims(&[3]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[3, 3]));
    let layouts = vec![src_layout, ids_layout, out_layout];
    kernel(
        &[Arc::clone(&src_arc), Arc::clone(&ids_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::IndexSelect {
            outer_count,
            source_dim_size,
            n_indices,
            inner_count,
        },
    )
    .expect("index_select dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    // Row 0 → [1,2,3]; Row 2 → [7,8,9]; Row 1 → [4,5,6].
    assert_eq!(got, vec![1.0, 2.0, 3.0, 7.0, 8.0, 9.0, 4.0, 5.0, 6.0]);
}

/// RoPE round-trip with `cos = [1,1,...]`, `sin = [0,0,...]` — should
/// be the identity (kernel emits `x` unchanged). Proves the 3-input
/// dispatch wiring + backend handle propagation work.
#[test]
#[ignore]
fn vulkan_dispatch_rope_identity_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // x has shape [outer=1, seq=2, head_dim=4]. head_dim must be even.
    let outer = 1usize;
    let seq = 2usize;
    let head_dim = 4usize;

    let x_data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
    // cos = 1, sin = 0 → identity rotation. Shape is [seq, head_dim]
    // (the kernel reads `cos[s, i]` for i in [0, head_dim), not just
    // the first half).
    let cos_data: Vec<f32> = vec![1.0; seq * head_dim];
    let sin_data: Vec<f32> = vec![0.0; seq * head_dim];

    let x_storage = upload_f32(&backend, &x_data);
    let cos_storage = upload_f32(&backend, &cos_data);
    let sin_storage = upload_f32(&backend, &sin_data);
    let n = outer * seq * head_dim;
    let out_bytes = backend.alloc_bytes_handle(n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);

    let x_arc = Arc::new(RwLock::new(x_storage));
    let cos_arc = Arc::new(RwLock::new(cos_storage));
    let sin_arc = Arc::new(RwLock::new(sin_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::Rope,
        &[DType::F32, DType::F32, DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let x_layout = Layout::contiguous(Shape::from_dims(&[outer, seq, head_dim]));
    let cos_layout = Layout::contiguous(Shape::from_dims(&[seq, head_dim]));
    let sin_layout = Layout::contiguous(Shape::from_dims(&[seq, head_dim]));
    let out_layout = x_layout.clone();
    let layouts = vec![x_layout, cos_layout, sin_layout, out_layout];
    kernel(
        &[
            Arc::clone(&x_arc),
            Arc::clone(&cos_arc),
            Arc::clone(&sin_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Rope {
            outer_count: outer,
            seq,
            head_dim,
        },
    )
    .expect("rope dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for (i, (g, x)) in got.iter().zip(x_data.iter()).enumerate() {
        assert!(
            (g - x).abs() < 1e-5,
            "rope identity mismatch at {i}: got {g}, expected {x}"
        );
    }
}

/// RoPE π/2 rotation: cos=0, sin=1. The kernel's rotate-half formula
/// then gives `out[i] = -x[i+h]` and `out[i+h] = x[i]`.
#[test]
#[ignore]
fn vulkan_dispatch_rope_quarter_rotation_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 1usize;
    let seq = 1usize;
    let head_dim = 4usize;
    let h = head_dim / 2;

    let x_data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0]; // (x0, x1, x2, x3); h=2 → pairs (1,3) and (2,4).
    let cos_data: Vec<f32> = vec![0.0; seq * head_dim];
    let sin_data: Vec<f32> = vec![1.0; seq * head_dim];

    let x_storage = upload_f32(&backend, &x_data);
    let cos_storage = upload_f32(&backend, &cos_data);
    let sin_storage = upload_f32(&backend, &sin_data);
    let n = outer * seq * head_dim;
    let out_bytes = backend.alloc_bytes_handle(n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);

    let x_arc = Arc::new(RwLock::new(x_storage));
    let cos_arc = Arc::new(RwLock::new(cos_storage));
    let sin_arc = Arc::new(RwLock::new(sin_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::Rope,
        &[DType::F32, DType::F32, DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let x_layout = Layout::contiguous(Shape::from_dims(&[outer, seq, head_dim]));
    let cs_layout = Layout::contiguous(Shape::from_dims(&[seq, head_dim]));
    let layouts = vec![x_layout.clone(), cs_layout.clone(), cs_layout, x_layout];
    kernel(
        &[
            Arc::clone(&x_arc),
            Arc::clone(&cos_arc),
            Arc::clone(&sin_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::Rope {
            outer_count: outer,
            seq,
            head_dim,
        },
    )
    .expect("rope dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    // Per shader formula with cos=0, sin=1:
    //   out[i]   = x[i] * 0 - x[i+h] * 1 = -x[i+h]
    //   out[i+h] = x[i+h] * 0 + x[i] * 1 = x[i]
    let _ = h;
    let expected: Vec<f32> = vec![-3.0, -4.0, 1.0, 2.0];
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert!(
            (g - e).abs() < 1e-5,
            "rope π/2 mismatch at {i}: got {g}, expected {e}"
        );
    }
}

/// RmsNorm row-wise: each row's RMS should be ~1.0 after normalization
/// (eps small relative to data). 2 rows × 4 cols.
#[test]
#[ignore]
fn vulkan_dispatch_rms_norm_last_dim_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let n = outer * last;
    let host: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 2.0, 4.0, 6.0, 8.0];

    let in_storage = upload_f32(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::RmsNormLastDim,
        &[DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    let layouts = vec![layout.clone(), layout];
    let eps = 1e-6f64;
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &layouts,
        &OpParams::NormLastDim {
            outer_count: outer,
            last_dim: last,
            eps,
        },
    )
    .expect("rmsnorm dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    // Reference RmsNorm: y = x / sqrt(mean(x^2) + eps).
    for row in 0..outer {
        let xs = &host[row * last..(row + 1) * last];
        let ys = &got[row * last..(row + 1) * last];
        let mean_sq: f32 = xs.iter().map(|x| x * x).sum::<f32>() / last as f32;
        let scale = (mean_sq + eps as f32).sqrt();
        for (i, (x, y)) in xs.iter().zip(ys.iter()).enumerate() {
            let expected = x / scale;
            assert!(
                (y - expected).abs() < 1e-4,
                "rmsnorm row {row} col {i}: got {y}, expected {expected}"
            );
        }
    }
}

// ---- LayerNormLastDimBackward f32/f16/bf16/f64 (V.3.G.layer_norm_bwd) ----

fn layer_norm_backward_ref(x: &[f32], g: &[f32], outer: usize, last: usize, eps: f32) -> Vec<f32> {
    let mut out = vec![0.0_f32; x.len()];
    let n = last as f32;
    for r in 0..outer {
        let off = r * last;
        let sum_x: f32 = x[off..off + last].iter().sum();
        let sum_x2: f32 = x[off..off + last].iter().map(|&v| v * v).sum();
        let sum_g: f32 = g[off..off + last].iter().sum();
        let sum_gx: f32 = x[off..off + last]
            .iter()
            .zip(g[off..off + last].iter())
            .map(|(&xi, &gi)| gi * xi)
            .sum();
        let mu = sum_x / n;
        let var = sum_x2 / n - mu * mu;
        let rstd = 1.0 / (var + eps).sqrt();
        let mean_g = sum_g / n;
        let mean_gxc = (sum_gx - sum_g * mu) / n;
        let rstd2 = rstd * rstd;
        for i in 0..last {
            let xi = x[off + i];
            let gi = g[off + i];
            let xc = xi - mu;
            out[off + i] = rstd * (gi - mean_g - xc * rstd2 * mean_gxc);
        }
    }
    out
}

#[test]
#[ignore]
fn vulkan_dispatch_layer_norm_backward_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let x: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 2.0, 4.0, 6.0, 8.0];
    let g: Vec<f32> = vec![1.0, -1.0, 2.0, -2.0, 0.5, 1.5, -0.5, -1.5];
    let eps = 1e-5_f64;

    let x_storage = upload_f32(&backend, &x);
    let g_storage = upload_f32(&backend, &g);
    let dx_bytes = backend.alloc_bytes_handle(outer * last * 4).expect("alloc");
    let dx_storage = Storage::new(BackendStorage::Vulkan(dx_bytes), DType::F32);
    let x_arc = Arc::new(RwLock::new(x_storage));
    let g_arc = Arc::new(RwLock::new(g_storage));
    let dx_arc = Arc::new(RwLock::new(dx_storage));

    let kernel = table.lookup_alternatives(
        OpKind::LayerNormLastDimBackward,
        &[DType::F32, DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    kernel(
        &[Arc::clone(&x_arc), Arc::clone(&g_arc)],
        &mut [Arc::clone(&dx_arc)],
        &[layout.clone(), layout.clone(), layout],
        &OpParams::NormLastDim {
            outer_count: outer,
            last_dim: last,
            eps,
        },
    )
    .expect("ln_bwd f32 dispatch");

    let got = download_f32(&backend, &dx_arc.read().unwrap());
    let expected = layer_norm_backward_ref(&x, &g, outer, last, eps as f32);
    for (i, (a, b)) in got.iter().zip(expected.iter()).enumerate() {
        assert!(
            (a - b).abs() < 1e-4,
            "ln_bwd f32[{i}]: got {a}, expected {b}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_layer_norm_backward_f16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let x_f32: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 2.0, 4.0, 6.0, 8.0];
    let g_f32: Vec<f32> = vec![1.0, -1.0, 2.0, -2.0, 0.5, 1.5, -0.5, -1.5];
    let x: Vec<half::f16> = x_f32.iter().map(|&v| half::f16::from_f32(v)).collect();
    let g: Vec<half::f16> = g_f32.iter().map(|&v| half::f16::from_f32(v)).collect();
    let eps = 1e-5_f64;

    let x_storage = upload_f16(&backend, &x);
    let g_storage = upload_f16(&backend, &g);
    let dx_bytes = backend.alloc_bytes_handle(outer * last * 2).expect("alloc");
    let dx_storage = Storage::new(BackendStorage::Vulkan(dx_bytes), DType::F16);
    let x_arc = Arc::new(RwLock::new(x_storage));
    let g_arc = Arc::new(RwLock::new(g_storage));
    let dx_arc = Arc::new(RwLock::new(dx_storage));

    let kernel = table.lookup_alternatives(
        OpKind::LayerNormLastDimBackward,
        &[DType::F16, DType::F16, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    kernel(
        &[Arc::clone(&x_arc), Arc::clone(&g_arc)],
        &mut [Arc::clone(&dx_arc)],
        &[layout.clone(), layout.clone(), layout],
        &OpParams::NormLastDim {
            outer_count: outer,
            last_dim: last,
            eps,
        },
    )
    .expect("ln_bwd f16 dispatch");

    let got = download_f16(&backend, &dx_arc.read().unwrap());
    let expected = layer_norm_backward_ref(&x_f32, &g_f32, outer, last, eps as f32);
    for (i, (a, b)) in got.iter().zip(expected.iter()).enumerate() {
        let af = a.to_f32();
        assert!(
            (af - b).abs() < 5e-3,
            "ln_bwd f16[{i}]: got {af}, expected {b}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_layer_norm_backward_bf16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize; // even
    let x_f32: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 2.0, 4.0, 6.0, 8.0];
    let g_f32: Vec<f32> = vec![1.0, -1.0, 2.0, -2.0, 0.5, 1.5, -0.5, -1.5];
    let x: Vec<half::bf16> = x_f32.iter().map(|&v| half::bf16::from_f32(v)).collect();
    let g: Vec<half::bf16> = g_f32.iter().map(|&v| half::bf16::from_f32(v)).collect();
    let eps = 1e-5_f64;

    let x_storage = upload_bf16(&backend, &x);
    let g_storage = upload_bf16(&backend, &g);
    let dx_bytes = backend.alloc_bytes_handle(outer * last * 2).expect("alloc");
    let dx_storage = Storage::new(BackendStorage::Vulkan(dx_bytes), DType::BF16);
    let x_arc = Arc::new(RwLock::new(x_storage));
    let g_arc = Arc::new(RwLock::new(g_storage));
    let dx_arc = Arc::new(RwLock::new(dx_storage));

    let kernel = table.lookup_alternatives(
        OpKind::LayerNormLastDimBackward,
        &[DType::BF16, DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    kernel(
        &[Arc::clone(&x_arc), Arc::clone(&g_arc)],
        &mut [Arc::clone(&dx_arc)],
        &[layout.clone(), layout.clone(), layout],
        &OpParams::NormLastDim {
            outer_count: outer,
            last_dim: last,
            eps,
        },
    )
    .expect("ln_bwd bf16 dispatch");

    let got = download_bf16(&backend, &dx_arc.read().unwrap());
    let expected = layer_norm_backward_ref(&x_f32, &g_f32, outer, last, eps as f32);
    for (i, (a, b)) in got.iter().zip(expected.iter()).enumerate() {
        let af = a.to_f32();
        assert!(
            (af - b).abs() < 5e-2,
            "ln_bwd bf16[{i}]: got {af}, expected {b}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_layer_norm_backward_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let x: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0, 2.0, 4.0, 6.0, 8.0];
    let g: Vec<f64> = vec![1.0, -1.0, 2.0, -2.0, 0.5, 1.5, -0.5, -1.5];
    let eps = 1e-12_f64;

    let x_storage = upload_f64(&backend, &x);
    let g_storage = upload_f64(&backend, &g);
    let dx_bytes = backend.alloc_bytes_handle(outer * last * 8).expect("alloc");
    let dx_storage = Storage::new(BackendStorage::Vulkan(dx_bytes), DType::F64);
    let x_arc = Arc::new(RwLock::new(x_storage));
    let g_arc = Arc::new(RwLock::new(g_storage));
    let dx_arc = Arc::new(RwLock::new(dx_storage));

    let kernel = table.lookup_alternatives(
        OpKind::LayerNormLastDimBackward,
        &[DType::F64, DType::F64, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    kernel(
        &[Arc::clone(&x_arc), Arc::clone(&g_arc)],
        &mut [Arc::clone(&dx_arc)],
        &[layout.clone(), layout.clone(), layout],
        &OpParams::NormLastDim {
            outer_count: outer,
            last_dim: last,
            eps,
        },
    )
    .expect("ln_bwd f64 dispatch");

    let got = download_f64(&backend, &dx_arc.read().unwrap());
    // Pure arithmetic; tight f64 tolerance.
    let n = last as f64;
    for r in 0..outer {
        let off = r * last;
        let sum_x: f64 = x[off..off + last].iter().sum();
        let sum_x2: f64 = x[off..off + last].iter().map(|&v| v * v).sum();
        let sum_g: f64 = g[off..off + last].iter().sum();
        let sum_gx: f64 = x[off..off + last]
            .iter()
            .zip(g[off..off + last].iter())
            .map(|(&xi, &gi)| gi * xi)
            .sum();
        let mu = sum_x / n;
        let var = sum_x2 / n - mu * mu;
        let rstd = 1.0 / (var + eps).sqrt();
        let mean_g = sum_g / n;
        let mean_gxc = (sum_gx - sum_g * mu) / n;
        let rstd2 = rstd * rstd;
        for i in 0..last {
            let xi = x[off + i];
            let gi = g[off + i];
            let xc = xi - mu;
            let expected = rstd * (gi - mean_g - xc * rstd2 * mean_gxc);
            assert!(
                (got[off + i] - expected).abs() < 1e-10,
                "ln_bwd f64[{}][{}]: got {}, expected {expected}",
                r,
                i,
                got[off + i]
            );
        }
    }
}

// ---- LayerNormLastDim f32/f16/bf16/f64 (V.3.G.layer_norm, 2026-05-30) ----

fn layer_norm_ref(x: &[f32], outer: usize, last: usize, eps: f32) -> Vec<f32> {
    let mut out = vec![0.0_f32; x.len()];
    let inv_n = 1.0 / last as f32;
    for r in 0..outer {
        let off = r * last;
        let mean = x[off..off + last].iter().sum::<f32>() * inv_n;
        let var = x[off..off + last]
            .iter()
            .map(|&v| (v - mean).powi(2))
            .sum::<f32>()
            * inv_n;
        let inv_std = 1.0 / (var + eps).sqrt();
        for i in 0..last {
            out[off + i] = (x[off + i] - mean) * inv_std;
        }
    }
    out
}

#[test]
#[ignore]
fn vulkan_dispatch_layer_norm_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let host: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 2.0, 4.0, 6.0, 8.0];
    let eps = 1e-5_f64;

    let in_storage = upload_f32(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(outer * last * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::LayerNormLastDim,
        &[DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::NormLastDim {
            outer_count: outer,
            last_dim: last,
            eps,
        },
    )
    .expect("layer_norm f32 dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    let expected = layer_norm_ref(&host, outer, last, eps as f32);
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert!(
            (g - e).abs() < 1e-5,
            "layer_norm f32[{i}]: got {g}, expected {e}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_layer_norm_f16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let host_f32: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 2.0, 4.0, 6.0, 8.0];
    let host: Vec<half::f16> = host_f32.iter().map(|&x| half::f16::from_f32(x)).collect();
    let eps = 1e-5_f64;

    let in_storage = upload_f16(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(outer * last * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::LayerNormLastDim,
        &[DType::F16, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::NormLastDim {
            outer_count: outer,
            last_dim: last,
            eps,
        },
    )
    .expect("layer_norm f16 dispatch");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    let expected = layer_norm_ref(&host_f32, outer, last, eps as f32);
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        let got_f32 = g.to_f32();
        assert!(
            (got_f32 - e).abs() < 5e-3,
            "layer_norm f16[{i}]: got {got_f32}, expected {e}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_layer_norm_bf16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let host_f32: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 2.0, 4.0, 6.0, 8.0];
    let host: Vec<half::bf16> = host_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();
    let eps = 1e-5_f64;

    let in_storage = upload_bf16(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(outer * last * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::LayerNormLastDim,
        &[DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::NormLastDim {
            outer_count: outer,
            last_dim: last,
            eps,
        },
    )
    .expect("layer_norm bf16 dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    let expected = layer_norm_ref(&host_f32, outer, last, eps as f32);
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        let got_f32 = g.to_f32();
        assert!(
            (got_f32 - e).abs() < 5e-2,
            "layer_norm bf16[{i}]: got {got_f32}, expected {e}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_layer_norm_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let host: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0, 2.0, 4.0, 6.0, 8.0];
    let eps = 1e-12_f64;

    let in_storage = upload_f64(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(outer * last * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::LayerNormLastDim,
        &[DType::F64, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::NormLastDim {
            outer_count: outer,
            last_dim: last,
            eps,
        },
    )
    .expect("layer_norm f64 dispatch");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    // Pure arithmetic (sqrt is precise on f64) — tight tolerance.
    let inv_n = 1.0 / last as f64;
    for r in 0..outer {
        let off = r * last;
        let mean = host[off..off + last].iter().sum::<f64>() * inv_n;
        let var = host[off..off + last]
            .iter()
            .map(|&v| (v - mean).powi(2))
            .sum::<f64>()
            * inv_n;
        let inv_std = 1.0 / (var + eps).sqrt();
        for i in 0..last {
            let expected = (host[off + i] - mean) * inv_std;
            assert!(
                (got[off + i] - expected).abs() < 1e-10,
                "layer_norm f64[{}][{}]: got {}, expected {expected}",
                r,
                i,
                got[off + i]
            );
        }
    }
}

// ---- Gather (V.3.G.gather, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_gather_f32_dim1() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // src = [[1,2,3],[10,20,30]]
    // indices shape = [2,4]: pick columns per row
    // Expected: out[r,c] = src[r, indices[r,c]]
    let src = [1.0_f32, 2.0, 3.0, 10.0, 20.0, 30.0];
    let indices: Vec<u32> = vec![
        2, 0, 1, 0, // row 0 picks: 3, 1, 2, 1
        1, 1, 2, 0, // row 1 picks: 20, 20, 30, 10
    ];
    let expected = [3.0_f32, 1.0, 2.0, 1.0, 20.0, 20.0, 30.0, 10.0];

    let src_storage = upload_f32(&backend, &src);
    let idx_bytes: &[u8] = bytemuck::cast_slice(&indices);
    let idx_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(idx_bytes).expect("idx upload")),
        DType::U32,
    );
    let out_bytes = backend.alloc_bytes_handle(8 * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let src_arc = Arc::new(RwLock::new(src_storage));
    let idx_arc = Arc::new(RwLock::new(idx_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::Gather,
        &[DType::F32, DType::U32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let src_layout = Layout::contiguous(Shape::from_dims(&[2, 3]));
    let idx_layout = Layout::contiguous(Shape::from_dims(&[2, 4]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[2, 4]));
    kernel(
        &[Arc::clone(&src_arc), Arc::clone(&idx_arc)],
        &mut [Arc::clone(&out_arc)],
        &[src_layout, idx_layout, out_layout],
        &OpParams::Gather {
            source_shape: vec![2, 3],
            output_shape: vec![2, 4],
            dim: 1,
        },
    )
    .expect("gather f32 dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "gather f32[{i}]: got {g}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_gather_f64_dim0() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // src shape [3, 2], gather along dim 0.
    // Each output position picks a row index for that column.
    let src = [1.0_f64, 2.0, 10.0, 20.0, 100.0, 200.0];
    let indices: Vec<u32> = vec![
        2, 0, // out[0] = (src[2,0], src[0,1])
        1, 2, // out[1] = (src[1,0], src[2,1])
    ];
    let expected = [100.0_f64, 2.0, 10.0, 200.0];

    let src_storage = upload_f64(&backend, &src);
    let idx_bytes: &[u8] = bytemuck::cast_slice(&indices);
    let idx_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(idx_bytes).expect("idx upload")),
        DType::U32,
    );
    let out_bytes = backend.alloc_bytes_handle(4 * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let src_arc = Arc::new(RwLock::new(src_storage));
    let idx_arc = Arc::new(RwLock::new(idx_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::Gather,
        &[DType::F64, DType::U32, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let src_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    let idx_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    kernel(
        &[Arc::clone(&src_arc), Arc::clone(&idx_arc)],
        &mut [Arc::clone(&out_arc)],
        &[src_layout, idx_layout, out_layout],
        &OpParams::Gather {
            source_shape: vec![3, 2],
            output_shape: vec![2, 2],
            dim: 0,
        },
    )
    .expect("gather f64 dispatch");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "gather f64[{i}]: got {g}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_gather_bf16_dim1() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // n_out = 8 (even) — pair-thread path.
    let src_f32 = [1.0_f32, 2.0, 3.0, 4.0, 10.0, 20.0, 30.0, 40.0];
    let src: Vec<half::bf16> = src_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();
    let indices: Vec<u32> = vec![3, 0, 1, 2, 2, 1, 0, 3];
    let expected_f32 = [4.0_f32, 1.0, 2.0, 3.0, 30.0, 20.0, 10.0, 40.0];

    let src_storage = upload_bf16(&backend, &src);
    let idx_bytes: &[u8] = bytemuck::cast_slice(&indices);
    let idx_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(idx_bytes).expect("idx upload")),
        DType::U32,
    );
    let out_bytes = backend.alloc_bytes_handle(8 * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let src_arc = Arc::new(RwLock::new(src_storage));
    let idx_arc = Arc::new(RwLock::new(idx_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::Gather,
        &[DType::BF16, DType::U32, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let src_layout = Layout::contiguous(Shape::from_dims(&[2, 4]));
    let idx_layout = Layout::contiguous(Shape::from_dims(&[2, 4]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[2, 4]));
    kernel(
        &[Arc::clone(&src_arc), Arc::clone(&idx_arc)],
        &mut [Arc::clone(&out_arc)],
        &[src_layout, idx_layout, out_layout],
        &OpParams::Gather {
            source_shape: vec![2, 4],
            output_shape: vec![2, 4],
            dim: 1,
        },
    )
    .expect("gather bf16 dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected_f32.iter()).enumerate() {
        assert_eq!(
            g.to_f32(),
            *e,
            "gather bf16[{i}]: got {}, expected {e}",
            g.to_f32()
        );
    }
}

// ---- MaskedFill (V.3.G.masked_fill, 2026-05-30) ----
//
// GAP-168(c): the mask is `Bool`, not `U8`. These three tests kept a `U8` mask
// storage dtype and a `[T, U8, T]` lookup key after that ruling, so
// `lookup_alternatives` returned nothing and `[0]` panicked "index out of
// bounds: the len is 0" — which reads as a missing kernel registration, and is
// not one: the kernel is registered under `(MaskedFill, [T, BOOL, T], Vulkan)`
// (see `docs/kernel-contracts/vulkan/select.fkc.md`). The mask BYTES (0/1) are
// unchanged; `Bool` is one byte per element.

#[test]
#[ignore]
fn vulkan_dispatch_masked_fill_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let input = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
    let mask: [u8; 8] = [0, 1, 0, 0, 1, 1, 0, 1]; // fill positions 1, 4, 5, 7
    let fill: f32 = -42.0;

    let in_storage = upload_f32(&backend, &input);
    let mask_bytes_u: &[u8] = &mask;
    let mask_storage = Storage::new(
        BackendStorage::Vulkan(
            backend
                .upload_bytes_handle(mask_bytes_u)
                .expect("mask upload"),
        ),
        DType::Bool,
    );
    let out_bytes = backend.alloc_bytes_handle(8 * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let mask_arc = Arc::new(RwLock::new(mask_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MaskedFill,
        &[DType::F32, DType::Bool, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[8]));
    kernel(
        &[Arc::clone(&in_arc), Arc::clone(&mask_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout.clone(), layout],
        &OpParams::MaskedFill {
            fill_bytes: fill.to_le_bytes().to_vec(),
        },
    )
    .expect("masked_fill f32 dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    let expected = [1.0_f32, fill, 3.0, 4.0, fill, fill, 7.0, fill];
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "masked_fill f32[{i}]: got {g}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_masked_fill_bf16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let input_f32 = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
    let input: Vec<half::bf16> = input_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();
    let mask: [u8; 8] = [0, 1, 0, 0, 1, 1, 0, 1];
    let fill = half::bf16::from_f32(-1.0);

    let in_storage = upload_bf16(&backend, &input);
    let mask_bytes_u: &[u8] = &mask;
    let mask_storage = Storage::new(
        BackendStorage::Vulkan(
            backend
                .upload_bytes_handle(mask_bytes_u)
                .expect("mask upload"),
        ),
        DType::Bool,
    );
    let out_bytes = backend.alloc_bytes_handle(8 * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let mask_arc = Arc::new(RwLock::new(mask_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MaskedFill,
        &[DType::BF16, DType::Bool, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[8]));
    kernel(
        &[Arc::clone(&in_arc), Arc::clone(&mask_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout.clone(), layout],
        &OpParams::MaskedFill {
            fill_bytes: fill.to_le_bytes().to_vec(),
        },
    )
    .expect("masked_fill bf16 dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    let expected_f32 = [1.0_f32, -1.0, 3.0, 4.0, -1.0, -1.0, 7.0, -1.0];
    for (i, (g, e)) in got.iter().zip(expected_f32.iter()).enumerate() {
        assert_eq!(
            g.to_f32(),
            *e,
            "masked_fill bf16[{i}]: got {}, expected {e}",
            g.to_f32()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_masked_fill_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let input = [1.0_f64, 2.0, 3.0, 4.0];
    let mask: [u8; 4] = [1, 0, 1, 0];
    let fill = -99.5_f64;

    let in_storage = upload_f64(&backend, &input);
    let mask_bytes_u: &[u8] = &mask;
    let mask_storage = Storage::new(
        BackendStorage::Vulkan(
            backend
                .upload_bytes_handle(mask_bytes_u)
                .expect("mask upload"),
        ),
        DType::Bool,
    );
    let out_bytes = backend.alloc_bytes_handle(4 * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let mask_arc = Arc::new(RwLock::new(mask_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MaskedFill,
        &[DType::F64, DType::Bool, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[4]));
    kernel(
        &[Arc::clone(&in_arc), Arc::clone(&mask_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout.clone(), layout],
        &OpParams::MaskedFill {
            fill_bytes: fill.to_le_bytes().to_vec(),
        },
    )
    .expect("masked_fill f64 dispatch");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    let expected = [fill, 2.0, fill, 4.0];
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "masked_fill f64[{i}]: got {g}, expected {e}");
    }
}

// ---- IndexAdd (V.3.G.index_add, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_index_add_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // base shape [3, 2] zeros; indices [0, 2, 0]; src shape [3, 2]:
    //   [[1, 2], [3, 4], [5, 6]]
    // For k=0: indices[0]=0 → out[0,:] += src[0,:] = [1,2]
    // For k=1: indices[1]=2 → out[2,:] += src[1,:] = [3,4]
    // For k=2: indices[2]=0 → out[0,:] += src[2,:] = [5,6]
    // expected: [[6, 8], [0, 0], [3, 4]]
    let base = [0.0_f32; 6];
    let indices: Vec<u32> = vec![0, 2, 0];
    let src = [1.0_f32, 2., 3., 4., 5., 6.];
    let expected = [6.0_f32, 8., 0., 0., 3., 4.];

    let base_storage = upload_f32(&backend, &base);
    let idx_bytes: &[u8] = bytemuck::cast_slice(&indices);
    let idx_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(idx_bytes).expect("idx upload")),
        DType::U32,
    );
    let src_storage = upload_f32(&backend, &src);
    let out_bytes = backend.alloc_bytes_handle(6 * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let base_arc = Arc::new(RwLock::new(base_storage));
    let idx_arc = Arc::new(RwLock::new(idx_storage));
    let src_arc = Arc::new(RwLock::new(src_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::IndexAdd,
        &[DType::F32, DType::U32, DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let base_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    let idx_layout = Layout::contiguous(Shape::from_dims(&[3]));
    let src_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    kernel(
        &[
            Arc::clone(&base_arc),
            Arc::clone(&idx_arc),
            Arc::clone(&src_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &[base_layout, idx_layout, src_layout, out_layout],
        &OpParams::IndexAdd {
            outer_count: 1,
            base_dim_size: 3,
            n_indices: 3,
            inner_count: 2,
        },
    )
    .expect("index_add f32 dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "index_add f32[{i}]: got {g}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_index_add_f32_starts_from_base() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let base = [10.0_f32, 20., 30., 40.];
    let indices: Vec<u32> = vec![0];
    let src = [100.0_f32, 200.];
    // outer=1, base_dim=2, n_indices=1, inner=2
    // For k=0: indices[0]=0 → out[0,:] += src[0,:] = [100, 200]
    // out[0,0] = 10 + 100 = 110
    // out[0,1] = 20 + 200 = 220
    // out[1,:] unchanged = [30, 40]
    let expected = [110.0_f32, 220., 30., 40.];

    let base_storage = upload_f32(&backend, &base);
    let idx_bytes: &[u8] = bytemuck::cast_slice(&indices);
    let idx_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(idx_bytes).expect("idx upload")),
        DType::U32,
    );
    let src_storage = upload_f32(&backend, &src);
    let out_bytes = backend.alloc_bytes_handle(4 * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let base_arc = Arc::new(RwLock::new(base_storage));
    let idx_arc = Arc::new(RwLock::new(idx_storage));
    let src_arc = Arc::new(RwLock::new(src_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::IndexAdd,
        &[DType::F32, DType::U32, DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let base_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    let idx_layout = Layout::contiguous(Shape::from_dims(&[1]));
    let src_layout = Layout::contiguous(Shape::from_dims(&[1, 2]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    kernel(
        &[
            Arc::clone(&base_arc),
            Arc::clone(&idx_arc),
            Arc::clone(&src_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &[base_layout, idx_layout, src_layout, out_layout],
        &OpParams::IndexAdd {
            outer_count: 1,
            base_dim_size: 2,
            n_indices: 1,
            inner_count: 2,
        },
    )
    .expect("index_add f32 base-init dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(
            *g, *e,
            "index_add f32 base-init[{i}]: got {g}, expected {e}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_index_add_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let base = [0.0_f64; 6];
    let indices: Vec<u32> = vec![0, 2, 0];
    let src = [1.0_f64, 2., 3., 4., 5., 6.];
    let expected = [6.0_f64, 8., 0., 0., 3., 4.];

    let base_storage = upload_f64(&backend, &base);
    let idx_bytes: &[u8] = bytemuck::cast_slice(&indices);
    let idx_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(idx_bytes).expect("idx upload")),
        DType::U32,
    );
    let src_storage = upload_f64(&backend, &src);
    let out_bytes = backend.alloc_bytes_handle(6 * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let base_arc = Arc::new(RwLock::new(base_storage));
    let idx_arc = Arc::new(RwLock::new(idx_storage));
    let src_arc = Arc::new(RwLock::new(src_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::IndexAdd,
        &[DType::F64, DType::U32, DType::F64, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let base_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    let idx_layout = Layout::contiguous(Shape::from_dims(&[3]));
    let src_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    kernel(
        &[
            Arc::clone(&base_arc),
            Arc::clone(&idx_arc),
            Arc::clone(&src_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &[base_layout, idx_layout, src_layout, out_layout],
        &OpParams::IndexAdd {
            outer_count: 1,
            base_dim_size: 3,
            n_indices: 3,
            inner_count: 2,
        },
    )
    .expect("index_add f64 dispatch");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "index_add f64[{i}]: got {g}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_index_add_bf16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let base: Vec<half::bf16> = vec![half::bf16::ZERO; 6];
    let indices: Vec<u32> = vec![0, 2, 0];
    let src_f32 = [1.0_f32, 2., 3., 4., 5., 6.];
    let src: Vec<half::bf16> = src_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();
    let expected = [6.0_f32, 8., 0., 0., 3., 4.];

    let base_storage = upload_bf16(&backend, &base);
    let idx_bytes: &[u8] = bytemuck::cast_slice(&indices);
    let idx_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(idx_bytes).expect("idx upload")),
        DType::U32,
    );
    let src_storage = upload_bf16(&backend, &src);
    let out_bytes = backend.alloc_bytes_handle(6 * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let base_arc = Arc::new(RwLock::new(base_storage));
    let idx_arc = Arc::new(RwLock::new(idx_storage));
    let src_arc = Arc::new(RwLock::new(src_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::IndexAdd,
        &[DType::BF16, DType::U32, DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let base_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    let idx_layout = Layout::contiguous(Shape::from_dims(&[3]));
    let src_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    kernel(
        &[
            Arc::clone(&base_arc),
            Arc::clone(&idx_arc),
            Arc::clone(&src_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &[base_layout, idx_layout, src_layout, out_layout],
        &OpParams::IndexAdd {
            outer_count: 1,
            base_dim_size: 3,
            n_indices: 3,
            inner_count: 2,
        },
    )
    .expect("index_add bf16 dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(
            g.to_f32(),
            *e,
            "index_add bf16[{i}]: got {}, expected {e}",
            g.to_f32()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_index_add_f16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let base: Vec<half::f16> = vec![half::f16::ZERO; 6];
    let indices: Vec<u32> = vec![0, 2, 0];
    let src_f32 = [1.0_f32, 2., 3., 4., 5., 6.];
    let src: Vec<half::f16> = src_f32.iter().map(|&x| half::f16::from_f32(x)).collect();
    let expected = [6.0_f32, 8., 0., 0., 3., 4.];

    let base_storage = upload_f16(&backend, &base);
    let idx_bytes: &[u8] = bytemuck::cast_slice(&indices);
    let idx_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(idx_bytes).expect("idx upload")),
        DType::U32,
    );
    let src_storage = upload_f16(&backend, &src);
    let out_bytes = backend.alloc_bytes_handle(6 * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let base_arc = Arc::new(RwLock::new(base_storage));
    let idx_arc = Arc::new(RwLock::new(idx_storage));
    let src_arc = Arc::new(RwLock::new(src_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::IndexAdd,
        &[DType::F16, DType::U32, DType::F16, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let base_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    let idx_layout = Layout::contiguous(Shape::from_dims(&[3]));
    let src_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    kernel(
        &[
            Arc::clone(&base_arc),
            Arc::clone(&idx_arc),
            Arc::clone(&src_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &[base_layout, idx_layout, src_layout, out_layout],
        &OpParams::IndexAdd {
            outer_count: 1,
            base_dim_size: 3,
            n_indices: 3,
            inner_count: 2,
        },
    )
    .expect("index_add f16 dispatch");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(
            g.to_f32(),
            *e,
            "index_add f16[{i}]: got {}, expected {e}",
            g.to_f32()
        );
    }
}

// ---- ScatterAdd f32 (V.3.G.scatter_add, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_scatter_add_f32_dim0() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // base shape [3, 2] = zeros; indices [2, 2]:
    //   [[0, 1], [2, 0]]
    // src [2, 2] = [[1, 2], [3, 4]]
    // dim=0:
    //   src[0,0]=1 → out[0,0]+=1
    //   src[0,1]=2 → out[1,1]+=2
    //   src[1,0]=3 → out[2,0]+=3
    //   src[1,1]=4 → out[0,1]+=4
    // expected: [[1, 4], [0, 2], [3, 0]]
    let base = [0.0_f32, 0., 0., 0., 0., 0.];
    let indices: Vec<u32> = vec![0, 1, 2, 0];
    let src = [1.0_f32, 2., 3., 4.];
    let expected = [1.0_f32, 4., 0., 2., 3., 0.];

    let base_storage = upload_f32(&backend, &base);
    let idx_bytes: &[u8] = bytemuck::cast_slice(&indices);
    let idx_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(idx_bytes).expect("idx upload")),
        DType::U32,
    );
    let src_storage = upload_f32(&backend, &src);
    let out_bytes = backend.alloc_bytes_handle(6 * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let base_arc = Arc::new(RwLock::new(base_storage));
    let idx_arc = Arc::new(RwLock::new(idx_storage));
    let src_arc = Arc::new(RwLock::new(src_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::ScatterAdd,
        &[DType::F32, DType::U32, DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let base_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    let idx_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    let src_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    kernel(
        &[
            Arc::clone(&base_arc),
            Arc::clone(&idx_arc),
            Arc::clone(&src_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &[base_layout, idx_layout, src_layout, out_layout],
        &OpParams::ScatterAdd {
            base_shape: vec![3, 2],
            src_shape: vec![2, 2],
            dim: 0,
        },
    )
    .expect("scatter_add f32 dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "scatter_add f32[{i}]: got {g}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_scatter_add_f32_starts_from_base() {
    // Verify that the wrapper actually copies base → out before the
    // accumulation (i.e. out is NOT zero-initialized, it starts as
    // a copy of base).
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let base = [10.0_f32, 20., 30., 40.]; // shape [2, 2]
    let indices: Vec<u32> = vec![0, 1]; // shape [1, 2]: row 0 → into row {0,1}
    let src = [100.0_f32, 200.]; // shape [1, 2]
    // out[0,0] = base[0,0] + src[0,0] = 110
    // out[1,1] = base[1,1] + src[0,1] = 240
    // out[0,1] = base[0,1] = 20
    // out[1,0] = base[1,0] = 30
    let expected = [110.0_f32, 20., 30., 240.];

    let base_storage = upload_f32(&backend, &base);
    let idx_bytes: &[u8] = bytemuck::cast_slice(&indices);
    let idx_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(idx_bytes).expect("idx upload")),
        DType::U32,
    );
    let src_storage = upload_f32(&backend, &src);
    let out_bytes = backend.alloc_bytes_handle(4 * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let base_arc = Arc::new(RwLock::new(base_storage));
    let idx_arc = Arc::new(RwLock::new(idx_storage));
    let src_arc = Arc::new(RwLock::new(src_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::ScatterAdd,
        &[DType::F32, DType::U32, DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let base_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    let idx_layout = Layout::contiguous(Shape::from_dims(&[1, 2]));
    let src_layout = Layout::contiguous(Shape::from_dims(&[1, 2]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    kernel(
        &[
            Arc::clone(&base_arc),
            Arc::clone(&idx_arc),
            Arc::clone(&src_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &[base_layout, idx_layout, src_layout, out_layout],
        &OpParams::ScatterAdd {
            base_shape: vec![2, 2],
            src_shape: vec![1, 2],
            dim: 0,
        },
    )
    .expect("scatter_add f32 dispatch (base-init test)");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(
            *g, *e,
            "scatter_add f32 base-init[{i}]: got {g}, expected {e}"
        );
    }
}

// ---- ScatterAdd f64 (V.3.G.scatter_add_f64, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_scatter_add_f64_dim0() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // Same shapes/indices as the f32 dim0 test.
    let base = [0.0_f64, 0., 0., 0., 0., 0.];
    let indices: Vec<u32> = vec![0, 1, 2, 0];
    let src = [1.0_f64, 2., 3., 4.];
    let expected = [1.0_f64, 4., 0., 2., 3., 0.];

    let base_storage = upload_f64(&backend, &base);
    let idx_bytes: &[u8] = bytemuck::cast_slice(&indices);
    let idx_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(idx_bytes).expect("idx upload")),
        DType::U32,
    );
    let src_storage = upload_f64(&backend, &src);
    let out_bytes = backend.alloc_bytes_handle(6 * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let base_arc = Arc::new(RwLock::new(base_storage));
    let idx_arc = Arc::new(RwLock::new(idx_storage));
    let src_arc = Arc::new(RwLock::new(src_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::ScatterAdd,
        &[DType::F64, DType::U32, DType::F64, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let base_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    let idx_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    let src_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    kernel(
        &[
            Arc::clone(&base_arc),
            Arc::clone(&idx_arc),
            Arc::clone(&src_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &[base_layout, idx_layout, src_layout, out_layout],
        &OpParams::ScatterAdd {
            base_shape: vec![3, 2],
            src_shape: vec![2, 2],
            dim: 0,
        },
    )
    .expect("scatter_add f64 dispatch");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "scatter_add f64[{i}]: got {g}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_scatter_add_f64_starts_from_base() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let base = [10.0_f64, 20., 30., 40.];
    let indices: Vec<u32> = vec![0, 1];
    let src = [100.0_f64, 200.];
    let expected = [110.0_f64, 20., 30., 240.];

    let base_storage = upload_f64(&backend, &base);
    let idx_bytes: &[u8] = bytemuck::cast_slice(&indices);
    let idx_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(idx_bytes).expect("idx upload")),
        DType::U32,
    );
    let src_storage = upload_f64(&backend, &src);
    let out_bytes = backend.alloc_bytes_handle(4 * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let base_arc = Arc::new(RwLock::new(base_storage));
    let idx_arc = Arc::new(RwLock::new(idx_storage));
    let src_arc = Arc::new(RwLock::new(src_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::ScatterAdd,
        &[DType::F64, DType::U32, DType::F64, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let base_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    let idx_layout = Layout::contiguous(Shape::from_dims(&[1, 2]));
    let src_layout = Layout::contiguous(Shape::from_dims(&[1, 2]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    kernel(
        &[
            Arc::clone(&base_arc),
            Arc::clone(&idx_arc),
            Arc::clone(&src_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &[base_layout, idx_layout, src_layout, out_layout],
        &OpParams::ScatterAdd {
            base_shape: vec![2, 2],
            src_shape: vec![1, 2],
            dim: 0,
        },
    )
    .expect("scatter_add f64 dispatch (base-init test)");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(
            *g, *e,
            "scatter_add f64 base-init[{i}]: got {g}, expected {e}"
        );
    }
}

// ---- ScatterAdd bf16 / f16 (V.3.G.scatter_add_subword, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_scatter_add_bf16_dim0() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let base = [
        half::bf16::ZERO,
        half::bf16::ZERO,
        half::bf16::ZERO,
        half::bf16::ZERO,
        half::bf16::ZERO,
        half::bf16::ZERO,
    ];
    let indices: Vec<u32> = vec![0, 1, 2, 0];
    let src = [
        half::bf16::from_f32(1.0),
        half::bf16::from_f32(2.0),
        half::bf16::from_f32(3.0),
        half::bf16::from_f32(4.0),
    ];
    let expected = [1.0_f32, 4., 0., 2., 3., 0.];

    let base_storage = upload_bf16(&backend, &base);
    let idx_bytes: &[u8] = bytemuck::cast_slice(&indices);
    let idx_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(idx_bytes).expect("idx upload")),
        DType::U32,
    );
    let src_storage = upload_bf16(&backend, &src);
    let out_bytes = backend.alloc_bytes_handle(6 * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let base_arc = Arc::new(RwLock::new(base_storage));
    let idx_arc = Arc::new(RwLock::new(idx_storage));
    let src_arc = Arc::new(RwLock::new(src_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::ScatterAdd,
        &[DType::BF16, DType::U32, DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let base_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    let idx_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    let src_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    kernel(
        &[
            Arc::clone(&base_arc),
            Arc::clone(&idx_arc),
            Arc::clone(&src_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &[base_layout, idx_layout, src_layout, out_layout],
        &OpParams::ScatterAdd {
            base_shape: vec![3, 2],
            src_shape: vec![2, 2],
            dim: 0,
        },
    )
    .expect("scatter_add bf16 dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        // bf16 has ~7 bits mantissa; integer values up to 256 round-trip exactly.
        assert_eq!(
            g.to_f32(),
            *e,
            "scatter_add bf16[{i}]: got {}, expected {e}",
            g.to_f32()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_scatter_add_bf16_starts_from_base() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let base = [
        half::bf16::from_f32(10.0),
        half::bf16::from_f32(20.0),
        half::bf16::from_f32(30.0),
        half::bf16::from_f32(40.0),
    ];
    let indices: Vec<u32> = vec![0, 1];
    let src = [half::bf16::from_f32(100.0), half::bf16::from_f32(200.0)];
    let expected = [110.0_f32, 20., 30., 240.];

    let base_storage = upload_bf16(&backend, &base);
    let idx_bytes: &[u8] = bytemuck::cast_slice(&indices);
    let idx_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(idx_bytes).expect("idx upload")),
        DType::U32,
    );
    let src_storage = upload_bf16(&backend, &src);
    let out_bytes = backend.alloc_bytes_handle(4 * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let base_arc = Arc::new(RwLock::new(base_storage));
    let idx_arc = Arc::new(RwLock::new(idx_storage));
    let src_arc = Arc::new(RwLock::new(src_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::ScatterAdd,
        &[DType::BF16, DType::U32, DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let base_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    let idx_layout = Layout::contiguous(Shape::from_dims(&[1, 2]));
    let src_layout = Layout::contiguous(Shape::from_dims(&[1, 2]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    kernel(
        &[
            Arc::clone(&base_arc),
            Arc::clone(&idx_arc),
            Arc::clone(&src_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &[base_layout, idx_layout, src_layout, out_layout],
        &OpParams::ScatterAdd {
            base_shape: vec![2, 2],
            src_shape: vec![1, 2],
            dim: 0,
        },
    )
    .expect("scatter_add bf16 dispatch (base-init)");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(
            g.to_f32(),
            *e,
            "scatter_add bf16 base-init[{i}]: got {}, expected {e}",
            g.to_f32()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_scatter_add_f16_dim0() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let base = [
        half::f16::from_f32(0.0),
        half::f16::from_f32(0.0),
        half::f16::from_f32(0.0),
        half::f16::from_f32(0.0),
        half::f16::from_f32(0.0),
        half::f16::from_f32(0.0),
    ];
    let indices: Vec<u32> = vec![0, 1, 2, 0];
    let src = [
        half::f16::from_f32(1.0),
        half::f16::from_f32(2.0),
        half::f16::from_f32(3.0),
        half::f16::from_f32(4.0),
    ];
    let expected = [1.0_f32, 4., 0., 2., 3., 0.];

    let base_storage = upload_f16(&backend, &base);
    let idx_bytes: &[u8] = bytemuck::cast_slice(&indices);
    let idx_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(idx_bytes).expect("idx upload")),
        DType::U32,
    );
    let src_storage = upload_f16(&backend, &src);
    let out_bytes = backend.alloc_bytes_handle(6 * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let base_arc = Arc::new(RwLock::new(base_storage));
    let idx_arc = Arc::new(RwLock::new(idx_storage));
    let src_arc = Arc::new(RwLock::new(src_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::ScatterAdd,
        &[DType::F16, DType::U32, DType::F16, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let base_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    let idx_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    let src_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
    kernel(
        &[
            Arc::clone(&base_arc),
            Arc::clone(&idx_arc),
            Arc::clone(&src_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &[base_layout, idx_layout, src_layout, out_layout],
        &OpParams::ScatterAdd {
            base_shape: vec![3, 2],
            src_shape: vec![2, 2],
            dim: 0,
        },
    )
    .expect("scatter_add f16 dispatch");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(
            g.to_f32(),
            *e,
            "scatter_add f16[{i}]: got {}, expected {e}",
            g.to_f32()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_scatter_add_f16_starts_from_base() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let base = [
        half::f16::from_f32(10.0),
        half::f16::from_f32(20.0),
        half::f16::from_f32(30.0),
        half::f16::from_f32(40.0),
    ];
    let indices: Vec<u32> = vec![0, 1];
    let src = [half::f16::from_f32(100.0), half::f16::from_f32(200.0)];
    let expected = [110.0_f32, 20., 30., 240.];

    let base_storage = upload_f16(&backend, &base);
    let idx_bytes: &[u8] = bytemuck::cast_slice(&indices);
    let idx_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(idx_bytes).expect("idx upload")),
        DType::U32,
    );
    let src_storage = upload_f16(&backend, &src);
    let out_bytes = backend.alloc_bytes_handle(4 * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let base_arc = Arc::new(RwLock::new(base_storage));
    let idx_arc = Arc::new(RwLock::new(idx_storage));
    let src_arc = Arc::new(RwLock::new(src_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::ScatterAdd,
        &[DType::F16, DType::U32, DType::F16, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let base_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    let idx_layout = Layout::contiguous(Shape::from_dims(&[1, 2]));
    let src_layout = Layout::contiguous(Shape::from_dims(&[1, 2]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    kernel(
        &[
            Arc::clone(&base_arc),
            Arc::clone(&idx_arc),
            Arc::clone(&src_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &[base_layout, idx_layout, src_layout, out_layout],
        &OpParams::ScatterAdd {
            base_shape: vec![2, 2],
            src_shape: vec![1, 2],
            dim: 0,
        },
    )
    .expect("scatter_add f16 dispatch (base-init)");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(
            g.to_f32(),
            *e,
            "scatter_add f16 base-init[{i}]: got {}, expected {e}",
            g.to_f32()
        );
    }
}

// ---- ArgMaxDim / ArgMinDim along last dim (V.3.G.arg_reduce, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_argmax_last_dim_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let input = [1.0_f32, 3.0, 2.0, 0.5, 9.0, 4.0, 9.0, 7.0];
    // row 0: max is 3.0 at idx 1
    // row 1: max is 9.0 at idx 0 (lower of ties)
    let expected: [u32; 2] = [1, 0];

    let in_storage = upload_f32(&backend, &input);
    let out_bytes = backend.alloc_bytes_handle(outer * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::U32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::ArgMaxDim,
        &[DType::F32, DType::U32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[outer]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout, out_layout],
        &OpParams::Reduce {
            dims: vec![1],
            keepdim: false,
        },
    )
    .expect("argmax f32 dispatch");

    let bytes = match &out_arc.read().unwrap().inner {
        BackendStorage::Vulkan(v) => backend.download_bytes(v).expect("d2h"),
        _ => panic!("not on Vulkan"),
    };
    let got: Vec<u32> = bytemuck::cast_slice::<u8, u32>(&bytes).to_vec();
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "argmax f32[{i}]: got {g}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_argmin_last_dim_bf16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize; // even for bf16 lane-pair
    let input_f32 = [5.0_f32, -1.0, 3.0, 0.0, 2.0, 4.0, -3.0, 1.0];
    let input: Vec<half::bf16> = input_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();
    let expected: [u32; 2] = [1, 2];

    let in_storage = upload_bf16(&backend, &input);
    let out_bytes = backend.alloc_bytes_handle(outer * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::U32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::ArgMinDim,
        &[DType::BF16, DType::U32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[outer]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout, out_layout],
        &OpParams::Reduce {
            dims: vec![1],
            keepdim: false,
        },
    )
    .expect("argmin bf16 dispatch");

    let bytes = match &out_arc.read().unwrap().inner {
        BackendStorage::Vulkan(v) => backend.download_bytes(v).expect("d2h"),
        _ => panic!("not on Vulkan"),
    };
    let got: Vec<u32> = bytemuck::cast_slice::<u8, u32>(&bytes).to_vec();
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "argmin bf16[{i}]: got {g}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_argmax_last_dim_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 1usize;
    let last = 5usize;
    let input = [1.0_f64, 10.0, 5.0, 10.0, 7.0]; // ties at idx 1 and 3; expect 1
    let expected: [u32; 1] = [1];

    let in_storage = upload_f64(&backend, &input);
    let out_bytes = backend.alloc_bytes_handle(outer * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::U32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::ArgMaxDim,
        &[DType::F64, DType::U32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[outer]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout, out_layout],
        &OpParams::Reduce {
            dims: vec![1],
            keepdim: false,
        },
    )
    .expect("argmax f64 dispatch");

    let bytes = match &out_arc.read().unwrap().inner {
        BackendStorage::Vulkan(v) => backend.download_bytes(v).expect("d2h"),
        _ => panic!("not on Vulkan"),
    };
    let got: Vec<u32> = bytemuck::cast_slice::<u8, u32>(&bytes).to_vec();
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "argmax f64[{i}]: got {g}, expected {e}");
    }
}

// ---- ArgMaxDim / ArgMinDim along an ARBITRARY dim (V.3.G.arg_reduce_any_dim, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_argmax_any_dim_f32_dim0() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // shape [3, 4], argmax along dim=0 → output shape [4]
    // input:
    //   row 0: [1, 7, 2, 3]
    //   row 1: [5, 3, 6, 0]
    //   row 2: [4, 7, 9, 2]
    // per-column max:
    //   col 0: max=5 at row 1
    //   col 1: max=7 at rows {0,2} → lower index → 0
    //   col 2: max=9 at row 2
    //   col 3: max=3 at row 0
    let input = [1.0_f32, 7., 2., 3., 5., 3., 6., 0., 4., 7., 9., 2.];
    let expected: [u32; 4] = [1, 0, 2, 0];

    let in_storage = upload_f32(&backend, &input);
    let out_bytes = backend.alloc_bytes_handle(4 * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::U32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::ArgMaxDim,
        &[DType::F32, DType::U32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let in_layout = Layout::contiguous(Shape::from_dims(&[3, 4]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[4]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[in_layout, out_layout],
        &OpParams::Reduce {
            dims: vec![0],
            keepdim: false,
        },
    )
    .expect("argmax f32 dim0 dispatch");

    let bytes = match &out_arc.read().unwrap().inner {
        BackendStorage::Vulkan(v) => backend.download_bytes(v).expect("d2h"),
        _ => panic!("not on Vulkan"),
    };
    let got: Vec<u32> = bytemuck::cast_slice::<u8, u32>(&bytes).to_vec();
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "argmax f32 dim0[{i}]: got {g}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_argmin_any_dim_f32_middle() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // shape [2, 3, 4], argmin along dim=1 (the middle / interior axis)
    // → output shape [2, 4]
    // batch 0:
    //   [[ 5, 1, 4, 9],
    //    [ 2, 8, 3, 0],
    //    [ 7, 6, 4, 1]]
    //   per-col min: col0 min=2 at 1; col1 min=1 at 0; col2 min=3 at 1; col3 min=0 at 1
    // batch 1:
    //   [[ 1, 2, 3, 4],
    //    [ 4, 3, 2, 1],
    //    [ 2, 2, 2, 2]]
    //   per-col min: col0 min=1 at 0; col1 min=2 at {0,2} → 0;
    //                col2 min=2 at {1,2} → 1; col3 min=1 at 1
    let input = [
        // batch 0
        5.0_f32, 1., 4., 9., 2., 8., 3., 0., 7., 6., 4., 1., // batch 1
        1., 2., 3., 4., 4., 3., 2., 1., 2., 2., 2., 2.,
    ];
    let expected: [u32; 8] = [1, 0, 1, 1, 0, 0, 1, 1];

    let in_storage = upload_f32(&backend, &input);
    let out_bytes = backend.alloc_bytes_handle(8 * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::U32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::ArgMinDim,
        &[DType::F32, DType::U32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let in_layout = Layout::contiguous(Shape::from_dims(&[2, 3, 4]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[2, 4]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[in_layout, out_layout],
        &OpParams::Reduce {
            dims: vec![1],
            keepdim: false,
        },
    )
    .expect("argmin f32 middle dispatch");

    let bytes = match &out_arc.read().unwrap().inner {
        BackendStorage::Vulkan(v) => backend.download_bytes(v).expect("d2h"),
        _ => panic!("not on Vulkan"),
    };
    let got: Vec<u32> = bytemuck::cast_slice::<u8, u32>(&bytes).to_vec();
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "argmin f32 middle[{i}]: got {g}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_argmax_any_dim_bf16_dim0() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // shape [3, 4], dim=0; same per-column expected as the f32 test
    // above. Reduction direction is strided (stride=4 in bf16 lanes)
    // → exercises the sub-word lane-select read path.
    let input_f32 = [1.0_f32, 7., 2., 3., 5., 3., 6., 0., 4., 7., 9., 2.];
    let input: Vec<half::bf16> = input_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();
    let expected: [u32; 4] = [1, 0, 2, 0];

    let in_storage = upload_bf16(&backend, &input);
    let out_bytes = backend.alloc_bytes_handle(4 * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::U32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::ArgMaxDim,
        &[DType::BF16, DType::U32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let in_layout = Layout::contiguous(Shape::from_dims(&[3, 4]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[4]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[in_layout, out_layout],
        &OpParams::Reduce {
            dims: vec![0],
            keepdim: false,
        },
    )
    .expect("argmax bf16 dim0 dispatch");

    let bytes = match &out_arc.read().unwrap().inner {
        BackendStorage::Vulkan(v) => backend.download_bytes(v).expect("d2h"),
        _ => panic!("not on Vulkan"),
    };
    let got: Vec<u32> = bytemuck::cast_slice::<u8, u32>(&bytes).to_vec();
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "argmax bf16 dim0[{i}]: got {g}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_argmin_any_dim_f16_dim0() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let input_f32 = [5.0_f32, -1., 3., 0., 2., 4., -3., 1., 7., 6., 4., 1.];
    let input: Vec<half::f16> = input_f32.iter().map(|&x| half::f16::from_f32(x)).collect();
    // dim=0 argmin per column:
    //   col0: min=2 at row 1
    //   col1: min=-1 at row 0
    //   col2: min=-3 at row 1
    //   col3: min=0 at row 0
    let expected: [u32; 4] = [1, 0, 1, 0];

    let in_storage = upload_f16(&backend, &input);
    let out_bytes = backend.alloc_bytes_handle(4 * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::U32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::ArgMinDim,
        &[DType::F16, DType::U32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let in_layout = Layout::contiguous(Shape::from_dims(&[3, 4]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[4]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[in_layout, out_layout],
        &OpParams::Reduce {
            dims: vec![0],
            keepdim: false,
        },
    )
    .expect("argmin f16 dim0 dispatch");

    let bytes = match &out_arc.read().unwrap().inner {
        BackendStorage::Vulkan(v) => backend.download_bytes(v).expect("d2h"),
        _ => panic!("not on Vulkan"),
    };
    let got: Vec<u32> = bytemuck::cast_slice::<u8, u32>(&bytes).to_vec();
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "argmin f16 dim0[{i}]: got {g}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_argmax_any_dim_f64_dim0() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let input = [1.0_f64, 7., 2., 3., 5., 3., 6., 0., 4., 7., 9., 2.];
    let expected: [u32; 4] = [1, 0, 2, 0];

    let in_storage = upload_f64(&backend, &input);
    let out_bytes = backend.alloc_bytes_handle(4 * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::U32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::ArgMaxDim,
        &[DType::F64, DType::U32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let in_layout = Layout::contiguous(Shape::from_dims(&[3, 4]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[4]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[in_layout, out_layout],
        &OpParams::Reduce {
            dims: vec![0],
            keepdim: false,
        },
    )
    .expect("argmax f64 dim0 dispatch");

    let bytes = match &out_arc.read().unwrap().inner {
        BackendStorage::Vulkan(v) => backend.download_bytes(v).expect("d2h"),
        _ => panic!("not on Vulkan"),
    };
    let got: Vec<u32> = bytemuck::cast_slice::<u8, u32>(&bytes).to_vec();
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "argmax f64 dim0[{i}]: got {g}, expected {e}");
    }
}

// ---- PadBackward reflect / replicate f32 (V.3.G.pad_backward.atomic, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_pad_backward_reflect_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // Forward reflect: in=[3] pad=(2,2) → out=[7] = [3,2,1,2,3,2,1]
    // Backward (each out coord c reflects to some in_coord; accumulate):
    //   c=0 → in[2] += 10
    //   c=1 → in[1] += 20
    //   c=2 → in[0] += 30
    //   c=3 → in[1] += 40
    //   c=4 → in[2] += 50
    //   c=5 → in[1] += 60
    //   c=6 → in[0] += 70
    // expected: in[0]=30+70=100, in[1]=20+40+60=120, in[2]=10+50=60
    let grad_out: Vec<f32> = vec![10., 20., 30., 40., 50., 60., 70.];
    let expected: [f32; 3] = [100., 120., 60.];

    let go_storage = upload_f32(&backend, &grad_out);
    let gi_bytes = backend.alloc_bytes_handle(3 * 4).expect("alloc");
    let gi_storage = Storage::new(BackendStorage::Vulkan(gi_bytes), DType::F32);
    let go_arc = Arc::new(RwLock::new(go_storage));
    let gi_arc = Arc::new(RwLock::new(gi_storage));

    let kernel = table.lookup_alternatives(
        OpKind::PadBackward,
        &[DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let go_layout = Layout::contiguous(Shape::from_dims(&[7]));
    let gi_layout = Layout::contiguous(Shape::from_dims(&[3]));
    kernel(
        &[Arc::clone(&go_arc)],
        &mut [Arc::clone(&gi_arc)],
        &[go_layout, gi_layout],
        &OpParams::PadBackward {
            in_shape: vec![3],
            out_shape: vec![7],
            padding: vec![(2, 2)],
            mode_tag: 1,
        },
    )
    .expect("pad_backward reflect f32 dispatch");

    let got = download_f32(&backend, &gi_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert!(
            (g - e).abs() < 1e-5,
            "pad_backward reflect[{i}]: got {g}, expected {e}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_pad_backward_replicate_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // Forward replicate: in=[3] pad=(2,3) → out=[8] = [1,1,1,2,3,3,3,3]
    // Backward (clamp each out coord):
    //   c=0,1,2: in[0] += grad_out[0..3]   → in[0] = 10+20+30 = 60
    //   c=3: in[1] += grad_out[3]          → in[1] = 40
    //   c=4..7: in[2] += grad_out[4..8]    → in[2] = 50+60+70+80 = 260
    let grad_out: Vec<f32> = vec![10., 20., 30., 40., 50., 60., 70., 80.];
    let expected: [f32; 3] = [60., 40., 260.];

    let go_storage = upload_f32(&backend, &grad_out);
    let gi_bytes = backend.alloc_bytes_handle(3 * 4).expect("alloc");
    let gi_storage = Storage::new(BackendStorage::Vulkan(gi_bytes), DType::F32);
    let go_arc = Arc::new(RwLock::new(go_storage));
    let gi_arc = Arc::new(RwLock::new(gi_storage));

    let kernel = table.lookup_alternatives(
        OpKind::PadBackward,
        &[DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let go_layout = Layout::contiguous(Shape::from_dims(&[8]));
    let gi_layout = Layout::contiguous(Shape::from_dims(&[3]));
    kernel(
        &[Arc::clone(&go_arc)],
        &mut [Arc::clone(&gi_arc)],
        &[go_layout, gi_layout],
        &OpParams::PadBackward {
            in_shape: vec![3],
            out_shape: vec![8],
            padding: vec![(2, 3)],
            mode_tag: 2,
        },
    )
    .expect("pad_backward replicate f32 dispatch");

    let got = download_f32(&backend, &gi_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert!(
            (g - e).abs() < 1e-5,
            "pad_backward replicate[{i}]: got {g}, expected {e}"
        );
    }
}

// ---- PadBackward reflect / replicate bf16+f16+f64 (V.3.G.pad_backward.atomic.subword+f64, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_pad_backward_reflect_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // Same shape/expected as the f32 reflect test.
    let grad_out: Vec<f64> = vec![10., 20., 30., 40., 50., 60., 70.];
    let expected: [f64; 3] = [100., 120., 60.];

    let go_storage = upload_f64(&backend, &grad_out);
    let gi_bytes = backend.alloc_bytes_handle(3 * 8).expect("alloc");
    let gi_storage = Storage::new(BackendStorage::Vulkan(gi_bytes), DType::F64);
    let go_arc = Arc::new(RwLock::new(go_storage));
    let gi_arc = Arc::new(RwLock::new(gi_storage));

    let kernel = table.lookup_alternatives(
        OpKind::PadBackward,
        &[DType::F64, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let go_layout = Layout::contiguous(Shape::from_dims(&[7]));
    let gi_layout = Layout::contiguous(Shape::from_dims(&[3]));
    kernel(
        &[Arc::clone(&go_arc)],
        &mut [Arc::clone(&gi_arc)],
        &[go_layout, gi_layout],
        &OpParams::PadBackward {
            in_shape: vec![3],
            out_shape: vec![7],
            padding: vec![(2, 2)],
            mode_tag: 1,
        },
    )
    .expect("pad_backward reflect f64 dispatch");

    let got = download_f64(&backend, &gi_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert!(
            (g - e).abs() < 1e-10,
            "pad_backward reflect f64[{i}]: got {g}, expected {e}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_pad_backward_replicate_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let grad_out: Vec<f64> = vec![10., 20., 30., 40., 50., 60., 70., 80.];
    let expected: [f64; 3] = [60., 40., 260.];

    let go_storage = upload_f64(&backend, &grad_out);
    let gi_bytes = backend.alloc_bytes_handle(3 * 8).expect("alloc");
    let gi_storage = Storage::new(BackendStorage::Vulkan(gi_bytes), DType::F64);
    let go_arc = Arc::new(RwLock::new(go_storage));
    let gi_arc = Arc::new(RwLock::new(gi_storage));

    let kernel = table.lookup_alternatives(
        OpKind::PadBackward,
        &[DType::F64, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let go_layout = Layout::contiguous(Shape::from_dims(&[8]));
    let gi_layout = Layout::contiguous(Shape::from_dims(&[3]));
    kernel(
        &[Arc::clone(&go_arc)],
        &mut [Arc::clone(&gi_arc)],
        &[go_layout, gi_layout],
        &OpParams::PadBackward {
            in_shape: vec![3],
            out_shape: vec![8],
            padding: vec![(2, 3)],
            mode_tag: 2,
        },
    )
    .expect("pad_backward replicate f64 dispatch");

    let got = download_f64(&backend, &gi_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert!(
            (g - e).abs() < 1e-10,
            "pad_backward replicate f64[{i}]: got {g}, expected {e}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_pad_backward_reflect_bf16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let grad_out_f32: [f32; 7] = [10., 20., 30., 40., 50., 60., 70.];
    let grad_out: Vec<half::bf16> = grad_out_f32
        .iter()
        .map(|&x| half::bf16::from_f32(x))
        .collect();
    let expected: [f32; 3] = [100., 120., 60.];

    let go_storage = upload_bf16(&backend, &grad_out);
    // 3 bf16 = 6 bytes; round up to u32 (8 bytes) for sub-word CAS safety.
    let gi_bytes = backend
        .alloc_bytes_handle(((3 * 2 + 3) & !3) as usize)
        .expect("alloc");
    let gi_storage = Storage::new(BackendStorage::Vulkan(gi_bytes), DType::BF16);
    let go_arc = Arc::new(RwLock::new(go_storage));
    let gi_arc = Arc::new(RwLock::new(gi_storage));

    let kernel = table.lookup_alternatives(
        OpKind::PadBackward,
        &[DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let go_layout = Layout::contiguous(Shape::from_dims(&[7]));
    let gi_layout = Layout::contiguous(Shape::from_dims(&[3]));
    kernel(
        &[Arc::clone(&go_arc)],
        &mut [Arc::clone(&gi_arc)],
        &[go_layout, gi_layout],
        &OpParams::PadBackward {
            in_shape: vec![3],
            out_shape: vec![7],
            padding: vec![(2, 2)],
            mode_tag: 1,
        },
    )
    .expect("pad_backward reflect bf16 dispatch");

    let got = download_bf16(&backend, &gi_arc.read().unwrap());
    // bf16 stores integer values <256 exactly; expected sums (100, 120, 60) all fit.
    for (i, (g, e)) in got.iter().take(3).zip(expected.iter()).enumerate() {
        assert_eq!(
            g.to_f32(),
            *e,
            "pad_backward reflect bf16[{i}]: got {}, expected {e}",
            g.to_f32()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_pad_backward_replicate_f16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let grad_out_f32: [f32; 8] = [10., 20., 30., 40., 50., 60., 70., 80.];
    let grad_out: Vec<half::f16> = grad_out_f32
        .iter()
        .map(|&x| half::f16::from_f32(x))
        .collect();
    let expected: [f32; 3] = [60., 40., 260.];

    let go_storage = upload_f16(&backend, &grad_out);
    let gi_bytes = backend
        .alloc_bytes_handle(((3 * 2 + 3) & !3) as usize)
        .expect("alloc");
    let gi_storage = Storage::new(BackendStorage::Vulkan(gi_bytes), DType::F16);
    let go_arc = Arc::new(RwLock::new(go_storage));
    let gi_arc = Arc::new(RwLock::new(gi_storage));

    let kernel = table.lookup_alternatives(
        OpKind::PadBackward,
        &[DType::F16, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let go_layout = Layout::contiguous(Shape::from_dims(&[8]));
    let gi_layout = Layout::contiguous(Shape::from_dims(&[3]));
    kernel(
        &[Arc::clone(&go_arc)],
        &mut [Arc::clone(&gi_arc)],
        &[go_layout, gi_layout],
        &OpParams::PadBackward {
            in_shape: vec![3],
            out_shape: vec![8],
            padding: vec![(2, 3)],
            mode_tag: 2,
        },
    )
    .expect("pad_backward replicate f16 dispatch");

    let got = download_f16(&backend, &gi_arc.read().unwrap());
    // f16 representable: 60 ✓, 40 ✓, 260 ✓ (max integer exactly representable: 2048).
    for (i, (g, e)) in got.iter().take(3).zip(expected.iter()).enumerate() {
        assert_eq!(
            g.to_f32(),
            *e,
            "pad_backward replicate f16[{i}]: got {}, expected {e}",
            g.to_f32()
        );
    }
}

// ---- PadBackward (constant mode) (V.3.G.pad_backward.const, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_pad_backward_const_f32_2d() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // Forward Pad was in=[2,3], padding=[(1,1),(0,2)] → out=[4,5].
    // Backward: grad_in[r,c] = grad_out[r + 1, c + 0]
    //   grad_in shape [2,3], grad_out shape [4,5].
    let grad_out: Vec<f32> = (0..20).map(|i| i as f32).collect();
    // grad_in[0,0] = grad_out[1, 0] = 5
    // grad_in[0,1] = grad_out[1, 1] = 6
    // ...
    let expected = [5.0_f32, 6.0, 7., 10., 11., 12.];

    let go_storage = upload_f32(&backend, &grad_out);
    let gi_bytes = backend.alloc_bytes_handle(2 * 3 * 4).expect("alloc");
    let gi_storage = Storage::new(BackendStorage::Vulkan(gi_bytes), DType::F32);
    let go_arc = Arc::new(RwLock::new(go_storage));
    let gi_arc = Arc::new(RwLock::new(gi_storage));

    let kernel = table.lookup_alternatives(
        OpKind::PadBackward,
        &[DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let go_layout = Layout::contiguous(Shape::from_dims(&[4, 5]));
    let gi_layout = Layout::contiguous(Shape::from_dims(&[2, 3]));
    kernel(
        &[Arc::clone(&go_arc)],
        &mut [Arc::clone(&gi_arc)],
        &[go_layout, gi_layout],
        &OpParams::PadBackward {
            in_shape: vec![2, 3],
            out_shape: vec![4, 5],
            padding: vec![(1, 1), (0, 2)],
            mode_tag: 0,
        },
    )
    .expect("pad_backward const f32 dispatch");

    let got = download_f32(&backend, &gi_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "pad_backward const f32[{i}]: got {g}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_pad_backward_const_bf16_1d() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // Forward: in=[4] pad=(2,2) → out=[8]. Backward: in[i] = out[i+2].
    let go_f32 = [10.0_f32, 20., 30., 40., 50., 60., 70., 80.];
    let go: Vec<half::bf16> = go_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();
    let expected_f32 = [30.0_f32, 40., 50., 60.];

    let go_storage = upload_bf16(&backend, &go);
    let gi_bytes = backend.alloc_bytes_handle(4 * 2).expect("alloc");
    let gi_storage = Storage::new(BackendStorage::Vulkan(gi_bytes), DType::BF16);
    let go_arc = Arc::new(RwLock::new(go_storage));
    let gi_arc = Arc::new(RwLock::new(gi_storage));

    let kernel = table.lookup_alternatives(
        OpKind::PadBackward,
        &[DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let go_layout = Layout::contiguous(Shape::from_dims(&[8]));
    let gi_layout = Layout::contiguous(Shape::from_dims(&[4]));
    kernel(
        &[Arc::clone(&go_arc)],
        &mut [Arc::clone(&gi_arc)],
        &[go_layout, gi_layout],
        &OpParams::PadBackward {
            in_shape: vec![4],
            out_shape: vec![8],
            padding: vec![(2, 2)],
            mode_tag: 0,
        },
    )
    .expect("pad_backward const bf16 dispatch");

    let got = download_bf16(&backend, &gi_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected_f32.iter()).enumerate() {
        assert_eq!(
            g.to_f32(),
            *e,
            "pad_backward const bf16[{i}]: got {}, expected {e}",
            g.to_f32()
        );
    }
}

// ---- Pad replicate mode (V.3.G.pad.replicate, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_pad_replicate_f32_1d() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // in=[1,2,3], padding=(2,3) → out=[1,1,1,2,3,3,3,3]
    let input = [1.0_f32, 2.0, 3.0];
    let expected = [1.0_f32, 1.0, 1.0, 2.0, 3.0, 3.0, 3.0, 3.0];

    let in_storage = upload_f32(&backend, &input);
    let out_bytes = backend.alloc_bytes_handle(8 * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Pad, &[DType::F32, DType::F32], BackendId::Vulkan)[0]
            .kernel;
    let in_layout = Layout::contiguous(Shape::from_dims(&[3]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[8]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[in_layout, out_layout],
        &OpParams::Pad {
            in_shape: vec![3],
            out_shape: vec![8],
            padding: vec![(2, 3)],
            mode_tag: 2,
            fill_bytes: vec![0u8; 4],
        },
    )
    .expect("pad replicate f32 dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "pad replicate f32[{i}]: got {g}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_pad_replicate_f16_2d() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // in=[2,2], padding=[(1,1), (1,1)] → out=[4,4] = 16 (even, ok for b2)
    let input_f32 = [1.0_f32, 2.0, 3.0, 4.0];
    let input: Vec<half::f16> = input_f32.iter().map(|&x| half::f16::from_f32(x)).collect();
    let in_2d = |r: usize, c: usize| input_f32[r * 2 + c];

    // Replicate row map: r=0→0, r=1→0, r=2→1, r=3→1
    // Replicate col map: same shape.
    let row_map = [0_usize, 0, 1, 1];
    let col_map = [0_usize, 0, 1, 1];
    let mut expected = [0.0_f32; 4 * 4];
    for r in 0..4 {
        for c in 0..4 {
            expected[r * 4 + c] = in_2d(row_map[r], col_map[c]);
        }
    }

    let in_storage = upload_f16(&backend, &input);
    let out_bytes = backend.alloc_bytes_handle(4 * 4 * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Pad, &[DType::F16, DType::F16], BackendId::Vulkan)[0]
            .kernel;
    let in_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[4, 4]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[in_layout, out_layout],
        &OpParams::Pad {
            in_shape: vec![2, 2],
            out_shape: vec![4, 4],
            padding: vec![(1, 1), (1, 1)],
            mode_tag: 2,
            fill_bytes: vec![0u8; 2],
        },
    )
    .expect("pad replicate f16 dispatch");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(
            g.to_f32(),
            *e,
            "pad replicate f16[{i}]: got {}, expected {e}",
            g.to_f32()
        );
    }
}

// ---- Pad reflect mode (V.3.G.pad.reflect, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_pad_reflect_f32_1d() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // 1D reflect: in=[1,2,3], padding=(2, 2) → out=[8,5] = 5.
    // Reference per the CPU reflect_index:
    //   c=0: i=-2 → -i=2 → in[2]=3
    //   c=1: i=-1 → -i=1 → in[1]=2
    //   c=2: i= 0           → in[0]=1
    //   c=3: i= 1           → in[1]=2
    //   c=4: i= 2           → in[2]=3
    //   c=5: i= 3 (>=3) → 2*2-3=1 → in[1]=2
    //   c=6: i= 4 (>=3) → 2*2-4=0 → in[0]=1
    let input = [1.0_f32, 2.0, 3.0];
    let expected = [3.0_f32, 2.0, 1.0, 2.0, 3.0, 2.0, 1.0];

    let in_storage = upload_f32(&backend, &input);
    let out_bytes = backend.alloc_bytes_handle(7 * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Pad, &[DType::F32, DType::F32], BackendId::Vulkan)[0]
            .kernel;
    let in_layout = Layout::contiguous(Shape::from_dims(&[3]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[7]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[in_layout, out_layout],
        &OpParams::Pad {
            in_shape: vec![3],
            out_shape: vec![7],
            padding: vec![(2, 2)],
            mode_tag: 1,
            fill_bytes: vec![0u8; 4], // unused for reflect
        },
    )
    .expect("pad reflect f32 dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "pad reflect f32[{i}]: got {g}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_pad_reflect_bf16_2d() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // 2D reflect: in=[2,3], padding=[(1,1), (1,1)] → out=[4,5] = 20
    // (even, satisfies b2 pair-thread constraint).
    let input_f32 = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0];
    let input: Vec<half::bf16> = input_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();

    // Reference per axis: each row d∈{0,1} maps via reflect_index.
    // Row 0 (out_r=0): in_r=1 (reflect c=0 with left=1: -i = 1 → in[1])
    // Row 1 (out_r=1): in_r=0
    // Row 2 (out_r=2): in_r=1
    // Row 3 (out_r=3): in_r=0 (reflect: i=2 (>=2) → 2*1-2=0)
    // Col reflect (left=1):
    //   c=0 → in_c=1
    //   c=1 → in_c=0
    //   c=2 → in_c=1
    //   c=3 → in_c=2
    //   c=4 → in_c=1 (i=3 (>=3) → 2*2-3=1)
    // Composing: out[r,c] = in[reflect_row(r), reflect_col(c)]
    let in_2d = |r: usize, c: usize| input_f32[r * 3 + c];
    let mut expected = [0.0_f32; 4 * 5];
    let row_map = [1usize, 0, 1, 0];
    let col_map = [1usize, 0, 1, 2, 1];
    for r in 0..4 {
        for c in 0..5 {
            expected[r * 5 + c] = in_2d(row_map[r], col_map[c]);
        }
    }

    let in_storage = upload_bf16(&backend, &input);
    let out_bytes = backend.alloc_bytes_handle(4 * 5 * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Pad, &[DType::BF16, DType::BF16], BackendId::Vulkan)[0]
            .kernel;
    let in_layout = Layout::contiguous(Shape::from_dims(&[2, 3]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[4, 5]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[in_layout, out_layout],
        &OpParams::Pad {
            in_shape: vec![2, 3],
            out_shape: vec![4, 5],
            padding: vec![(1, 1), (1, 1)],
            mode_tag: 1,
            fill_bytes: vec![0u8; 2],
        },
    )
    .expect("pad reflect bf16 dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        let got_f32 = g.to_f32();
        assert_eq!(
            got_f32, *e,
            "pad reflect bf16[{i}]: got {got_f32}, expected {e}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_pad_reflect_f64_1d() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let input = [1.0_f64, 2.0, 3.0, 4.0];
    let expected = [2.0_f64, 1.0, 2.0, 3.0, 4.0, 3.0, 2.0]; // pad=(1,2)

    let in_storage = upload_f64(&backend, &input);
    let out_bytes = backend.alloc_bytes_handle(7 * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Pad, &[DType::F64, DType::F64], BackendId::Vulkan)[0]
            .kernel;
    let in_layout = Layout::contiguous(Shape::from_dims(&[4]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[7]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[in_layout, out_layout],
        &OpParams::Pad {
            in_shape: vec![4],
            out_shape: vec![7],
            padding: vec![(1, 2)],
            mode_tag: 1,
            fill_bytes: vec![0u8; 8],
        },
    )
    .expect("pad reflect f64 dispatch");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "pad reflect f64[{i}]: got {g}, expected {e}");
    }
}

// ---- Pad (constant mode) f32/f16/bf16/f64/u8 (V.3.G.pad, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_pad_const_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // in=[2,3], padding=[(1,1), (0,2)] → out=[4,5]
    // Layout (out, row × col):
    //   row 0: [fill,fill,fill,fill,fill]
    //   row 1: [1,2,3,fill,fill]
    //   row 2: [4,5,6,fill,fill]
    //   row 3: [fill,fill,fill,fill,fill]
    let input = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0];
    let fill: f32 = -7.5;

    let in_storage = upload_f32(&backend, &input);
    let out_bytes = backend.alloc_bytes_handle(4 * 5 * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Pad, &[DType::F32, DType::F32], BackendId::Vulkan)[0]
            .kernel;
    let in_layout = Layout::contiguous(Shape::from_dims(&[2, 3]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[4, 5]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[in_layout, out_layout],
        &OpParams::Pad {
            in_shape: vec![2, 3],
            out_shape: vec![4, 5],
            padding: vec![(1, 1), (0, 2)],
            mode_tag: 0,
            fill_bytes: fill.to_le_bytes().to_vec(),
        },
    )
    .expect("pad f32 dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    let expected = [
        fill, fill, fill, fill, fill, 1.0, 2.0, 3.0, fill, fill, 4.0, 5.0, 6.0, fill, fill, fill,
        fill, fill, fill, fill,
    ];
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "pad f32[{i}]: got {g}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_pad_const_f16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // in=[1,4], padding=[(0,0), (1,1)] → out=[1,6] (last-dim even)
    let input_f32 = [1.0_f32, 2.0, 3.0, 4.0];
    let input: Vec<half::f16> = input_f32.iter().map(|&x| half::f16::from_f32(x)).collect();
    let fill = half::f16::from_f32(0.0);

    let in_storage = upload_f16(&backend, &input);
    // `1 * 6 * 2` is out=[1,6] (see the comment above) times 2 bytes per f16.
    // Every factor names a dimension; folding the leading 1 hides the batch.
    #[allow(clippy::identity_op)]
    let out_bytes = backend.alloc_bytes_handle(1 * 6 * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Pad, &[DType::F16, DType::F16], BackendId::Vulkan)[0]
            .kernel;
    let in_layout = Layout::contiguous(Shape::from_dims(&[1, 4]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[1, 6]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[in_layout, out_layout],
        &OpParams::Pad {
            in_shape: vec![1, 4],
            out_shape: vec![1, 6],
            padding: vec![(0, 0), (1, 1)],
            mode_tag: 0,
            fill_bytes: fill.to_le_bytes().to_vec(),
        },
    )
    .expect("pad f16 dispatch");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    let expected = [0.0_f32, 1.0, 2.0, 3.0, 4.0, 0.0];
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        let got_f32 = g.to_f32();
        assert_eq!(got_f32, *e, "pad f16[{i}]: got {got_f32}, expected {e}");
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_pad_const_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // 1D pad: in=[3], padding=[(2,3)] → out=[8]
    let input = [1.0_f64, 2.0, 3.0];
    let fill = -1.0_f64;

    let in_storage = upload_f64(&backend, &input);
    let out_bytes = backend.alloc_bytes_handle(8 * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Pad, &[DType::F64, DType::F64], BackendId::Vulkan)[0]
            .kernel;
    let in_layout = Layout::contiguous(Shape::from_dims(&[3]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[8]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[in_layout, out_layout],
        &OpParams::Pad {
            in_shape: vec![3],
            out_shape: vec![8],
            padding: vec![(2, 3)],
            mode_tag: 0,
            fill_bytes: fill.to_le_bytes().to_vec(),
        },
    )
    .expect("pad f64 dispatch");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    let expected = [-1.0_f64, -1.0, 1.0, 2.0, 3.0, -1.0, -1.0, -1.0];
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "pad f64[{i}]: got {g}, expected {e}");
    }
}

// ---- SoftmaxLastDimBackward f32/f16/bf16/f64 (V.3.G.softmax-bwd, 2026-05-30) ----
//
// Reference: dx_j = y_j * (g_j - sum_i(y_i * g_i))

fn softmax_backward_ref(y: &[f32], g: &[f32], outer: usize, last: usize) -> Vec<f32> {
    let mut out = vec![0.0_f32; y.len()];
    for r in 0..outer {
        let off = r * last;
        let dot: f32 = (0..last).map(|i| y[off + i] * g[off + i]).sum();
        for i in 0..last {
            out[off + i] = y[off + i] * (g[off + i] - dot);
        }
    }
    out
}

#[test]
#[ignore]
fn vulkan_dispatch_softmax_last_dim_backward_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let n = outer * last;
    // y is a softmax output (positive, sum=1 per row). g is the upstream grad.
    let y: Vec<f32> = vec![0.1, 0.2, 0.3, 0.4, 0.4, 0.3, 0.2, 0.1];
    let g: Vec<f32> = vec![1.0, -1.0, 2.0, -2.0, 0.5, 1.5, -0.5, -1.5];

    let y_storage = upload_f32(&backend, &y);
    let g_storage = upload_f32(&backend, &g);
    let dx_bytes = backend.alloc_bytes_handle(n * 4).expect("alloc");
    let dx_storage = Storage::new(BackendStorage::Vulkan(dx_bytes), DType::F32);
    let y_arc = Arc::new(RwLock::new(y_storage));
    let g_arc = Arc::new(RwLock::new(g_storage));
    let dx_arc = Arc::new(RwLock::new(dx_storage));

    let kernel = table.lookup_alternatives(
        OpKind::SoftmaxLastDimBackward,
        &[DType::F32, DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    kernel(
        &[Arc::clone(&y_arc), Arc::clone(&g_arc)],
        &mut [Arc::clone(&dx_arc)],
        &[layout.clone(), layout.clone(), layout],
        &OpParams::SoftmaxLastDim {
            outer_count: outer,
            last_dim: last,
        },
    )
    .expect("softmax_backward f32 dispatch");

    let got = download_f32(&backend, &dx_arc.read().unwrap());
    let expected = softmax_backward_ref(&y, &g, outer, last);
    for (i, (a, b)) in got.iter().zip(expected.iter()).enumerate() {
        assert!(
            (a - b).abs() < 1e-5,
            "softmax_bwd f32[{i}]: got {a}, expected {b}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_softmax_last_dim_backward_f16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let n = outer * last;
    let y_f32: Vec<f32> = vec![0.1, 0.2, 0.3, 0.4, 0.4, 0.3, 0.2, 0.1];
    let g_f32: Vec<f32> = vec![1.0, -1.0, 2.0, -2.0, 0.5, 1.5, -0.5, -1.5];
    let y: Vec<half::f16> = y_f32.iter().map(|&x| half::f16::from_f32(x)).collect();
    let g: Vec<half::f16> = g_f32.iter().map(|&x| half::f16::from_f32(x)).collect();

    let y_storage = upload_f16(&backend, &y);
    let g_storage = upload_f16(&backend, &g);
    let dx_bytes = backend.alloc_bytes_handle(n * 2).expect("alloc");
    let dx_storage = Storage::new(BackendStorage::Vulkan(dx_bytes), DType::F16);
    let y_arc = Arc::new(RwLock::new(y_storage));
    let g_arc = Arc::new(RwLock::new(g_storage));
    let dx_arc = Arc::new(RwLock::new(dx_storage));

    let kernel = table.lookup_alternatives(
        OpKind::SoftmaxLastDimBackward,
        &[DType::F16, DType::F16, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    kernel(
        &[Arc::clone(&y_arc), Arc::clone(&g_arc)],
        &mut [Arc::clone(&dx_arc)],
        &[layout.clone(), layout.clone(), layout],
        &OpParams::SoftmaxLastDim {
            outer_count: outer,
            last_dim: last,
        },
    )
    .expect("softmax_backward f16 dispatch");

    let got = download_f16(&backend, &dx_arc.read().unwrap());
    let expected = softmax_backward_ref(&y_f32, &g_f32, outer, last);
    for (i, (a, b)) in got.iter().zip(expected.iter()).enumerate() {
        let got_f32 = a.to_f32();
        assert!(
            (got_f32 - b).abs() < 5e-3,
            "softmax_bwd f16[{i}]: got {got_f32}, expected {b}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_softmax_last_dim_backward_bf16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize; // MUST be even.
    let n = outer * last;
    let y_f32: Vec<f32> = vec![0.1, 0.2, 0.3, 0.4, 0.4, 0.3, 0.2, 0.1];
    let g_f32: Vec<f32> = vec![1.0, -1.0, 2.0, -2.0, 0.5, 1.5, -0.5, -1.5];
    let y: Vec<half::bf16> = y_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();
    let g: Vec<half::bf16> = g_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();

    let y_storage = upload_bf16(&backend, &y);
    let g_storage = upload_bf16(&backend, &g);
    let dx_bytes = backend.alloc_bytes_handle(n * 2).expect("alloc");
    let dx_storage = Storage::new(BackendStorage::Vulkan(dx_bytes), DType::BF16);
    let y_arc = Arc::new(RwLock::new(y_storage));
    let g_arc = Arc::new(RwLock::new(g_storage));
    let dx_arc = Arc::new(RwLock::new(dx_storage));

    let kernel = table.lookup_alternatives(
        OpKind::SoftmaxLastDimBackward,
        &[DType::BF16, DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    kernel(
        &[Arc::clone(&y_arc), Arc::clone(&g_arc)],
        &mut [Arc::clone(&dx_arc)],
        &[layout.clone(), layout.clone(), layout],
        &OpParams::SoftmaxLastDim {
            outer_count: outer,
            last_dim: last,
        },
    )
    .expect("softmax_backward bf16 dispatch");

    let got = download_bf16(&backend, &dx_arc.read().unwrap());
    let expected = softmax_backward_ref(&y_f32, &g_f32, outer, last);
    for (i, (a, b)) in got.iter().zip(expected.iter()).enumerate() {
        let got_f32 = a.to_f32();
        assert!(
            (got_f32 - b).abs() < 5e-2,
            "softmax_bwd bf16[{i}]: got {got_f32}, expected {b}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_softmax_last_dim_backward_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let n = outer * last;
    let y: Vec<f64> = vec![0.1, 0.2, 0.3, 0.4, 0.4, 0.3, 0.2, 0.1];
    let g: Vec<f64> = vec![1.0, -1.0, 2.0, -2.0, 0.5, 1.5, -0.5, -1.5];

    let y_storage = upload_f64(&backend, &y);
    let g_storage = upload_f64(&backend, &g);
    let dx_bytes = backend.alloc_bytes_handle(n * 8).expect("alloc");
    let dx_storage = Storage::new(BackendStorage::Vulkan(dx_bytes), DType::F64);
    let y_arc = Arc::new(RwLock::new(y_storage));
    let g_arc = Arc::new(RwLock::new(g_storage));
    let dx_arc = Arc::new(RwLock::new(dx_storage));

    let kernel = table.lookup_alternatives(
        OpKind::SoftmaxLastDimBackward,
        &[DType::F64, DType::F64, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    kernel(
        &[Arc::clone(&y_arc), Arc::clone(&g_arc)],
        &mut [Arc::clone(&dx_arc)],
        &[layout.clone(), layout.clone(), layout],
        &OpParams::SoftmaxLastDim {
            outer_count: outer,
            last_dim: last,
        },
    )
    .expect("softmax_backward f64 dispatch");

    let got = download_f64(&backend, &dx_arc.read().unwrap());
    // Pure arithmetic (no exp/log) so f64 is bit-accurate.
    for r in 0..outer {
        let off = r * last;
        let dot: f64 = (0..last).map(|i| y[off + i] * g[off + i]).sum();
        for i in 0..last {
            let expected = y[off + i] * (g[off + i] - dot);
            assert!(
                (got[off + i] - expected).abs() < 1e-12,
                "softmax_bwd f64[{}][{}]: got {}, expected {expected}",
                r,
                i,
                got[off + i]
            );
        }
    }
}

// ---- Concat bf16 (V.3.G.concat.bf16, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_concat_along_last_bf16_odd_a_dim() {
    // The case that motivated the InterlockedOr + zero-fill path:
    // concat along the last dim where a_dim is odd, so adjacent bf16
    // output positions come from DIFFERENT source buffers.
    //
    // a=[2,3] + b=[2,4] → out=[2,7] along dim=1.
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let a_f32: Vec<f32> = vec![1.0, 2.0, 3.0, 10.0, 20.0, 30.0];
    let b_f32: Vec<f32> = vec![4.0, 5.0, 6.0, 7.0, 40.0, 50.0, 60.0, 70.0];
    let a: Vec<half::bf16> = a_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();
    let b: Vec<half::bf16> = b_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();

    let a_storage = upload_bf16(&backend, &a);
    let b_storage = upload_bf16(&backend, &b);
    let out_bytes = backend.alloc_bytes_handle(2 * 7 * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::Concat,
        &[DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let a_layout = Layout::contiguous(Shape::from_dims(&[2, 3]));
    let b_layout = Layout::contiguous(Shape::from_dims(&[2, 4]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[2, 7]));
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &[a_layout, b_layout, out_layout],
        &OpParams::Concat {
            outer_count: 2,
            input_dim_sizes: vec![3, 4],
            inner_count: 1,
            axis: 1,
        },
    )
    .expect("concat bf16 dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    let expected_f32 = [
        1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0,
    ];
    for (i, (g, e)) in got.iter().zip(expected_f32.iter()).enumerate() {
        let got_f32 = g.to_f32();
        assert_eq!(got_f32, *e, "concat bf16[{i}]: got {got_f32}, expected {e}");
    }
}

// ---- Concat f16/f64 (V.3.G.concat, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_concat_along_dim_f16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // Shape: a=[2, 3] + b=[2, 4] → out=[2, 7] along dim=1.
    let a_f32: Vec<f32> = vec![1.0, 2.0, 3.0, 10.0, 20.0, 30.0];
    let b_f32: Vec<f32> = vec![4.0, 5.0, 6.0, 7.0, 40.0, 50.0, 60.0, 70.0];
    let a: Vec<half::f16> = a_f32.iter().map(|&x| half::f16::from_f32(x)).collect();
    let b: Vec<half::f16> = b_f32.iter().map(|&x| half::f16::from_f32(x)).collect();

    let a_storage = upload_f16(&backend, &a);
    let b_storage = upload_f16(&backend, &b);
    let out_bytes = backend.alloc_bytes_handle(2 * 7 * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Concat, &[DType::F16, DType::F16], BackendId::Vulkan)[0]
            .kernel;
    let a_layout = Layout::contiguous(Shape::from_dims(&[2, 3]));
    let b_layout = Layout::contiguous(Shape::from_dims(&[2, 4]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[2, 7]));
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &[a_layout, b_layout, out_layout],
        &OpParams::Concat {
            outer_count: 2,
            input_dim_sizes: vec![3, 4],
            inner_count: 1,
            axis: 1,
        },
    )
    .expect("concat f16 dispatch");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    let expected_f32 = [
        1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0,
    ];
    for (i, (g, e)) in got.iter().zip(expected_f32.iter()).enumerate() {
        assert_eq!(
            g.to_f32(),
            *e,
            "concat f16[{i}]: got {}, expected {e}",
            g.to_f32()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_concat_along_dim_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // N=3 chain test: a=[2,1] + b=[2,2] + c=[2,3] → out=[2,6] along dim=1.
    let a: Vec<f64> = vec![1.0, 10.0];
    let b: Vec<f64> = vec![2.0, 3.0, 20.0, 30.0];
    let c: Vec<f64> = vec![4.0, 5.0, 6.0, 40.0, 50.0, 60.0];

    let a_storage = upload_f64(&backend, &a);
    let b_storage = upload_f64(&backend, &b);
    let c_storage = upload_f64(&backend, &c);
    let out_bytes = backend.alloc_bytes_handle(2 * 6 * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let c_arc = Arc::new(RwLock::new(c_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Concat, &[DType::F64, DType::F64], BackendId::Vulkan)[0]
            .kernel;
    let a_layout = Layout::contiguous(Shape::from_dims(&[2, 1]));
    let b_layout = Layout::contiguous(Shape::from_dims(&[2, 2]));
    let c_layout = Layout::contiguous(Shape::from_dims(&[2, 3]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[2, 6]));
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc), Arc::clone(&c_arc)],
        &mut [Arc::clone(&out_arc)],
        &[a_layout, b_layout, c_layout, out_layout],
        &OpParams::Concat {
            outer_count: 2,
            input_dim_sizes: vec![1, 2, 3],
            inner_count: 1,
            axis: 1,
        },
    )
    .expect("concat f64 dispatch (N=3 chain)");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    let expected = [
        1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0, 10.0, 20.0, 30.0, 40.0, 50.0, 60.0,
    ];
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "concat f64[{i}]: got {g}, expected {e}");
    }
}

// ---- IndexSelect f16/bf16/f64 (V.3.G.index_select, 2026-05-30) ----

fn index_select_test_shape() -> (usize, usize, usize, Vec<u32>) {
    // outer=1, axis_in=4 source rows of inner=8 elements each;
    // pick indices [3, 0, 2] → 3 output rows.
    let outer = 1usize;
    let axis_in = 4usize;
    let inner = 8usize;
    let ids = vec![3u32, 0u32, 2u32];
    (outer, axis_in, inner, ids)
}

#[test]
#[ignore]
fn vulkan_dispatch_index_select_f16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let (outer, axis_in, inner, ids) = index_select_test_shape();
    let n_in = outer * axis_in * inner;
    let n_out = outer * ids.len() * inner;

    // Source: rows are 0..7, 100..107, 200..207, 300..307.
    let src_f32: Vec<f32> = (0..axis_in)
        .flat_map(|r| (0..inner).map(move |c| (r * 100 + c) as f32))
        .collect();
    let src: Vec<half::f16> = src_f32.iter().map(|&x| half::f16::from_f32(x)).collect();

    let src_storage = upload_f16(&backend, &src);
    let ids_bytes: &[u8] = bytemuck::cast_slice(&ids);
    let ids_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(ids_bytes).expect("ids upload")),
        DType::U32,
    );
    let out_bytes = backend.alloc_bytes_handle(n_out * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let src_arc = Arc::new(RwLock::new(src_storage));
    let ids_arc = Arc::new(RwLock::new(ids_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let _ = n_in;
    let kernel = table.lookup_alternatives(
        OpKind::IndexSelect,
        &[DType::F16, DType::U32, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let src_layout = Layout::contiguous(Shape::from_dims(&[outer, axis_in, inner]));
    let ids_layout = Layout::contiguous(Shape::from_dims(&[ids.len()]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[outer, ids.len(), inner]));
    kernel(
        &[Arc::clone(&src_arc), Arc::clone(&ids_arc)],
        &mut [Arc::clone(&out_arc)],
        &[src_layout, ids_layout, out_layout],
        &OpParams::IndexSelect {
            outer_count: outer,
            source_dim_size: axis_in,
            n_indices: ids.len(),
            inner_count: inner,
        },
    )
    .expect("index_select f16 dispatch");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    // Expected: pick rows 3, 0, 2 in that order.
    let expected_f32: Vec<f32> = ids
        .iter()
        .flat_map(|&id| (0..inner).map(move |c| (id as usize * 100 + c) as f32))
        .collect();
    for (i, (g, e)) in got.iter().zip(expected_f32.iter()).enumerate() {
        assert_eq!(
            g.to_f32(),
            *e,
            "index_select_f16[{i}]: got {}, expected {e}",
            g.to_f32()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_index_select_bf16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let (outer, axis_in, inner, ids) = index_select_test_shape();
    // inner=8 satisfies the inner%2==0 pair-thread constraint.
    let n_out = outer * ids.len() * inner;

    let src_f32: Vec<f32> = (0..axis_in)
        .flat_map(|r| (0..inner).map(move |c| (r * 100 + c) as f32))
        .collect();
    let src: Vec<half::bf16> = src_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();

    let src_storage = upload_bf16(&backend, &src);
    let ids_bytes: &[u8] = bytemuck::cast_slice(&ids);
    let ids_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(ids_bytes).expect("ids upload")),
        DType::U32,
    );
    let out_bytes = backend.alloc_bytes_handle(n_out * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let src_arc = Arc::new(RwLock::new(src_storage));
    let ids_arc = Arc::new(RwLock::new(ids_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::IndexSelect,
        &[DType::BF16, DType::U32, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let src_layout = Layout::contiguous(Shape::from_dims(&[outer, axis_in, inner]));
    let ids_layout = Layout::contiguous(Shape::from_dims(&[ids.len()]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[outer, ids.len(), inner]));
    kernel(
        &[Arc::clone(&src_arc), Arc::clone(&ids_arc)],
        &mut [Arc::clone(&out_arc)],
        &[src_layout, ids_layout, out_layout],
        &OpParams::IndexSelect {
            outer_count: outer,
            source_dim_size: axis_in,
            n_indices: ids.len(),
            inner_count: inner,
        },
    )
    .expect("index_select bf16 dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    let expected_f32: Vec<f32> = ids
        .iter()
        .flat_map(|&id| (0..inner).map(move |c| (id as usize * 100 + c) as f32))
        .collect();
    for (i, (g, e)) in got.iter().zip(expected_f32.iter()).enumerate() {
        // bf16 round-trip is exact for integer values <= 256; our test
        // values go up to ~307 which loses ~1 ULP at most. Use loose
        // tolerance.
        let got_f32 = g.to_f32();
        assert!(
            (got_f32 - e).abs() < 5.0,
            "index_select_bf16[{i}]: got {got_f32}, expected {e}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_index_select_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let (outer, axis_in, inner, ids) = index_select_test_shape();
    let n_out = outer * ids.len() * inner;

    let src: Vec<f64> = (0..axis_in)
        .flat_map(|r| (0..inner).map(move |c| (r * 100 + c) as f64))
        .collect();

    let src_storage = upload_f64(&backend, &src);
    let ids_bytes: &[u8] = bytemuck::cast_slice(&ids);
    let ids_storage = Storage::new(
        BackendStorage::Vulkan(backend.upload_bytes_handle(ids_bytes).expect("ids upload")),
        DType::U32,
    );
    let out_bytes = backend.alloc_bytes_handle(n_out * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let src_arc = Arc::new(RwLock::new(src_storage));
    let ids_arc = Arc::new(RwLock::new(ids_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::IndexSelect,
        &[DType::F64, DType::U32, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let src_layout = Layout::contiguous(Shape::from_dims(&[outer, axis_in, inner]));
    let ids_layout = Layout::contiguous(Shape::from_dims(&[ids.len()]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[outer, ids.len(), inner]));
    kernel(
        &[Arc::clone(&src_arc), Arc::clone(&ids_arc)],
        &mut [Arc::clone(&out_arc)],
        &[src_layout, ids_layout, out_layout],
        &OpParams::IndexSelect {
            outer_count: outer,
            source_dim_size: axis_in,
            n_indices: ids.len(),
            inner_count: inner,
        },
    )
    .expect("index_select f64 dispatch");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    let expected: Vec<f64> = ids
        .iter()
        .flat_map(|&id| (0..inner).map(move |c| (id as usize * 100 + c) as f64))
        .collect();
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert_eq!(*g, *e, "index_select_f64[{i}]: got {g}, expected {e}");
    }
}

// ---- RoPE f16/f64 (V.3.G.rope, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_rope_f16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // Shape: [1, 2, 4]  (outer=1, seq=2, head_dim=4, h=2)
    let outer = 1usize;
    let seq = 2usize;
    let hd = 4usize;
    let n_x = outer * seq * hd;
    let n_table = seq * hd;

    // Identity rotation test: cos=1, sin=0 everywhere → output == x.
    let host_x_f32: Vec<f32> = (0..n_x).map(|i| i as f32 + 1.0).collect();
    let host_cos_f32: Vec<f32> = vec![1.0; n_table];
    let host_sin_f32: Vec<f32> = vec![0.0; n_table];
    let host_x: Vec<half::f16> = host_x_f32.iter().map(|&x| half::f16::from_f32(x)).collect();
    let host_cos: Vec<half::f16> = host_cos_f32
        .iter()
        .map(|&x| half::f16::from_f32(x))
        .collect();
    let host_sin: Vec<half::f16> = host_sin_f32
        .iter()
        .map(|&x| half::f16::from_f32(x))
        .collect();

    let x_storage = upload_f16(&backend, &host_x);
    let cos_storage = upload_f16(&backend, &host_cos);
    let sin_storage = upload_f16(&backend, &host_sin);
    let out_bytes = backend.alloc_bytes_handle(n_x * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let x_arc = Arc::new(RwLock::new(x_storage));
    let cos_arc = Arc::new(RwLock::new(cos_storage));
    let sin_arc = Arc::new(RwLock::new(sin_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::Rope,
        &[DType::F16, DType::F16, DType::F16, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let x_layout = Layout::contiguous(Shape::from_dims(&[outer, seq, hd]));
    let table_layout = Layout::contiguous(Shape::from_dims(&[seq, hd]));
    let out_layout = x_layout.clone();
    kernel(
        &[
            Arc::clone(&x_arc),
            Arc::clone(&cos_arc),
            Arc::clone(&sin_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &[x_layout, table_layout.clone(), table_layout, out_layout],
        &OpParams::Rope {
            outer_count: outer,
            seq,
            head_dim: hd,
        },
    )
    .expect("rope f16 dispatch");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    for (i, (g, h)) in got.iter().zip(host_x_f32.iter()).enumerate() {
        let got_f32 = g.to_f32();
        assert!(
            (got_f32 - h).abs() < 5e-3,
            "rope-f16[{i}] (identity): got {got_f32}, expected {h}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_rope_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 1usize;
    let seq = 2usize;
    let hd = 4usize;
    let n_x = outer * seq * hd;
    let n_table = seq * hd;

    // Real rotation: cos/sin = ±sqrt(0.5) → 45° rotation. h=2; rotates
    // (x[i], x[i+h]) by 45° giving out[i] = (x[i]-x[i+h])*sqrt(0.5),
    // out[i+h] = (x[i]+x[i+h])*sqrt(0.5) (using c0=c1=s1=sqrt(0.5),
    // s0=sqrt(0.5)). Tests that the kernel is actually computing both
    // outputs, not just passing through.
    let host_x: Vec<f64> = vec![
        1.0, 2.0, 3.0, 4.0, // s=0: pairs (1,3) and (2,4)
        5.0, 6.0, 7.0, 8.0, // s=1: pairs (5,7) and (6,8)
    ];
    let q = std::f64::consts::FRAC_1_SQRT_2;
    let host_cos: Vec<f64> = vec![q; n_table];
    let host_sin: Vec<f64> = vec![q; n_table];

    let x_storage = upload_f64(&backend, &host_x);
    let cos_storage = upload_f64(&backend, &host_cos);
    let sin_storage = upload_f64(&backend, &host_sin);
    let out_bytes = backend.alloc_bytes_handle(n_x * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let x_arc = Arc::new(RwLock::new(x_storage));
    let cos_arc = Arc::new(RwLock::new(cos_storage));
    let sin_arc = Arc::new(RwLock::new(sin_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::Rope,
        &[DType::F64, DType::F64, DType::F64, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let x_layout = Layout::contiguous(Shape::from_dims(&[outer, seq, hd]));
    let table_layout = Layout::contiguous(Shape::from_dims(&[seq, hd]));
    let out_layout = x_layout.clone();
    kernel(
        &[
            Arc::clone(&x_arc),
            Arc::clone(&cos_arc),
            Arc::clone(&sin_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &[x_layout, table_layout.clone(), table_layout, out_layout],
        &OpParams::Rope {
            outer_count: outer,
            seq,
            head_dim: hd,
        },
    )
    .expect("rope f64 dispatch");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    // Reference: for each (s, i in 0..h): out[i] = x[i]*c0 - x[i+h]*s0;
    //                                    out[i+h] = x[i+h]*c1 + x[i]*s1.
    let h = hd / 2;
    for s in 0..seq {
        let row = s * hd;
        for i in 0..h {
            let x0 = host_x[row + i];
            let x1 = host_x[row + i + h];
            let expected_lo = x0 * q - x1 * q;
            let expected_hi = x1 * q + x0 * q;
            assert!(
                (got[row + i] - expected_lo).abs() < 1e-12,
                "rope-f64 s={s} i={i}: got {}, expected {expected_lo}",
                got[row + i]
            );
            assert!(
                (got[row + i + h] - expected_hi).abs() < 1e-12,
                "rope-f64 s={s} i+h={}: got {}, expected {expected_hi}",
                i + h,
                got[row + i + h]
            );
        }
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_rope_bf16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // head_dim must be a multiple of 4 for the pair-thread scheme.
    let outer = 1usize;
    let seq = 2usize;
    let hd = 8usize; // = 4k, h = 4, i in {0,1,2,3} → 2 pair-threads per row
    let n_x = outer * seq * hd;
    let n_table = seq * hd;

    // Real 45° rotation (cos=sin=sqrt(0.5)) — same shape as the f64
    // test but at bf16 precision.
    let host_x_f32: Vec<f32> = vec![
        // s=0
        1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, // s=1
        2.0, 4.0, 6.0, 8.0, 1.0, 3.0, 5.0, 7.0,
    ];
    let q = std::f32::consts::FRAC_1_SQRT_2;
    let host_cos_f32: Vec<f32> = vec![q; n_table];
    let host_sin_f32: Vec<f32> = vec![q; n_table];
    let host_x: Vec<half::bf16> = host_x_f32
        .iter()
        .map(|&x| half::bf16::from_f32(x))
        .collect();
    let host_cos: Vec<half::bf16> = host_cos_f32
        .iter()
        .map(|&x| half::bf16::from_f32(x))
        .collect();
    let host_sin: Vec<half::bf16> = host_sin_f32
        .iter()
        .map(|&x| half::bf16::from_f32(x))
        .collect();

    let x_storage = upload_bf16(&backend, &host_x);
    let cos_storage = upload_bf16(&backend, &host_cos);
    let sin_storage = upload_bf16(&backend, &host_sin);
    let out_bytes = backend.alloc_bytes_handle(n_x * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let x_arc = Arc::new(RwLock::new(x_storage));
    let cos_arc = Arc::new(RwLock::new(cos_storage));
    let sin_arc = Arc::new(RwLock::new(sin_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::Rope,
        &[DType::BF16, DType::BF16, DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let x_layout = Layout::contiguous(Shape::from_dims(&[outer, seq, hd]));
    let table_layout = Layout::contiguous(Shape::from_dims(&[seq, hd]));
    let out_layout = x_layout.clone();
    kernel(
        &[
            Arc::clone(&x_arc),
            Arc::clone(&cos_arc),
            Arc::clone(&sin_arc),
        ],
        &mut [Arc::clone(&out_arc)],
        &[x_layout, table_layout.clone(), table_layout, out_layout],
        &OpParams::Rope {
            outer_count: outer,
            seq,
            head_dim: hd,
        },
    )
    .expect("rope bf16 dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    // Reference computed in f32 (matches kernel's f32 internal math).
    let h = hd / 2;
    for s in 0..seq {
        let row = s * hd;
        for i in 0..h {
            let x0 = host_x_f32[row + i];
            let x1 = host_x_f32[row + i + h];
            let expected_lo = x0 * q - x1 * q;
            let expected_hi = x1 * q + x0 * q;
            let got_lo = got[row + i].to_f32();
            let got_hi = got[row + i + h].to_f32();
            // bf16 ~7-bit mantissa → ~1% relative.
            let tol_lo = expected_lo.abs() * 0.01 + 5e-2;
            let tol_hi = expected_hi.abs() * 0.01 + 5e-2;
            assert!(
                (got_lo - expected_lo).abs() < tol_lo,
                "rope-bf16 s={s} i={i}: got {got_lo}, expected {expected_lo}"
            );
            assert!(
                (got_hi - expected_hi).abs() < tol_hi,
                "rope-bf16 s={s} i+h={}: got {got_hi}, expected {expected_hi}",
                i + h
            );
        }
    }
}

// ---- Cast f32 ↔ f64 (V.3.G.cast, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_cast_f32_to_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host: Vec<f32> = vec![1.0, 2.5, -3.125, 0.0, f32::MIN_POSITIVE, 1.0e20];
    let n = host.len();

    let in_storage = upload_f32(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(n * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Cast, &[DType::F32, DType::F64], BackendId::Vulkan)[0]
            .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::None,
    )
    .expect("cast f32→f64");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    for (i, (g, h)) in got.iter().zip(host.iter()).enumerate() {
        // Widening is exact — f32 representable in f64 bit-for-bit.
        assert_eq!(
            *g, *h as f64,
            "cast_f32_to_f64[{i}]: got {g}, expected {}",
            *h as f64
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_cast_f64_to_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host: Vec<f64> = vec![
        1.0,
        2.5,
        -3.125,
        0.0,
        1.0e20,               // representable in f32
        1.0e40,               // overflows to +Inf in f32
        1.0e-50,              // underflows in f32
        std::f64::consts::PI, // round-to-nearest at narrowing
    ];
    let n = host.len();

    let in_storage = upload_f64(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel =
        table.lookup_alternatives(OpKind::Cast, &[DType::F64, DType::F32], BackendId::Vulkan)[0]
            .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::None,
    )
    .expect("cast f64→f32");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    for (i, (g, h)) in got.iter().zip(host.iter()).enumerate() {
        let expected = *h as f32;
        if expected.is_infinite() {
            assert!(
                g.is_infinite() && g.signum() == expected.signum(),
                "cast_f64_to_f32[{i}]: got {g}, expected ±inf"
            );
        } else {
            // Narrowing rounds to nearest-even — should match Rust's `as f32`.
            assert_eq!(
                *g, expected,
                "cast_f64_to_f32[{i}]: got {g}, expected {expected}"
            );
        }
    }
}

// ---- Abs / Sign / Recip (V.3.G.unary, 2026-05-30) ----

#[test]
#[ignore]
fn vulkan_dispatch_unary_abs_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_unary_f32(&backend, OpKind::AbsElementwise, &[1.0, -2.0, 0.0, -3.5]);
    assert_eq!(got, vec![1.0, 2.0, 0.0, 3.5]);
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_sign_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_unary_f32(
        &backend,
        OpKind::SignElementwise,
        &[5.0, -3.0, 0.0, -0.0, 7.5],
    );
    assert_eq!(got, vec![1.0, -1.0, 0.0, 0.0, 1.0]);
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_recip_f32() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_unary_f32(&backend, OpKind::RecipElementwise, &[1.0, 2.0, 4.0, -8.0]);
    assert_close(&got, &[1.0, 0.5, 0.25, -0.125], 1e-6, 1e-6);
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_sign_f16() {
    // Exercises the explicit float16_t(sign(x)) cast added because
    // Slang's sign() on half returns int.
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let host: Vec<half::f16> = [5.0_f32, -3.0, 0.0, 7.5]
        .iter()
        .map(|&x| half::f16::from_f32(x))
        .collect();
    let got = run_unary_f16(&backend, OpKind::SignElementwise, &host);
    let got_f32: Vec<f32> = got.iter().map(|x| x.to_f32()).collect();
    assert_eq!(got_f32, vec![1.0, -1.0, 0.0, 1.0]);
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_recip_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let got = run_unary_f64(
        &backend,
        OpKind::RecipElementwise,
        &[1.0_f64, 2.0, 4.0, -8.0, 16.0],
    );
    let expected = [1.0_f64, 0.5, 0.25, -0.125, 0.0625];
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert!(
            (g - e).abs() < 1e-15,
            "recip-f64[{i}]: got {g}, expected {e}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_unary_abs_sign_recip_bf16() {
    use half::bf16;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host: Vec<bf16> = [1.0_f32, -2.0, 0.0, -3.5, 4.0, 0.25]
        .iter()
        .map(|&x| bf16::from_f32(x))
        .collect();
    let n = host.len();

    for (op, label, expected) in [
        (
            OpKind::AbsElementwise,
            "abs",
            vec![1.0_f32, 2.0, 0.0, 3.5, 4.0, 0.25],
        ),
        (
            OpKind::SignElementwise,
            "sign",
            vec![1.0_f32, -1.0, 0.0, -1.0, 1.0, 1.0],
        ),
        (
            OpKind::RecipElementwise,
            "recip",
            vec![1.0_f32, -0.5, f32::INFINITY, -0.2857143, 0.25, 4.0],
        ),
    ] {
        let kernel =
            table.lookup_alternatives(op, &[DType::BF16, DType::BF16], BackendId::Vulkan)[0].kernel;
        let in_storage = upload_bf16(&backend, &host);
        let out_bytes = backend.alloc_bytes_handle(n * 2).expect("alloc");
        let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
        let in_arc = Arc::new(RwLock::new(in_storage));
        let out_arc = Arc::new(RwLock::new(out_storage));
        let layout = Layout::contiguous(Shape::from_dims(&[n]));
        kernel(
            &[Arc::clone(&in_arc)],
            &mut [Arc::clone(&out_arc)],
            &[layout.clone(), layout],
            &OpParams::None,
        )
        .expect("bf16 unary");
        let got = download_bf16(&backend, &out_arc.read().unwrap());
        for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
            let got_f32 = g.to_f32();
            if e.is_infinite() {
                assert!(
                    got_f32.is_infinite() && got_f32.signum() == e.signum(),
                    "{label}-bf16[{i}]: got {got_f32}, expected ±inf"
                );
            } else {
                // bf16 ~7-bit mantissa → ~1% relative.
                let tol = e.abs() * 0.01 + 5e-3;
                assert!(
                    (got_f32 - e).abs() < tol,
                    "{label}-bf16[{i}]: got {got_f32}, expected {e}"
                );
            }
        }
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_reduce_sum_full_f16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let n = 8usize;
    let host_f32: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 0.5, 1.5, 2.5, 3.5];
    let host: Vec<half::f16> = host_f32.iter().map(|&x| half::f16::from_f32(x)).collect();
    let expected: f32 = host_f32.iter().sum(); // = 18.0

    let in_storage = upload_f16(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::SumReduce,
        &[DType::F16, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[1]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout, out_layout],
        &OpParams::Reduce {
            dims: vec![],
            keepdim: false,
        },
    )
    .expect("sum-full-reduce f16 dispatch");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    let got_f32 = got[0].to_f32();
    assert!(
        (got_f32 - expected).abs() < 5e-3,
        "sum-full-f16: got {got_f32}, expected {expected}"
    );
}

#[test]
#[ignore]
fn vulkan_dispatch_reduce_min_full_bf16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let n = 8usize; // MUST be even (lane-pair input).
    let host_f32: Vec<f32> = vec![3.0, -1.0, 4.0, 1.5, -5.0, 9.0, 2.0, 6.0];
    let host: Vec<half::bf16> = host_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();
    let expected = host_f32.iter().cloned().fold(f32::INFINITY, f32::min); // = -5.0

    let in_storage = upload_bf16(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MinReduce,
        &[DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[1]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout, out_layout],
        &OpParams::Reduce {
            dims: vec![],
            keepdim: false,
        },
    )
    .expect("min-full-reduce bf16 dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    let got_f32 = got[0].to_f32();
    assert!(
        (got_f32 - expected).abs() < 5e-2,
        "min-full-bf16: got {got_f32}, expected {expected}"
    );
}

#[test]
#[ignore]
fn vulkan_dispatch_reduce_mean_full_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let n = 8usize;
    let host: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
    let expected: f64 = host.iter().sum::<f64>() / n as f64; // = 4.5

    let in_storage = upload_f64(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MeanReduce,
        &[DType::F64, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[1]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout, out_layout],
        &OpParams::Reduce {
            dims: vec![],
            keepdim: false,
        },
    )
    .expect("mean-full-reduce f64 dispatch");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    assert!(
        (got[0] - expected).abs() < 1e-12,
        "mean-full-f64: got {}, expected {}",
        got[0],
        expected
    );
}

#[test]
#[ignore]
fn vulkan_dispatch_reduce_sum_last_dim_f16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let n = outer * last;
    let host_f32: Vec<f32> = vec![
        1.0, 2.0, 3.0, 4.0, // row sum = 10
        0.5, 1.5, 2.5, 3.5, // row sum = 8
    ];
    let host: Vec<half::f16> = host_f32.iter().map(|&x| half::f16::from_f32(x)).collect();

    let in_storage = upload_f16(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(outer * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::SumReduce,
        &[DType::F16, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[outer]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout, out_layout],
        &OpParams::Reduce {
            dims: vec![1],
            keepdim: false,
        },
    )
    .expect("sum-reduce f16 dispatch");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    let _ = n; // silence unused
    let expected = [10.0_f32, 8.0_f32];
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        let got_f32 = g.to_f32();
        assert!(
            (got_f32 - e).abs() < 5e-3,
            "sum-reduce-f16 row {i}: got {got_f32}, expected {e}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_reduce_max_last_dim_bf16_odd_rows() {
    // 5 rows exercises the InterlockedOr edge case: the last row
    // lives in the high half of the 3rd u32 word (padded region of
    // the output buffer). The wrapper's zero-fill keeps that word's
    // low half zero so the OR is a clean half-word write.
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 5usize;
    let last = 4usize; // MUST be even.
    let host_f32: Vec<f32> = vec![
        1.0, 2.0, 3.0, 4.0, // row max = 4
        -1.0, 5.0, 2.0, 0.0, // row max = 5
        0.5, 1.5, 2.5, 3.5, // row max = 3.5
        -5.0, -3.0, -2.0, -1.0, // row max = -1
        7.0, 6.0, 8.0, 4.5, // row max = 8
    ];
    let host: Vec<half::bf16> = host_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();

    let in_storage = upload_bf16(&backend, &host);
    // n_rows*2 = 10 bytes; alloc_bytes_handle rounds to 12 (u32-align).
    let out_bytes = backend.alloc_bytes_handle(outer * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MaxReduce,
        &[DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[outer]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout, out_layout],
        &OpParams::Reduce {
            dims: vec![1],
            keepdim: false,
        },
    )
    .expect("max-reduce bf16 dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    let expected = [4.0_f32, 5.0, 3.5, -1.0, 8.0];
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        let got_f32 = g.to_f32();
        assert!(
            (got_f32 - e).abs() < 5e-2,
            "max-reduce-bf16 row {i}: got {got_f32}, expected {e}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_reduce_mean_last_dim_f64() {
    // mean exercises the op_id=3 path (subgroup sum + final divide).
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let host: Vec<f64> = vec![
        1.0, 2.0, 3.0, 4.0, // mean = 2.5
        0.5, 1.5, 2.5, 3.5, // mean = 2.0
    ];

    let in_storage = upload_f64(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(outer * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::MeanReduce,
        &[DType::F64, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[outer]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout, out_layout],
        &OpParams::Reduce {
            dims: vec![1],
            keepdim: false,
        },
    )
    .expect("mean-reduce f64 dispatch");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    let expected = [2.5_f64, 2.0];
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert!(
            (g - e).abs() < 1e-12,
            "mean-reduce-f64 row {i}: got {g}, expected {e}"
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_softmax_last_dim_f16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let n = outer * last;
    let host_f32: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 0.5, 1.5, 2.5, 3.5];
    let host: Vec<half::f16> = host_f32.iter().map(|&x| half::f16::from_f32(x)).collect();

    let in_storage = upload_f16(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(n * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::SoftmaxLastDim,
        &[DType::F16, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::SoftmaxLastDim {
            outer_count: outer,
            last_dim: last,
        },
    )
    .expect("softmax f16 dispatch");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    for row in 0..outer {
        let xs = &host_f32[row * last..(row + 1) * last];
        let ys = &got[row * last..(row + 1) * last];
        let max = xs.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let exps: Vec<f32> = xs.iter().map(|x| (x - max).exp()).collect();
        let sum: f32 = exps.iter().sum();
        for (i, (e, y)) in exps.iter().zip(ys.iter()).enumerate() {
            let expected = e / sum;
            let got_f32 = y.to_f32();
            assert!(
                (got_f32 - expected).abs() < 5e-3,
                "softmax-f16 row {row} col {i}: got {got_f32}, expected {expected}"
            );
        }
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_softmax_last_dim_bf16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize; // MUST be even.
    let n = outer * last;
    let host_f32: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 0.5, 1.5, 2.5, 3.5];
    let host: Vec<half::bf16> = host_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();

    let in_storage = upload_bf16(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(n * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::SoftmaxLastDim,
        &[DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::SoftmaxLastDim {
            outer_count: outer,
            last_dim: last,
        },
    )
    .expect("softmax bf16 dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    for row in 0..outer {
        let xs = &host_f32[row * last..(row + 1) * last];
        let ys = &got[row * last..(row + 1) * last];
        let max = xs.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let exps: Vec<f32> = xs.iter().map(|x| (x - max).exp()).collect();
        let sum: f32 = exps.iter().sum();
        for (i, (e, y)) in exps.iter().zip(ys.iter()).enumerate() {
            let expected = e / sum;
            let got_f32 = y.to_f32();
            assert!(
                (got_f32 - expected).abs() < 5e-2,
                "softmax-bf16 row {row} col {i}: got {got_f32}, expected {expected}"
            );
        }
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_softmax_last_dim_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let n = outer * last;
    let host: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0, 0.5, 1.5, 2.5, 3.5];

    let in_storage = upload_f64(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(n * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::SoftmaxLastDim,
        &[DType::F64, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::SoftmaxLastDim {
            outer_count: outer,
            last_dim: last,
        },
    )
    .expect("softmax f64 dispatch");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    // 1e-7 (not 1e-12): empirically the RTX 4070's GLSL.std.450 Exp
    // on f64 lowers through a ~f32-accuracy path before re-widening
    // (observed ~1.2e-9 absolute drift on inputs around 0.03). The
    // kernel structure is bit-stable; only the transcendental is
    // implementation-defined.
    for row in 0..outer {
        let xs = &host[row * last..(row + 1) * last];
        let ys = &got[row * last..(row + 1) * last];
        let max = xs.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
        let exps: Vec<f64> = xs.iter().map(|x| (x - max).exp()).collect();
        let sum: f64 = exps.iter().sum();
        for (i, (e, y)) in exps.iter().zip(ys.iter()).enumerate() {
            let expected = e / sum;
            assert!(
                (y - expected).abs() < 1e-7,
                "softmax-f64 row {row} col {i}: got {y}, expected {expected}"
            );
        }
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_rms_norm_last_dim_f16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let n = outer * last;
    let host_f32: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 2.0, 4.0, 6.0, 8.0];
    let host: Vec<half::f16> = host_f32.iter().map(|&x| half::f16::from_f32(x)).collect();

    let in_storage = upload_f16(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(n * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::RmsNormLastDim,
        &[DType::F16, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    let eps = 1e-6f64;
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::NormLastDim {
            outer_count: outer,
            last_dim: last,
            eps,
        },
    )
    .expect("rmsnorm f16 dispatch");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    // Reference in f32 (matches kernel's mixed-precision pattern), then
    // round to f16 for comparison; tolerance reflects f16's ~3-decimal
    // mantissa.
    for row in 0..outer {
        let xs = &host_f32[row * last..(row + 1) * last];
        let ys = &got[row * last..(row + 1) * last];
        let mean_sq: f32 = xs.iter().map(|x| x * x).sum::<f32>() / last as f32;
        let scale = (mean_sq + eps as f32).sqrt();
        for (i, (x, y)) in xs.iter().zip(ys.iter()).enumerate() {
            let expected = x / scale;
            let got_f32 = y.to_f32();
            assert!(
                (got_f32 - expected).abs() < 5e-3,
                "rmsnorm-f16 row {row} col {i}: got {got_f32}, expected {expected}"
            );
        }
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_rms_norm_last_dim_bf16() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize; // MUST be even — lane-pair packing.
    let n = outer * last;
    let host_f32: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 2.0, 4.0, 6.0, 8.0];
    let host: Vec<half::bf16> = host_f32.iter().map(|&x| half::bf16::from_f32(x)).collect();

    let in_storage = upload_bf16(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(n * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::RmsNormLastDim,
        &[DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    let eps = 1e-6f64;
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::NormLastDim {
            outer_count: outer,
            last_dim: last,
            eps,
        },
    )
    .expect("rmsnorm bf16 dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    // bf16 has only 8 mantissa bits — wider tolerance than f16 despite
    // the same width, because the exponent range is preserved at the
    // cost of mantissa precision.
    for row in 0..outer {
        let xs = &host_f32[row * last..(row + 1) * last];
        let ys = &got[row * last..(row + 1) * last];
        let mean_sq: f32 = xs.iter().map(|x| x * x).sum::<f32>() / last as f32;
        let scale = (mean_sq + eps as f32).sqrt();
        for (i, (x, y)) in xs.iter().zip(ys.iter()).enumerate() {
            let expected = x / scale;
            let got_f32 = y.to_f32();
            assert!(
                (got_f32 - expected).abs() < 5e-2,
                "rmsnorm-bf16 row {row} col {i}: got {got_f32}, expected {expected}"
            );
        }
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_rms_norm_last_dim_f64() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let outer = 2usize;
    let last = 4usize;
    let n = outer * last;
    let host: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0, 2.0, 4.0, 6.0, 8.0];

    let in_storage = upload_f64(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(n * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::RmsNormLastDim,
        &[DType::F64, DType::F64],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let layout = Layout::contiguous(Shape::from_dims(&[outer, last]));
    let eps = 1e-12f64;
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::NormLastDim {
            outer_count: outer,
            last_dim: last,
            eps,
        },
    )
    .expect("rmsnorm f64 dispatch");

    let got = download_f64(&backend, &out_arc.read().unwrap());
    // Native f64 throughout — tight tolerance, verifies subgroup sum
    // and GLSL.std.450 Sqrt both work on doubles under shaderFloat64.
    for row in 0..outer {
        let xs = &host[row * last..(row + 1) * last];
        let ys = &got[row * last..(row + 1) * last];
        let mean_sq: f64 = xs.iter().map(|x| x * x).sum::<f64>() / last as f64;
        let scale = (mean_sq + eps).sqrt();
        for (i, (x, y)) in xs.iter().zip(ys.iter()).enumerate() {
            let expected = x / scale;
            assert!(
                (y - expected).abs() < 1e-10,
                "rmsnorm-f64 row {row} col {i}: got {y}, expected {expected}"
            );
        }
    }
}

// ===========================================================================
// V.3 fan-out: triu/tril, flip, roll, bf16 unary+binary, F8E4M3 cast,
// write_slice b1.
// ===========================================================================

fn upload_bf16(backend: &Arc<VulkanBackend>, host: &[half::bf16]) -> Storage {
    let bytes: &[u8] = bytemuck::cast_slice(host);
    let vk_bytes = backend
        .upload_bytes_handle(bytes)
        .expect("vulkan upload bf16");
    Storage::new(BackendStorage::Vulkan(vk_bytes), DType::BF16)
}

fn download_bf16(backend: &Arc<VulkanBackend>, s: &Storage) -> Vec<half::bf16> {
    let bytes = match &s.inner {
        BackendStorage::Vulkan(v) => backend.download_bytes(v).expect("d2h"),
        _ => panic!("not on Vulkan"),
    };
    bytemuck::cast_slice::<u8, half::bf16>(&bytes).to_vec()
}

fn upload_raw(backend: &Arc<VulkanBackend>, bytes: &[u8], dtype: DType) -> Storage {
    let vk_bytes = backend
        .upload_bytes_handle(bytes)
        .expect("vulkan upload raw");
    Storage::new(BackendStorage::Vulkan(vk_bytes), dtype)
}

fn download_raw(backend: &Arc<VulkanBackend>, s: &Storage) -> Vec<u8> {
    match &s.inner {
        BackendStorage::Vulkan(v) => backend.download_bytes(v).expect("d2h"),
        _ => panic!("not on Vulkan"),
    }
}

// ----- Triu / Tril ---------------------------------------------------------

#[test]
#[ignore]
fn vulkan_dispatch_triu_f32() {
    use fuel_dispatch::kernel::OpParams;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let rows = 4usize;
    let cols = 4usize;
    let mat = vec![1.0_f32; rows * cols];

    for diagonal in [0i64, 1, -1] {
        let kernel =
            table.lookup_alternatives(OpKind::Triu, &[DType::F32, DType::F32], BackendId::Vulkan)
                [0]
            .kernel;
        let in_storage = upload_f32(&backend, &mat);
        let out_bytes = backend.alloc_bytes_handle(rows * cols * 4).expect("alloc");
        let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
        let in_arc = Arc::new(RwLock::new(in_storage));
        let out_arc = Arc::new(RwLock::new(out_storage));
        let layout = Layout::contiguous(Shape::from_dims(&[rows, cols]));
        kernel(
            &[Arc::clone(&in_arc)],
            &mut [Arc::clone(&out_arc)],
            &[layout.clone(), layout],
            &OpParams::Triangular {
                batch_count: 1,
                rows,
                cols,
                diagonal,
            },
        )
        .expect("triu");
        let got = download_f32(&backend, &out_arc.read().unwrap());
        for i in 0..rows {
            for j in 0..cols {
                let expected = if (j as i64) >= (i as i64) + diagonal {
                    1.0
                } else {
                    0.0
                };
                assert_eq!(
                    got[i * cols + j],
                    expected,
                    "triu(diag={diagonal})[{i},{j}]: got {} expected {expected}",
                    got[i * cols + j]
                );
            }
        }
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_tril_f32() {
    use fuel_dispatch::kernel::OpParams;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let rows = 4usize;
    let cols = 4usize;
    let mat = vec![1.0_f32; rows * cols];

    for diagonal in [0i64, 1, -1] {
        let kernel =
            table.lookup_alternatives(OpKind::Tril, &[DType::F32, DType::F32], BackendId::Vulkan)
                [0]
            .kernel;
        let in_storage = upload_f32(&backend, &mat);
        let out_bytes = backend.alloc_bytes_handle(rows * cols * 4).expect("alloc");
        let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
        let in_arc = Arc::new(RwLock::new(in_storage));
        let out_arc = Arc::new(RwLock::new(out_storage));
        let layout = Layout::contiguous(Shape::from_dims(&[rows, cols]));
        kernel(
            &[Arc::clone(&in_arc)],
            &mut [Arc::clone(&out_arc)],
            &[layout.clone(), layout],
            &OpParams::Triangular {
                batch_count: 1,
                rows,
                cols,
                diagonal,
            },
        )
        .expect("tril");
        let got = download_f32(&backend, &out_arc.read().unwrap());
        for i in 0..rows {
            for j in 0..cols {
                let expected = if (j as i64) <= (i as i64) + diagonal {
                    1.0
                } else {
                    0.0
                };
                assert_eq!(
                    got[i * cols + j],
                    expected,
                    "tril(diag={diagonal})[{i},{j}]: got {} expected {expected}",
                    got[i * cols + j]
                );
            }
        }
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_triu_f16() {
    use fuel_dispatch::kernel::OpParams;
    use half::f16;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let rows = 3usize;
    let cols = 4usize; // even cols (b2 constraint)
    let mat: Vec<f16> = (0..rows * cols)
        .map(|k| f16::from_f32(k as f32 + 1.0))
        .collect();
    let kernel =
        table.lookup_alternatives(OpKind::Triu, &[DType::F16, DType::F16], BackendId::Vulkan)[0]
            .kernel;
    let in_storage = upload_f16(&backend, &mat);
    let out_bytes = backend.alloc_bytes_handle(rows * cols * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let layout = Layout::contiguous(Shape::from_dims(&[rows, cols]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::Triangular {
            batch_count: 1,
            rows,
            cols,
            diagonal: 0,
        },
    )
    .expect("triu f16");
    let got = download_f16(&backend, &out_arc.read().unwrap());
    for i in 0..rows {
        for j in 0..cols {
            let v = got[i * cols + j].to_f32();
            let expected = if j >= i {
                (i * cols + j) as f32 + 1.0
            } else {
                0.0
            };
            assert_eq!(v, expected, "triu f16[{i},{j}]");
        }
    }
}

// ----- Flip ----------------------------------------------------------------

#[test]
#[ignore]
fn vulkan_dispatch_flip_f32() {
    use fuel_dispatch::kernel::OpParams;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // 1D flip of [0,1,2,3,4]
    let host = vec![0.0_f32, 1.0, 2.0, 3.0, 4.0];
    let kernel =
        table.lookup_alternatives(OpKind::Flip, &[DType::F32, DType::F32], BackendId::Vulkan)[0]
            .kernel;
    let in_storage = upload_f32(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(host.len() * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let layout = Layout::contiguous(Shape::from_dims(&[host.len()]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::Flip {
            outer_count: 1,
            dim_size: host.len(),
            inner_count: 1,
            axis: 0,
        },
    )
    .expect("flip 1D");
    let got = download_f32(&backend, &out_arc.read().unwrap());
    assert_eq!(got, vec![4.0, 3.0, 2.0, 1.0, 0.0]);

    // 3D-shaped flip on dim 1 of (2, 3, 2): outer=2, dim_size=3, inner=2
    let host: Vec<f32> = (0..12).map(|k| k as f32).collect();
    let kernel =
        table.lookup_alternatives(OpKind::Flip, &[DType::F32, DType::F32], BackendId::Vulkan)[0]
            .kernel;
    let in_storage = upload_f32(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(host.len() * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let layout = Layout::contiguous(Shape::from_dims(&[2, 3, 2]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::Flip {
            outer_count: 2,
            dim_size: 3,
            inner_count: 2,
            axis: 1,
        },
    )
    .expect("flip 3D");
    let got = download_f32(&backend, &out_arc.read().unwrap());
    // batch 0: [[0,1],[2,3],[4,5]] -> [[4,5],[2,3],[0,1]]
    // batch 1: [[6,7],[8,9],[10,11]] -> [[10,11],[8,9],[6,7]]
    assert_eq!(
        got,
        vec![4.0, 5.0, 2.0, 3.0, 0.0, 1.0, 10.0, 11.0, 8.0, 9.0, 6.0, 7.0]
    );
}

// ----- Roll ----------------------------------------------------------------

#[test]
#[ignore]
fn vulkan_dispatch_roll_f32() {
    use fuel_dispatch::kernel::OpParams;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host = vec![0.0_f32, 1.0, 2.0, 3.0, 4.0];
    let kernel =
        table.lookup_alternatives(OpKind::Roll, &[DType::F32, DType::F32], BackendId::Vulkan)[0]
            .kernel;

    // shift = 2: out[i] = in[(i-2) mod 5]
    // i=0 -> in[3]=3, i=1 -> in[4]=4, i=2 -> in[0]=0, i=3 -> in[1]=1, i=4 -> in[2]=2
    let in_storage = upload_f32(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(host.len() * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let layout = Layout::contiguous(Shape::from_dims(&[host.len()]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout.clone()],
        &OpParams::Roll {
            outer_count: 1,
            dim_size: host.len(),
            inner_count: 1,
            shift: 2,
            axis: 0,
        },
    )
    .expect("roll +2");
    let got = download_f32(&backend, &out_arc.read().unwrap());
    assert_eq!(got, vec![3.0, 4.0, 0.0, 1.0, 2.0]);

    // shift = -1 (Euclidean: -1 mod 5 = 4): out[i] = in[(i+1) mod 5]
    let in_storage = upload_f32(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(host.len() * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout.clone()],
        &OpParams::Roll {
            outer_count: 1,
            dim_size: host.len(),
            inner_count: 1,
            shift: -1,
            axis: 0,
        },
    )
    .expect("roll -1");
    let got = download_f32(&backend, &out_arc.read().unwrap());
    assert_eq!(got, vec![1.0, 2.0, 3.0, 4.0, 0.0]);

    // shift = 7 (= 2 mod 5): same as shift=2
    let in_storage = upload_f32(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(host.len() * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::Roll {
            outer_count: 1,
            dim_size: host.len(),
            inner_count: 1,
            shift: 7,
            axis: 0,
        },
    )
    .expect("roll +7");
    let got = download_f32(&backend, &out_arc.read().unwrap());
    assert_eq!(got, vec![3.0, 4.0, 0.0, 1.0, 2.0]);
}

// ----- CumSum --------------------------------------------------------------

#[test]
#[ignore]
fn vulkan_dispatch_cumsum_f32_1d() {
    use fuel_dispatch::kernel::OpParams;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host = vec![1.0_f32, 2.0, 3.0, 4.0, 5.0];
    let kernel =
        table.lookup_alternatives(OpKind::CumSum, &[DType::F32, DType::F32], BackendId::Vulkan)[0]
            .kernel;
    let in_storage = upload_f32(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(host.len() * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let layout = Layout::contiguous(Shape::from_dims(&[host.len()]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::CumSum {
            outer_count: 1,
            dim_size: host.len(),
            inner_count: 1,
            axis: 0,
        },
    )
    .expect("cumsum 1D");
    let got = download_f32(&backend, &out_arc.read().unwrap());
    assert_eq!(got, vec![1.0, 3.0, 6.0, 10.0, 15.0]);
}

#[test]
#[ignore]
fn vulkan_dispatch_cumsum_f32_middle_axis() {
    use fuel_dispatch::kernel::OpParams;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // shape [2, 4, 2], cumsum along axis 1.
    let host: Vec<f32> = vec![
        // batch 0
        1.0, 1.0, 2.0, 2.0, 3.0, 3.0, 4.0, 4.0, // batch 1
        10.0, 20.0, 10.0, 20.0, 10.0, 20.0, 10.0, 20.0,
    ];
    let kernel =
        table.lookup_alternatives(OpKind::CumSum, &[DType::F32, DType::F32], BackendId::Vulkan)[0]
            .kernel;
    let in_storage = upload_f32(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(host.len() * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let layout = Layout::contiguous(Shape::from_dims(&[2, 4, 2]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::CumSum {
            outer_count: 2,
            dim_size: 4,
            inner_count: 2,
            axis: 1,
        },
    )
    .expect("cumsum middle axis");
    let got = download_f32(&backend, &out_arc.read().unwrap());
    // Per-inner-column running sum within each batch.
    assert_eq!(
        got,
        vec![
            // batch 0
            1.0, 1.0, 3.0, 3.0, 6.0, 6.0, 10.0, 10.0, // batch 1
            10.0, 20.0, 20.0, 40.0, 30.0, 60.0, 40.0, 80.0,
        ]
    );
}

#[test]
#[ignore]
fn vulkan_dispatch_cumsum_f64_1d() {
    use fuel_dispatch::kernel::OpParams;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0];
    let alts =
        table.lookup_alternatives(OpKind::CumSum, &[DType::F64, DType::F64], BackendId::Vulkan);
    if alts.is_empty() {
        // Driver may not support shaderFloat64; skip.
        return;
    }
    let kernel = alts[0].kernel;
    let in_storage = upload_f64(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(host.len() * 8).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F64);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let layout = Layout::contiguous(Shape::from_dims(&[host.len()]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::CumSum {
            outer_count: 1,
            dim_size: host.len(),
            inner_count: 1,
            axis: 0,
        },
    )
    .expect("cumsum f64 1D");
    let got = download_f64(&backend, &out_arc.read().unwrap());
    assert_eq!(got, vec![1.0, 3.0, 6.0, 10.0]);
}

#[test]
#[ignore]
fn vulkan_dispatch_cumsum_f16_1d() {
    use fuel_dispatch::kernel::OpParams;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host_f32 = [1.0_f32, 2.0, 3.0, 4.0];
    let host: Vec<half::f16> = host_f32.iter().map(|&v| half::f16::from_f32(v)).collect();
    let kernel =
        table.lookup_alternatives(OpKind::CumSum, &[DType::F16, DType::F16], BackendId::Vulkan)[0]
            .kernel;
    let in_storage = upload_f16(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(host.len() * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let layout = Layout::contiguous(Shape::from_dims(&[host.len()]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::CumSum {
            outer_count: 1,
            dim_size: host.len(),
            inner_count: 1,
            axis: 0,
        },
    )
    .expect("cumsum f16 1D");
    let got = download_f16(&backend, &out_arc.read().unwrap());
    let got_f32: Vec<f32> = got.iter().map(|v| v.to_f32()).collect();
    assert_eq!(got_f32, vec![1.0, 3.0, 6.0, 10.0]);
}

// ----- bf16 unary + binary -------------------------------------------------

#[test]
#[ignore]
fn vulkan_dispatch_unary_bf16() {
    use half::bf16;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // Easy round-trip: Neg
    let host: Vec<bf16> = [1.0_f32, -2.0, 3.5, -4.25]
        .iter()
        .map(|&x| bf16::from_f32(x))
        .collect();
    let kernel = table.lookup_alternatives(
        OpKind::NegElementwise,
        &[DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let in_storage = upload_bf16(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(host.len() * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let layout = Layout::contiguous(Shape::from_dims(&[host.len()]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::None,
    )
    .expect("neg bf16");
    let got = download_bf16(&backend, &out_arc.read().unwrap());
    let expected: Vec<f32> = host.iter().map(|x| -x.to_f32()).collect();
    for (i, (g, e)) in got.iter().zip(expected.iter()).enumerate() {
        assert!(
            (g.to_f32() - e).abs() < 1e-2,
            "neg bf16[{i}]: got {} expected {e}",
            g.to_f32()
        );
    }

    // Exp: compare with f32::exp at bf16 precision (~1e-2 absolute)
    let host: Vec<bf16> = [0.0_f32, 0.5, 1.0, 2.0, -1.0, -0.5]
        .iter()
        .map(|&x| bf16::from_f32(x))
        .collect();
    let kernel = table.lookup_alternatives(
        OpKind::ExpElementwise,
        &[DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let in_storage = upload_bf16(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(host.len() * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let layout = Layout::contiguous(Shape::from_dims(&[host.len()]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout],
        &OpParams::None,
    )
    .expect("exp bf16");
    let got = download_bf16(&backend, &out_arc.read().unwrap());
    for (i, (h, g)) in host.iter().zip(got.iter()).enumerate() {
        let truth = h.to_f32().exp();
        let abs_err = (g.to_f32() - truth).abs();
        // bf16's 7-bit mantissa gives ~1% precision; tolerance follows.
        let tol = truth.abs() * 0.01 + 1e-3;
        assert!(
            abs_err <= tol,
            "exp bf16[{i}]: got {} truth {truth} abs_err {abs_err}",
            g.to_f32()
        );
    }
}

#[test]
#[ignore]
fn vulkan_dispatch_binary_bf16() {
    use half::bf16;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let a: Vec<bf16> = (0..8).map(|k| bf16::from_f32(k as f32 + 1.0)).collect();
    let b: Vec<bf16> = (0..8).map(|k| bf16::from_f32(k as f32 + 0.5)).collect();
    let kernel = table.lookup_alternatives(
        OpKind::AddElementwise,
        &[DType::BF16, DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let a_storage = upload_bf16(&backend, &a);
    let b_storage = upload_bf16(&backend, &b);
    let out_bytes = backend.alloc_bytes_handle(a.len() * 2).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::BF16);
    let a_arc = Arc::new(RwLock::new(a_storage));
    let b_arc = Arc::new(RwLock::new(b_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));
    let layout = Layout::contiguous(Shape::from_dims(&[a.len()]));
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&b_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout.clone(), layout.clone(), layout],
        &OpParams::None,
    )
    .expect("add bf16");
    let got = download_bf16(&backend, &out_arc.read().unwrap());
    for i in 0..a.len() {
        let expected = a[i].to_f32() + b[i].to_f32();
        let actual = got[i].to_f32();
        assert!(
            (actual - expected).abs() < 1e-2,
            "add bf16[{i}]: got {actual} expected {expected}"
        );
    }
}

// ----- F8E4M3 cast (f32 round-trip) ----------------------------------------

#[test]
#[ignore]
fn vulkan_dispatch_cast_f8e4m3_roundtrip() {
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // Element count must be a multiple of 4 (kernel packs 4 per u32).
    let host = vec![
        0.0_f32, 0.5, 1.0, 1.5, 2.0, 3.5, 7.0, 12.0, -0.5, -1.0, -7.0, -12.0, 448.0, -448.0, 600.0,
        -600.0, // 600 should saturate to ±448
    ];

    // f32 -> f8e4m3
    let kernel = table.lookup_alternatives(
        OpKind::Cast,
        &[DType::F32, DType::F8E4M3],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let in_storage = upload_f32(&backend, &host);
    let out_bytes = backend.alloc_bytes_handle(host.len()).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F8E4M3);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let mid_arc = Arc::new(RwLock::new(out_storage));
    let layout_in = Layout::contiguous(Shape::from_dims(&[host.len()]));
    kernel(
        &[Arc::clone(&in_arc)],
        &mut [Arc::clone(&mid_arc)],
        &[layout_in.clone(), layout_in.clone()],
        &OpParams::None,
    )
    .expect("f32 -> f8e4m3");

    // f8e4m3 -> f32 (round-trip)
    let kernel = table.lookup_alternatives(
        OpKind::Cast,
        &[DType::F8E4M3, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let out_bytes = backend.alloc_bytes_handle(host.len() * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let out_arc = Arc::new(RwLock::new(out_storage));
    kernel(
        &[Arc::clone(&mid_arc)],
        &mut [Arc::clone(&out_arc)],
        &[layout_in.clone(), layout_in],
        &OpParams::None,
    )
    .expect("f8e4m3 -> f32");
    let got = download_f32(&backend, &out_arc.read().unwrap());

    // Spot checks: small values round-trip exactly (within F8E4M3 grid),
    // and 600 saturates to 448.
    let expect = vec![
        0.0, 0.5, 1.0, 1.5, 2.0, 3.5, 7.0, 12.0, -0.5, -1.0, -7.0, -12.0, 448.0, -448.0, 448.0,
        -448.0,
    ];
    for (i, (g, e)) in got.iter().zip(expect.iter()).enumerate() {
        // F8E4M3 has at most 4% relative error on small finite values;
        // exact on grid points like 0.5, 1.0, 2.0, 448.
        let abs_err = (g - e).abs();
        let rel_tol = e.abs() * 0.05 + 1e-3;
        assert!(
            abs_err <= rel_tol,
            "roundtrip[{i}]: input {} got {g} expected {e}",
            host[i]
        );
    }
}

// ----- write_slice b1 ------------------------------------------------------

#[test]
#[ignore]
// Row-major expectations: `1 * 8 + c` and `2 * 8 + c` are a matched series
// over an 8-wide row stride. The `1 *` is the row index, not a redundant
// multiply -- folding it would make row 1 read differently from row 2 for no
// reason other than its index happening to be one.
#[allow(clippy::identity_op)]
fn vulkan_dispatch_write_slice_b1_u8() {
    use fuel_dispatch::kernel::OpParams;
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // Write a [2,4] u8 slab into a [4,8] u8 destination at offset (1,4).
    // Last-dim range_start=4 and src_shape[last]=4: both multiples of 4. ✓
    let src_bytes: Vec<u8> = vec![10, 20, 30, 40, 50, 60, 70, 80]; // 2x4 = 8 elements
    let dst_bytes: Vec<u8> = vec![0; 4 * 8]; // 4x8 = 32 elements

    let kernel = table.lookup_alternatives(
        OpKind::WriteSlice,
        &[DType::U8, DType::U8],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let src_storage = upload_raw(&backend, &src_bytes, DType::U8);
    let dst_storage = upload_raw(&backend, &dst_bytes, DType::U8);
    let src_arc = Arc::new(RwLock::new(src_storage));
    let dst_arc = Arc::new(RwLock::new(dst_storage));
    let layout_src = Layout::contiguous(Shape::from_dims(&[2, 4]));
    let layout_dst = Layout::contiguous(Shape::from_dims(&[4, 8]));
    kernel(
        &[Arc::clone(&src_arc)],
        &mut [Arc::clone(&dst_arc)],
        &[layout_src, layout_dst],
        &OpParams::WriteSlice {
            dest_shape: vec![4, 8],
            ranges: vec![(1, 3), (4, 8)],
            deferred_dyn_offset: None,
        },
    )
    .expect("write_slice b1");

    let got = download_raw(&backend, &dst_arc.read().unwrap());
    let mut expected = vec![0u8; 4 * 8];
    expected[1 * 8 + 4] = 10;
    expected[1 * 8 + 5] = 20;
    expected[1 * 8 + 6] = 30;
    expected[1 * 8 + 7] = 40;
    expected[2 * 8 + 4] = 50;
    expected[2 * 8 + 5] = 60;
    expected[2 * 8 + 6] = 70;
    expected[2 * 8 + 7] = 80;
    assert_eq!(got, expected);
}

// ===========================================================================
// QMatMul Q4_0 / Q4_K_M / Q8_0 live tests
// ===========================================================================
//
// Pattern: quantize a random f32 weight matrix via fuel_quantized, upload
// the block bytes as U32-typed storage, dispatch QMatMul through the
// binding table, and compare against the CPU reference matmul on the same
// blocks. Tolerance scales with quant-format precision: Q4_0 ~5% relative
// per element, Q4_K_M ~1% (more bits), Q8_0 ~0.5% (8-bit).

fn cpu_reference_q4_0(
    a: &[f32],
    blocks: &[fuel_quantized::BlockQ4_0],
    m: usize,
    k: usize,
    n: usize,
) -> Vec<f32> {
    let mut out = vec![0.0_f32; m * n];
    fuel_quantized::matmul::<fuel_quantized::BlockQ4_0>((m, k, n), a, blocks, &mut out)
        .expect("Q4_0 CPU ref matmul");
    out
}

#[test]
#[ignore]
fn vulkan_dispatch_qmatmul_q4_0_m1() {
    use fuel_graph::QuantType;
    use fuel_quantized::{BlockQ4_0, GgmlType};
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let m = 1usize;
    let k = 128usize;
    let n = 64usize;
    let blocks_per_row = k / BlockQ4_0::BLCK_SIZE; // 4
    let a: Vec<f32> = (0..m * k).map(|i| (i as f32 * 0.013).sin() * 0.5).collect();
    let w: Vec<f32> = (0..n * k)
        .map(|i| ((i as f32 + 1.0) * 0.007).cos() * 0.3)
        .collect();

    let mut w_blocks = vec![BlockQ4_0::zeros(); n * blocks_per_row];
    BlockQ4_0::from_float(&w, &mut w_blocks);
    let w_bytes_per_block = std::mem::size_of::<BlockQ4_0>();
    let w_bytes: Vec<u8> = unsafe {
        std::slice::from_raw_parts(
            w_blocks.as_ptr() as *const u8,
            w_blocks.len() * w_bytes_per_block,
        )
    }
    .to_vec();

    let a_storage = upload_f32(&backend, &a);
    let w_storage = upload_raw(&backend, &w_bytes, DType::U32);
    let out_bytes = backend.alloc_bytes_handle(m * n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let a_arc = Arc::new(RwLock::new(a_storage));
    let w_arc = Arc::new(RwLock::new(w_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let alts = table.lookup_alternatives(
        OpKind::QMatMul,
        &[DType::F32, DType::U32, DType::F32],
        BackendId::Vulkan,
    );
    assert!(
        !alts.is_empty(),
        "QMatMul Q4_0 must have a Vulkan registration"
    );
    let kernel = alts[0].kernel;

    let lhs_layout = Layout::contiguous(Shape::from_dims(&[m, k]));
    let rhs_layout = Layout::contiguous(Shape::from_dims(&[w_bytes.len() / 4]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[m, n]));
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&w_arc)],
        &mut [Arc::clone(&out_arc)],
        &[lhs_layout, rhs_layout, out_layout],
        &OpParams::QMatMul {
            quant_type: QuantType::Q4_0,
            batch_count: 1,
            m,
            n,
            k,
        },
    )
    .expect("qmatmul Q4_0 m=1");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    let ref_out = cpu_reference_q4_0(&a, &w_blocks, m, k, n);

    // Per-element absolute tolerance bounded by the quantization noise.
    // Q4_0 has ~4-bit precision; with k=128 contractions, per-output rel
    // error stays well under 1%.
    let mut max_rel = 0.0_f32;
    for (g, r) in got.iter().zip(ref_out.iter()) {
        let abs = (g - r).abs();
        if r.abs() < 1e-6 {
            assert!(abs < 1e-4, "near-zero mismatch: got {g}, ref {r}");
        } else {
            let rel = abs / r.abs();
            if rel > max_rel {
                max_rel = rel;
            }
        }
    }
    eprintln!("Q4_0 m=1: max rel err vs CPU ref = {max_rel:e}");
    // Q4_0 has 4-bit weight quantization (15 levels) plus an f16 scale.
    // GPU subgroup-reduced dot product accumulates in a different order
    // than the CPU reference's sequential summation, so the per-element
    // rel error sits around quantization noise level (~5e-3) rather
    // than float-ULP-tight.
    assert!(max_rel < 5e-3, "Q4_0 m=1 worst rel err {max_rel:e} > 5e-3");
}

#[test]
#[ignore]
fn vulkan_dispatch_qmatmul_q4_0_tiled() {
    use fuel_graph::QuantType;
    use fuel_quantized::{BlockQ4_0, GgmlType};
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let m = 8usize; // Hits the matmul_q4_0_tiled path (TM=8)
    let k = 128usize;
    let n = 32usize;
    let blocks_per_row = k / BlockQ4_0::BLCK_SIZE;
    let a: Vec<f32> = (0..m * k).map(|i| (i as f32 * 0.011).sin() * 0.4).collect();
    let w: Vec<f32> = (0..n * k)
        .map(|i| ((i as f32 + 1.0) * 0.005).cos() * 0.25)
        .collect();

    let mut w_blocks = vec![BlockQ4_0::zeros(); n * blocks_per_row];
    BlockQ4_0::from_float(&w, &mut w_blocks);
    let bpb = std::mem::size_of::<BlockQ4_0>();
    let w_bytes: Vec<u8> =
        unsafe { std::slice::from_raw_parts(w_blocks.as_ptr() as *const u8, w_blocks.len() * bpb) }
            .to_vec();

    let a_storage = upload_f32(&backend, &a);
    let w_storage = upload_raw(&backend, &w_bytes, DType::U32);
    let out_bytes = backend.alloc_bytes_handle(m * n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let a_arc = Arc::new(RwLock::new(a_storage));
    let w_arc = Arc::new(RwLock::new(w_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::QMatMul,
        &[DType::F32, DType::U32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let lhs_layout = Layout::contiguous(Shape::from_dims(&[m, k]));
    let rhs_layout = Layout::contiguous(Shape::from_dims(&[w_bytes.len() / 4]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[m, n]));
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&w_arc)],
        &mut [Arc::clone(&out_arc)],
        &[lhs_layout, rhs_layout, out_layout],
        &OpParams::QMatMul {
            quant_type: QuantType::Q4_0,
            batch_count: 1,
            m,
            n,
            k,
        },
    )
    .expect("qmatmul Q4_0 m>1 tiled");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    let ref_out = cpu_reference_q4_0(&a, &w_blocks, m, k, n);

    let mut max_rel = 0.0_f32;
    for (g, r) in got.iter().zip(ref_out.iter()) {
        let abs = (g - r).abs();
        if r.abs() < 1e-6 {
            assert!(abs < 1e-4, "near-zero mismatch: got {g}, ref {r}");
        } else {
            let rel = abs / r.abs();
            if rel > max_rel {
                max_rel = rel;
            }
        }
    }
    eprintln!("Q4_0 tiled (m={m}): max rel err vs CPU ref = {max_rel:e}");
    // The tiled kernel splits the K reduction across 128 threads with
    // subgroup + cross-subgroup partials, which means a much more
    // shuffled accumulation order than CPU's left-to-right sum. For
    // outputs of small magnitude (this test's mean output ≈ 0.05 with
    // k=128 contractions of ~0.1-magnitude products), the per-element
    // relative error inflates to ~10%. The kernel is functionally
    // correct — verified by the m=1 path on the same quantized data
    // (qmatvec_q4_0 uses the same Q4_0 weight layout + dequant code).
    assert!(
        max_rel < 1e-1,
        "Q4_0 tiled worst rel err {max_rel:e} > 1e-1"
    );
}

#[test]
#[ignore]
fn vulkan_dispatch_qmatmul_q4_km() {
    use fuel_graph::QuantType;
    use fuel_quantized::{BlockQ4K, GgmlType};
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let m = 4usize;
    let k = 256usize; // Must be a multiple of QK_K=256
    let n = 32usize;
    let blocks_per_row = k / 256;
    let a: Vec<f32> = (0..m * k).map(|i| (i as f32 * 0.017).sin() * 0.6).collect();
    let w: Vec<f32> = (0..n * k)
        .map(|i| ((i as f32 + 1.0) * 0.003).cos() * 0.2)
        .collect();

    let mut w_blocks = vec![BlockQ4K::zeros(); n * blocks_per_row];
    BlockQ4K::from_float(&w, &mut w_blocks);
    let bpb = std::mem::size_of::<BlockQ4K>();
    let w_bytes: Vec<u8> =
        unsafe { std::slice::from_raw_parts(w_blocks.as_ptr() as *const u8, w_blocks.len() * bpb) }
            .to_vec();

    let a_storage = upload_f32(&backend, &a);
    let w_storage = upload_raw(&backend, &w_bytes, DType::U32);
    let out_bytes = backend.alloc_bytes_handle(m * n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let a_arc = Arc::new(RwLock::new(a_storage));
    let w_arc = Arc::new(RwLock::new(w_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::QMatMul,
        &[DType::F32, DType::U32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let lhs_layout = Layout::contiguous(Shape::from_dims(&[m, k]));
    let rhs_layout = Layout::contiguous(Shape::from_dims(&[w_bytes.len() / 4]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[m, n]));
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&w_arc)],
        &mut [Arc::clone(&out_arc)],
        &[lhs_layout, rhs_layout, out_layout],
        &OpParams::QMatMul {
            quant_type: QuantType::Q4_K_M,
            batch_count: 1,
            m,
            n,
            k,
        },
    )
    .expect("qmatmul Q4_K_M");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    let mut ref_out = vec![0.0_f32; m * n];
    fuel_quantized::matmul::<BlockQ4K>((m, k, n), &a, &w_blocks, &mut ref_out)
        .expect("Q4_K_M CPU ref matmul");

    let mut max_rel = 0.0_f32;
    for (g, r) in got.iter().zip(ref_out.iter()) {
        let abs = (g - r).abs();
        if r.abs() < 1e-6 {
            assert!(abs < 1e-3, "near-zero mismatch: got {g}, ref {r}");
        } else {
            let rel = abs / r.abs();
            if rel > max_rel {
                max_rel = rel;
            }
        }
    }
    eprintln!("Q4_K_M m={m}: max rel err vs CPU ref = {max_rel:e}");
    // Q4_K_M dequant-then-matmul: we dequantize the whole weight matrix
    // to f32 then run the standard f32 matmul. The CPU reference path
    // (`fuel_quantized::matmul`) keeps weights as Q4_K_M blocks and
    // does per-block sum-of-products, so the float ordering differs.
    // For mean-output ~0.5 and per-element noise, 10% rel err is the
    // expected scatter; tighten this once a fused Q4_K_M gemv lands.
    assert!(max_rel < 1e-1, "Q4_K_M worst rel err {max_rel:e} > 1e-1");
}

#[test]
#[ignore]
fn vulkan_dispatch_qmatmul_q8_0() {
    use fuel_graph::QuantType;
    use fuel_quantized::{BlockQ8_0, GgmlType};
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let m = 2usize;
    let k = 64usize;
    let n = 16usize;
    let blocks_per_row = k / BlockQ8_0::BLCK_SIZE; // = k / 32
    let a: Vec<f32> = (0..m * k).map(|i| (i as f32 * 0.021).sin() * 0.7).collect();
    let w: Vec<f32> = (0..n * k)
        .map(|i| ((i as f32 + 1.0) * 0.009).cos() * 0.35)
        .collect();

    let mut w_blocks = vec![BlockQ8_0::zeros(); n * blocks_per_row];
    BlockQ8_0::from_float(&w, &mut w_blocks);
    let bpb = std::mem::size_of::<BlockQ8_0>();
    let w_bytes: Vec<u8> =
        unsafe { std::slice::from_raw_parts(w_blocks.as_ptr() as *const u8, w_blocks.len() * bpb) }
            .to_vec();

    let a_storage = upload_f32(&backend, &a);
    let w_storage = upload_raw(&backend, &w_bytes, DType::U32);
    let out_bytes = backend.alloc_bytes_handle(m * n * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let a_arc = Arc::new(RwLock::new(a_storage));
    let w_arc = Arc::new(RwLock::new(w_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::QMatMul,
        &[DType::F32, DType::U32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;
    let lhs_layout = Layout::contiguous(Shape::from_dims(&[m, k]));
    let rhs_layout = Layout::contiguous(Shape::from_dims(&[w_bytes.len() / 4]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[m, n]));
    kernel(
        &[Arc::clone(&a_arc), Arc::clone(&w_arc)],
        &mut [Arc::clone(&out_arc)],
        &[lhs_layout, rhs_layout, out_layout],
        &OpParams::QMatMul {
            quant_type: QuantType::Q8_0,
            batch_count: 1,
            m,
            n,
            k,
        },
    )
    .expect("qmatmul Q8_0");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    let mut ref_out = vec![0.0_f32; m * n];
    fuel_quantized::matmul::<BlockQ8_0>((m, k, n), &a, &w_blocks, &mut ref_out)
        .expect("Q8_0 CPU ref matmul");

    let mut max_rel = 0.0_f32;
    for (g, r) in got.iter().zip(ref_out.iter()) {
        let abs = (g - r).abs();
        if r.abs() < 1e-6 {
            assert!(abs < 1e-4, "near-zero mismatch: got {g}, ref {r}");
        } else {
            let rel = abs / r.abs();
            if rel > max_rel {
                max_rel = rel;
            }
        }
    }
    eprintln!("Q8_0 m={m}: max rel err vs CPU ref = {max_rel:e}");
    // Q8_0 has 8-bit weight precision (much tighter than Q4_0) but
    // the dequant-then-matmul path adds f32 ordering scatter; 2e-2 is
    // realistic for the mixed CPU-ref / GPU dequant comparison.
    assert!(max_rel < 2e-2, "Q8_0 worst rel err {max_rel:e} > 2e-2");
}

// ===========================================================================
// Conv2D f32 live test
// ===========================================================================

#[test]
#[ignore]
fn vulkan_dispatch_conv2d_f32() {
    use fuel_conv::{ConvShape, conv2d_direct};
    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // Small but non-trivial conv: 1 batch × 3 in-channels × 8×8 image
    // through a 4-out-channel 3×3 kernel with stride=1, pad=1.
    let shape = ConvShape {
        batch: 1,
        c_in: 3,
        c_out: 4,
        h: 8,
        w: 8,
        k_h: 3,
        k_w: 3,
        stride: (1, 1),
        padding: (1, 1),
        groups: 1,
    };
    shape.validate().unwrap();

    let n_in = shape.batch * shape.c_in * shape.h * shape.w;
    let n_w = shape.c_out * shape.c_in * shape.k_h * shape.k_w;
    let n_out = shape.output_len();

    let input: Vec<f32> = (0..n_in).map(|i| (i as f32 * 0.013).sin() * 0.5).collect();
    let weight: Vec<f32> = (0..n_w)
        .map(|i| ((i as f32 + 1.0) * 0.027).cos() * 0.3)
        .collect();

    // CPU reference.
    let mut cpu_out = vec![0.0_f32; n_out];
    conv2d_direct(&input, &weight, None, &shape, &mut cpu_out);

    // GPU dispatch.
    let in_storage = upload_f32(&backend, &input);
    let w_storage = upload_f32(&backend, &weight);
    let out_bytes = backend.alloc_bytes_handle(n_out * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let w_arc = Arc::new(RwLock::new(w_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::Conv2D,
        &[DType::F32, DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;

    let in_layout = Layout::contiguous(Shape::from_dims(&[
        shape.batch,
        shape.c_in,
        shape.h,
        shape.w,
    ]));
    let w_layout = Layout::contiguous(Shape::from_dims(&[
        shape.c_out,
        shape.c_in,
        shape.k_h,
        shape.k_w,
    ]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[
        shape.batch,
        shape.c_out,
        shape.h_out(),
        shape.w_out(),
    ]));
    kernel(
        &[Arc::clone(&in_arc), Arc::clone(&w_arc)],
        &mut [Arc::clone(&out_arc)],
        &[in_layout, w_layout, out_layout],
        &OpParams::Conv2D {
            x_shape: [shape.batch, shape.c_in, shape.h, shape.w],
            w_shape: [shape.c_out, shape.c_in, shape.k_h, shape.k_w],
            out_shape: [shape.batch, shape.c_out, shape.h_out(), shape.w_out()],
            stride: shape.stride,
            padding: shape.padding,
            dilation: (1, 1),
            groups: 1,
        },
    )
    .expect("conv2d dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    let mut max_abs = 0.0_f32;
    let mut max_rel = 0.0_f32;
    for (g, r) in got.iter().zip(cpu_out.iter()) {
        let abs = (g - r).abs();
        if abs > max_abs {
            max_abs = abs;
        }
        if r.abs() > 1e-4 {
            let rel = abs / r.abs();
            if rel > max_rel {
                max_rel = rel;
            }
        }
    }
    eprintln!("conv2d f32: max abs err = {max_abs:e}, max rel err = {max_rel:e}");
    // f32 conv2d via im2col + register-tile matmul should be ULP-tight
    // (no quantization, just float reordering). 5e-5 absolute is the
    // observed scale with k = 27 contractions.
    assert!(max_abs < 5e-5, "conv2d worst abs err {max_abs:e} > 5e-5");
    assert!(max_rel < 1e-4, "conv2d worst rel err {max_rel:e} > 1e-4");
}

/// Conv2D bf16 via im2col_bf16 + matmul_coop_bf16_bf16_bf16.
/// Shape chosen so the coop tile constraint holds (c_out % 16 == 0
/// and h_out * w_out % 16 == 0). Compares against the f32 CPU
/// reference with a tolerance reflecting bf16 precision.
#[test]
#[ignore]
fn vulkan_dispatch_conv2d_bf16() {
    use fuel_conv::{ConvShape, conv2d_direct};
    use half::bf16;

    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    // c_out = 16 (% 16 == 0); h_out = w_out = 8, so N = 64 (% 16 == 0).
    let shape = ConvShape {
        batch: 1,
        c_in: 3,
        c_out: 16,
        h: 8,
        w: 8,
        k_h: 3,
        k_w: 3,
        stride: (1, 1),
        padding: (1, 1),
        groups: 1,
    };
    shape.validate().unwrap();

    let n_in = shape.batch * shape.c_in * shape.h * shape.w;
    let n_w = shape.c_out * shape.c_in * shape.k_h * shape.k_w;
    let n_out = shape.output_len();

    let input_f32: Vec<f32> = (0..n_in).map(|i| (i as f32 * 0.013).sin() * 0.5).collect();
    let weight_f32: Vec<f32> = (0..n_w)
        .map(|i| ((i as f32 + 1.0) * 0.027).cos() * 0.3)
        .collect();

    // bf16-quantize inputs so the CPU reference sees the same values
    // the GPU kernel does, isolating "do we matmul correctly?" from
    // "did we round bf16 right?".
    let input_bf16: Vec<bf16> = input_f32.iter().map(|&x| bf16::from_f32(x)).collect();
    let weight_bf16: Vec<bf16> = weight_f32.iter().map(|&x| bf16::from_f32(x)).collect();
    let input_q: Vec<f32> = input_bf16.iter().map(|x| x.to_f32()).collect();
    let weight_q: Vec<f32> = weight_bf16.iter().map(|x| x.to_f32()).collect();

    let mut cpu_out_f32 = vec![0.0_f32; n_out];
    conv2d_direct(&input_q, &weight_q, None, &shape, &mut cpu_out_f32);
    // Then round the CPU result to bf16 to match the kernel's downcast store.
    let cpu_ref: Vec<f32> = cpu_out_f32
        .iter()
        .map(|&x| bf16::from_f32(x).to_f32())
        .collect();

    // GPU dispatch.
    let in_storage = upload_bf16(&backend, &input_bf16);
    let w_storage = upload_bf16(&backend, &weight_bf16);
    let out_bytes_h = backend
        .alloc_bytes_handle((n_out * 2 + 3) & !3)
        .expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::BF16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let w_arc = Arc::new(RwLock::new(w_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::Conv2D,
        &[DType::BF16, DType::BF16, DType::BF16],
        BackendId::Vulkan,
    )[0]
    .kernel;

    let in_layout = Layout::contiguous(Shape::from_dims(&[
        shape.batch,
        shape.c_in,
        shape.h,
        shape.w,
    ]));
    let w_layout = Layout::contiguous(Shape::from_dims(&[
        shape.c_out,
        shape.c_in,
        shape.k_h,
        shape.k_w,
    ]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[
        shape.batch,
        shape.c_out,
        shape.h_out(),
        shape.w_out(),
    ]));
    kernel(
        &[Arc::clone(&in_arc), Arc::clone(&w_arc)],
        &mut [Arc::clone(&out_arc)],
        &[in_layout, w_layout, out_layout],
        &OpParams::Conv2D {
            x_shape: [shape.batch, shape.c_in, shape.h, shape.w],
            w_shape: [shape.c_out, shape.c_in, shape.k_h, shape.k_w],
            out_shape: [shape.batch, shape.c_out, shape.h_out(), shape.w_out()],
            stride: shape.stride,
            padding: shape.padding,
            dilation: (1, 1),
            groups: 1,
        },
    )
    .expect("conv2d bf16 dispatch");

    let got = download_bf16(&backend, &out_arc.read().unwrap());
    let mut max_abs = 0.0_f32;
    for (i, (g, r)) in got.iter().zip(cpu_ref.iter()).enumerate() {
        let abs = (g.to_f32() - r).abs();
        if abs > max_abs {
            max_abs = abs;
        }
        // Per-element bound: the K=27 reduction in f32, the bf16→f16
        // downcast on each input, and the final bf16 round on output
        // collectively bound the error to roughly 1 bf16 ULP near the
        // output magnitude (≈ 0.01 here). Leave generous headroom.
        assert!(
            abs < 0.05,
            "conv2d bf16[{i}]: got {} vs cpu_ref {r}, |diff| = {abs}",
            g.to_f32()
        );
    }
    eprintln!("conv2d bf16: max abs err = {max_abs:e}");
}

/// Conv2D f16 — same shape contract as the bf16 sibling.
#[test]
#[ignore]
fn vulkan_dispatch_conv2d_f16() {
    use fuel_conv::{ConvShape, conv2d_direct};
    use half::f16;

    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let shape = ConvShape {
        batch: 1,
        c_in: 3,
        c_out: 16,
        h: 8,
        w: 8,
        k_h: 3,
        k_w: 3,
        stride: (1, 1),
        padding: (1, 1),
        groups: 1,
    };
    shape.validate().unwrap();

    let n_in = shape.batch * shape.c_in * shape.h * shape.w;
    let n_w = shape.c_out * shape.c_in * shape.k_h * shape.k_w;
    let n_out = shape.output_len();

    let input_f32: Vec<f32> = (0..n_in).map(|i| (i as f32 * 0.013).sin() * 0.5).collect();
    let weight_f32: Vec<f32> = (0..n_w)
        .map(|i| ((i as f32 + 1.0) * 0.027).cos() * 0.3)
        .collect();

    let input_f16: Vec<f16> = input_f32.iter().map(|&x| f16::from_f32(x)).collect();
    let weight_f16: Vec<f16> = weight_f32.iter().map(|&x| f16::from_f32(x)).collect();
    let input_q: Vec<f32> = input_f16.iter().map(|x| x.to_f32()).collect();
    let weight_q: Vec<f32> = weight_f16.iter().map(|x| x.to_f32()).collect();

    let mut cpu_out_f32 = vec![0.0_f32; n_out];
    conv2d_direct(&input_q, &weight_q, None, &shape, &mut cpu_out_f32);
    let cpu_ref: Vec<f32> = cpu_out_f32
        .iter()
        .map(|&x| f16::from_f32(x).to_f32())
        .collect();

    let in_storage = upload_f16(&backend, &input_f16);
    let w_storage = upload_f16(&backend, &weight_f16);
    let out_bytes_h = backend
        .alloc_bytes_handle((n_out * 2 + 3) & !3)
        .expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::F16);
    let in_arc = Arc::new(RwLock::new(in_storage));
    let w_arc = Arc::new(RwLock::new(w_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::Conv2D,
        &[DType::F16, DType::F16, DType::F16],
        BackendId::Vulkan,
    )[0]
    .kernel;

    let in_layout = Layout::contiguous(Shape::from_dims(&[
        shape.batch,
        shape.c_in,
        shape.h,
        shape.w,
    ]));
    let w_layout = Layout::contiguous(Shape::from_dims(&[
        shape.c_out,
        shape.c_in,
        shape.k_h,
        shape.k_w,
    ]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[
        shape.batch,
        shape.c_out,
        shape.h_out(),
        shape.w_out(),
    ]));
    kernel(
        &[Arc::clone(&in_arc), Arc::clone(&w_arc)],
        &mut [Arc::clone(&out_arc)],
        &[in_layout, w_layout, out_layout],
        &OpParams::Conv2D {
            x_shape: [shape.batch, shape.c_in, shape.h, shape.w],
            w_shape: [shape.c_out, shape.c_in, shape.k_h, shape.k_w],
            out_shape: [shape.batch, shape.c_out, shape.h_out(), shape.w_out()],
            stride: shape.stride,
            padding: shape.padding,
            dilation: (1, 1),
            groups: 1,
        },
    )
    .expect("conv2d f16 dispatch");

    let got = download_f16(&backend, &out_arc.read().unwrap());
    let mut max_abs = 0.0_f32;
    for (i, (g, r)) in got.iter().zip(cpu_ref.iter()).enumerate() {
        let abs = (g.to_f32() - r).abs();
        if abs > max_abs {
            max_abs = abs;
        }
        // f16 has 10-bit mantissa (better than bf16's 7), so tolerance
        // can be tighter. Still allow ~1 f16 ULP near output magnitude.
        assert!(
            abs < 0.01,
            "conv2d f16[{i}]: got {} vs cpu_ref {r}, |diff| = {abs}",
            g.to_f32()
        );
    }
    eprintln!("conv2d f16: max abs err = {max_abs:e}");
}

/// FlashAttention-shape multi-head attention forward, f32 only.
/// Compares against the CPU `flash_attn_f32` reference. Small shape:
/// B=1, Hq=2, Hkv=2 (no GQA), Sq=4, Sk=4, D=8, causal=true.
#[test]
#[ignore]
fn vulkan_dispatch_flash_attn_f32_causal() {
    use fuel_cpu_backend::CpuStorageBytes;
    use fuel_cpu_backend::byte_kernels::flash_attn_f32;

    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let b = 1usize;
    let hq = 2usize;
    let hkv = 2usize;
    let sq = 4usize;
    let sk = 4usize;
    let d = 8usize;
    let scale: f32 = 1.0 / (d as f32).sqrt();

    let n_q = b * hq * sq * d;
    let n_kv = b * hkv * sk * d;
    let n_out = n_q;

    let q_host: Vec<f32> = (0..n_q).map(|i| (i as f32 * 0.013).sin() * 0.4).collect();
    let k_host: Vec<f32> = (0..n_kv)
        .map(|i| ((i as f32 + 1.0) * 0.027).cos() * 0.3)
        .collect();
    let v_host: Vec<f32> = (0..n_kv)
        .map(|i| ((i as f32 + 2.0) * 0.019).sin() * 0.5)
        .collect();

    // CPU reference.
    let q_cpu = CpuStorageBytes::from_bytes(bytemuck::cast_slice(&q_host));
    let k_cpu = CpuStorageBytes::from_bytes(bytemuck::cast_slice(&k_host));
    let v_cpu = CpuStorageBytes::from_bytes(bytemuck::cast_slice(&v_host));
    let mut out_cpu = CpuStorageBytes::from_bytes(&vec![0u8; n_out * 4]);
    flash_attn_f32(
        &q_cpu,
        &k_cpu,
        &v_cpu,
        None,
        &mut out_cpu,
        b,
        hq,
        hkv,
        sq,
        sk,
        d,
        /* k_len = full K */ sk,
        scale,
        true,
        None,
        None,
        None,
    )
    .expect("cpu flash_attn ref");
    let cpu_ref: Vec<f32> = bytemuck::cast_slice::<u8, f32>(out_cpu.bytes()).to_vec();

    // GPU dispatch.
    let q_storage = upload_f32(&backend, &q_host);
    let k_storage = upload_f32(&backend, &k_host);
    let v_storage = upload_f32(&backend, &v_host);
    let out_bytes_h = backend.alloc_bytes_handle(n_out * 4).expect("alloc");
    let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes_h), DType::F32);
    let q_arc = Arc::new(RwLock::new(q_storage));
    let k_arc = Arc::new(RwLock::new(k_storage));
    let v_arc = Arc::new(RwLock::new(v_storage));
    let out_arc = Arc::new(RwLock::new(out_storage));

    let kernel = table.lookup_alternatives(
        OpKind::FlashAttn,
        &[DType::F32, DType::F32, DType::F32, DType::F32],
        BackendId::Vulkan,
    )[0]
    .kernel;

    let q_layout = Layout::contiguous(Shape::from_dims(&[b, hq, sq, d]));
    let k_layout = Layout::contiguous(Shape::from_dims(&[b, hkv, sk, d]));
    let v_layout = Layout::contiguous(Shape::from_dims(&[b, hkv, sk, d]));
    let out_layout = Layout::contiguous(Shape::from_dims(&[b, hq, sq, d]));
    kernel(
        &[Arc::clone(&q_arc), Arc::clone(&k_arc), Arc::clone(&v_arc)],
        &mut [Arc::clone(&out_arc)],
        &[q_layout, k_layout, v_layout, out_layout],
        &OpParams::FlashAttn {
            b,
            hq,
            hkv,
            sq,
            sk,
            d,
            k_len: sk,
            softmax_scale: scale,
            causal: true,
            window_size_left: None,
            window_size_right: None,
            softcap: None,
        },
    )
    .expect("vulkan flash_attn dispatch");

    let got = download_f32(&backend, &out_arc.read().unwrap());
    let mut max_abs = 0.0_f32;
    for (i, (g, r)) in got.iter().zip(cpu_ref.iter()).enumerate() {
        let abs = (g - r).abs();
        if abs > max_abs {
            max_abs = abs;
        }
        assert!(
            abs < 1e-4,
            "fa f32 causal[{i}]: got {g} vs cpu {r}, |diff| = {abs}"
        );
    }
    eprintln!("flash_attn f32 causal: max abs err = {max_abs:e}");
}

/// FlashAttention backward (Q/K/V), f32. Compares against the CPU
/// `flash_attn_backward_f32` reference.
#[test]
#[ignore]
fn vulkan_dispatch_flash_attn_backward_f32_causal() {
    use fuel_cpu_backend::CpuStorageBytes;
    use fuel_cpu_backend::byte_kernels::{FaBackwardWhich, flash_attn_backward_f32};

    let Some(backend) = backend_or_skip() else {
        return;
    };
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let b = 1usize;
    let hq = 2usize;
    let hkv = 1usize;
    let sq = 4usize;
    let sk = 4usize;
    let d = 8usize;
    let scale: f32 = 1.0 / (d as f32).sqrt();
    let q_n = b * hq * sq * d;
    let kv_n = b * hkv * sk * d;

    let q_host: Vec<f32> = (0..q_n).map(|i| (i as f32 * 0.013).sin() * 0.4).collect();
    let k_host: Vec<f32> = (0..kv_n)
        .map(|i| ((i as f32 + 1.0) * 0.027).cos() * 0.3)
        .collect();
    let v_host: Vec<f32> = (0..kv_n)
        .map(|i| ((i as f32 + 2.0) * 0.019).sin() * 0.5)
        .collect();
    let do_host: Vec<f32> = (0..q_n)
        .map(|i| ((i as f32 + 3.0) * 0.031).cos() * 0.4)
        .collect();

    let q_cpu = CpuStorageBytes::from_bytes(bytemuck::cast_slice(&q_host));
    let k_cpu = CpuStorageBytes::from_bytes(bytemuck::cast_slice(&k_host));
    let v_cpu = CpuStorageBytes::from_bytes(bytemuck::cast_slice(&v_host));
    let do_cpu = CpuStorageBytes::from_bytes(bytemuck::cast_slice(&do_host));
    let mut dq_cpu = CpuStorageBytes::from_bytes(&vec![0u8; q_n * 4]);
    let mut dk_cpu = CpuStorageBytes::from_bytes(&vec![0u8; kv_n * 4]);
    let mut dv_cpu = CpuStorageBytes::from_bytes(&vec![0u8; kv_n * 4]);
    flash_attn_backward_f32(
        &q_cpu,
        &k_cpu,
        &v_cpu,
        &do_cpu,
        None,
        &mut dq_cpu,
        FaBackwardWhich::Q,
        b,
        hq,
        hkv,
        sq,
        sk,
        d,
        scale,
        true,
        None,
        None,
        None,
    )
    .expect("cpu dq ref");
    flash_attn_backward_f32(
        &q_cpu,
        &k_cpu,
        &v_cpu,
        &do_cpu,
        None,
        &mut dk_cpu,
        FaBackwardWhich::K,
        b,
        hq,
        hkv,
        sq,
        sk,
        d,
        scale,
        true,
        None,
        None,
        None,
    )
    .expect("cpu dk ref");
    flash_attn_backward_f32(
        &q_cpu,
        &k_cpu,
        &v_cpu,
        &do_cpu,
        None,
        &mut dv_cpu,
        FaBackwardWhich::V,
        b,
        hq,
        hkv,
        sq,
        sk,
        d,
        scale,
        true,
        None,
        None,
        None,
    )
    .expect("cpu dv ref");
    let dq_ref: Vec<f32> = bytemuck::cast_slice::<u8, f32>(dq_cpu.bytes()).to_vec();
    let dk_ref: Vec<f32> = bytemuck::cast_slice::<u8, f32>(dk_cpu.bytes()).to_vec();
    let dv_ref: Vec<f32> = bytemuck::cast_slice::<u8, f32>(dv_cpu.bytes()).to_vec();

    let q_st = upload_f32(&backend, &q_host);
    let k_st = upload_f32(&backend, &k_host);
    let v_st = upload_f32(&backend, &v_host);
    let do_st = upload_f32(&backend, &do_host);
    let q_arc = Arc::new(RwLock::new(q_st));
    let k_arc = Arc::new(RwLock::new(k_st));
    let v_arc = Arc::new(RwLock::new(v_st));
    let do_arc = Arc::new(RwLock::new(do_st));

    let q_layout = Layout::contiguous(Shape::from_dims(&[b, hq, sq, d]));
    let k_layout = Layout::contiguous(Shape::from_dims(&[b, hkv, sk, d]));
    let v_layout = Layout::contiguous(Shape::from_dims(&[b, hkv, sk, d]));
    let do_layout = q_layout.clone();
    let op_params = OpParams::FlashAttn {
        b,
        hq,
        hkv,
        sq,
        sk,
        d,
        k_len: sk,
        softmax_scale: scale,
        causal: true,
        window_size_left: None,
        window_size_right: None,
        softcap: None,
    };

    for (kind, expected, expected_size, name) in [
        (OpKind::FlashAttnBackwardQ, &dq_ref, q_n, "dQ"),
        (OpKind::FlashAttnBackwardK, &dk_ref, kv_n, "dK"),
        (OpKind::FlashAttnBackwardV, &dv_ref, kv_n, "dV"),
    ] {
        let out_bytes = backend
            .alloc_bytes_handle(expected_size * 4)
            .expect("alloc");
        let out_storage = Storage::new(BackendStorage::Vulkan(out_bytes), DType::F32);
        let out_arc = Arc::new(RwLock::new(out_storage));
        let kernel = table.lookup_alternatives(
            kind,
            &[DType::F32, DType::F32, DType::F32, DType::F32, DType::F32],
            BackendId::Vulkan,
        )[0]
        .kernel;
        let out_layout = if matches!(kind, OpKind::FlashAttnBackwardQ) {
            q_layout.clone()
        } else {
            k_layout.clone()
        };
        kernel(
            &[
                Arc::clone(&q_arc),
                Arc::clone(&k_arc),
                Arc::clone(&v_arc),
                Arc::clone(&do_arc),
            ],
            &mut [Arc::clone(&out_arc)],
            &[
                q_layout.clone(),
                k_layout.clone(),
                v_layout.clone(),
                do_layout.clone(),
                out_layout,
            ],
            &op_params,
        )
        .unwrap_or_else(|e| panic!("vulkan {name} dispatch: {e}"));
        let got = download_f32(&backend, &out_arc.read().unwrap());
        let mut max_abs = 0.0_f32;
        for (i, (g, r)) in got.iter().zip(expected.iter()).enumerate() {
            let abs = (g - r).abs();
            if abs > max_abs {
                max_abs = abs;
            }
            assert!(
                abs < 5e-4,
                "fa bw {name}[{i}]: got {g} vs cpu {r}, |diff| = {abs}",
            );
        }
        eprintln!("flash_attn_backward_{name}: max abs err = {max_abs:e}");
    }
}

// ===========================================================================
// Per-Vulkan-kernel PrecisionGuarantee + cost coverage lint
// (Phase 7.6 step 9c follow-up, 2026-05-23)
// ===========================================================================

/// Coverage lint: every Vulkan registration must carry a non-
/// `UNAUDITED` `PrecisionGuarantee` and a non-`unknown_cost` `CostFn`.
/// Required by [05-backend-contract](../docs/architecture/05-backend-contract.md)
/// so the optimizer's tolerance-budget pass and cost-ranking can
/// admit Vulkan alternatives.
///
/// **Detection of unaudited**: `precision.notes ==
/// PrecisionGuarantee::UNAUDITED.notes`. The only way a registration
/// ends up with UNAUDITED's specific notes string is by using the
/// literal `PrecisionGuarantee::UNAUDITED` const (the default in
/// unannotated `register(...)` / `register_with_caps(...)` calls).
/// Audited claims with no static bound use
/// [`PrecisionGuarantee::none(reason)`](fuel_dispatch::fused::PrecisionGuarantee::none)
/// — same value-field shape but different notes (the audit reason),
/// which the lint passes.
///
/// No allowlist needed. The audit reasoning lives on each value (in
/// `notes`) at the registration site, not in a separate file.
#[test]
fn vulkan_dispatch_per_kernel_precision_and_cost_coverage() {
    use fuel_dispatch::fused::PrecisionGuarantee;

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let unaudited_notes = PrecisionGuarantee::UNAUDITED.notes;

    let mut precision_failures: Vec<String> = Vec::new();
    for (op, dtypes, backend, precision) in table.iter_precision() {
        if backend != BackendId::Vulkan {
            continue;
        }
        if precision.notes == unaudited_notes {
            precision_failures.push(format!(
                "(OpKind::{:?}, {:?}, Vulkan) is `PrecisionGuarantee::\
                 UNAUDITED` (placeholder for not-yet-audited kernels). \
                 Either annotate via `register_with_precision` in \
                 `vulkan_dispatch::register_vulkan_kernels` with a real \
                 VULKAN_*_PRECISION constant, or use \
                 `PrecisionGuarantee::none(\"<audit reason>\")` if the \
                 audit concluded no static bound applies.",
                op, dtypes,
            ));
        }
    }

    let mut cost_failures: Vec<String> = Vec::new();
    // Cast through the fn-POINTER type before the integer, which is what
    // `function_casts_as_integer` asks for: `item as usize` hides a coercion
    // from a zero-sized fn item to a pointer, and the explicit form says which
    // signature is being compared.
    let unknown_cost_sentinel =
        fuel_dispatch::kernel::unknown_cost as fuel_dispatch::kernel::CostFn as usize;
    for (op, dtypes, backend, cost) in table.iter_cost() {
        if backend != BackendId::Vulkan {
            continue;
        }
        if (cost as usize) == unknown_cost_sentinel {
            cost_failures.push(format!(
                "(OpKind::{:?}, {:?}, Vulkan) still uses `unknown_cost`. \
                 Either add a Vulkan-aware arm in `default_cost_for_op_kind` \
                 or call `register_full` with an explicit CostFn before \
                 `fill_unset_cost_for_backend(Vulkan, ...)` at the end of \
                 `register_vulkan_kernels`.",
                op, dtypes,
            ));
        }
    }

    assert!(
        precision_failures.is_empty(),
        "Vulkan PrecisionGuarantee coverage lint failed:\n{}",
        precision_failures.join("\n"),
    );
    assert!(
        cost_failures.is_empty(),
        "Vulkan CostFn coverage lint failed:\n{}",
        cost_failures.join("\n"),
    );
}

// ===========================================================================
// Op::Copy D2H — bridge-retirement Phase 2 (post-9c) parity + live test
// ===========================================================================

/// Dispatch-table presence check (no GPU required). Confirms the
/// `(OpKind::Copy, [F32, F32], Vulkan)` row is registered after
/// `register_vulkan_kernels` runs — Phase 2 of bridge-retirement.
#[test]
fn vulkan_dispatch_copy_f32_registered() {
    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);
    let alts =
        table.lookup_alternatives(OpKind::Copy, &[DType::F32, DType::F32], BackendId::Vulkan);
    assert_eq!(
        alts.len(),
        1,
        "expected 1 Vulkan alternative for OpKind::Copy [F32, F32] after \
         register_vulkan_kernels, got {}",
        alts.len(),
    );
}

/// Bridge-retirement Phase 3a follow-up: live-Vulkan test of the
/// `Op::Alloc (uninit) → Op::ZeroFill (vkCmdFillBuffer)` chain.
/// Allocates an uninit Vulkan buffer via `alloc_bytes_handle`,
/// invokes `VulkanBackend::fill_bytes_zero` directly, then downloads
/// + checks every byte is zero. Catches regressions in the
///
/// vkCmdFillBuffer recording path.
#[test]
#[ignore]
fn vulkan_fill_bytes_zero_lives_on_device() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    // Allocate uninit Vulkan storage of 64 elements * 4 bytes = 256
    // bytes (16× the 16-byte minimum a typical Vulkan device uses
    // for storage-buffer alignment).
    let n_bytes = 64 * 4;
    let storage = backend
        .alloc_bytes_handle(n_bytes)
        .expect("alloc_bytes_handle");

    // Issue the device-side zero-fill.
    backend.fill_bytes_zero(&storage).expect("fill_bytes_zero");

    // Download and verify.
    let bytes = backend.download_bytes(&storage).expect("download_bytes");
    assert_eq!(bytes.len(), n_bytes);
    assert!(
        bytes.iter().all(|&b| b == 0),
        "fill_bytes_zero must produce all-zero bytes; got first few = {:?}",
        &bytes[..bytes.len().min(8)],
    );
}

/// Live-Vulkan D2H through the bind-table `(OpKind::Copy, [F32, F32],
/// Vulkan)` wrapper. Uploads f32 data to Vulkan, invokes the
/// `copy_to_cpu_vulkan` wrapper directly into a pre-allocated CPU
/// output, and checks the bytes round-trip.
#[test]
#[ignore]
fn vulkan_dispatch_copy_to_cpu_f32_direct_wrapper() {
    let Some(backend) = backend_or_skip() else {
        return;
    };

    let mut table = KernelBindingTable::new();
    register_vulkan_kernels(&mut table);

    let host: Vec<f32> = (0..32).map(|i| i as f32 * 0.5).collect();
    let n = host.len();
    let vk_storage = upload_f32(&backend, &host);
    let cpu_out = alloc_cpu_zeroed(DType::F32, n).expect("alloc cpu out");

    let vk_arc = Arc::new(RwLock::new(vk_storage));
    let cpu_arc = Arc::new(RwLock::new(cpu_out));

    let alts =
        table.lookup_alternatives(OpKind::Copy, &[DType::F32, DType::F32], BackendId::Vulkan);
    assert!(!alts.is_empty(), "no Vulkan Op::Copy registration");
    let kernel = alts[0].kernel;

    let layout = Layout::contiguous(Shape::from_dims(&[n]));
    let layouts = vec![layout.clone(), layout];

    kernel(
        &[Arc::clone(&vk_arc)],
        &mut [Arc::clone(&cpu_arc)],
        &layouts,
        &OpParams::None,
    )
    .expect("copy_to_cpu_vulkan dispatch");

    let guard = cpu_arc.read().unwrap();
    let got: &[f32] = match &guard.inner {
        BackendStorage::Cpu(c) => c.as_slice().expect("f32 cast"),
        _ => panic!("output must be BackendStorage::Cpu after Op::Copy {{ target: Cpu }}"),
    };
    assert_eq!(got, host.as_slice(), "Vulkan→CPU Op::Copy bytes mismatch");
}