From 359aee7f1d568718fd08d1ebb5168d2e2dad44d3 Mon Sep 17 00:00:00 2001 From: TheTom Date: Wed, 22 Jul 2026 21:47:10 -0500 Subject: [PATCH 1/5] feat(gemv_q8): plain MoE gather kernel + bf16 casts, f16-scale param rename Add ffai_moe_gather_q4, the no-activation sibling of the ReLU2 gather, so gated (SwiGLU) MoE paths can gather the gate and up expert stacks with in-kernel expert indexing. Add ffai_cast_f32_bf16 / ffai_cast_bf16_f32 elementwise casts for bf16 residual streams. Rename the misleading d_f32 parameter to d_f16 on the five kernels whose per-block scales are f16 (gemv_q4 coalesced/relu2/accum, moe_gather_q4 relu2/down) and document the scale dtype contract on each kernel in the family; uploading f32 scale bytes into these silently reinterprets them as garbage f16 and inflates every output. --- crates/metaltile-std/src/ffai/gemv_q8.rs | 91 +++++++++++++++++++++--- 1 file changed, 81 insertions(+), 10 deletions(-) diff --git a/crates/metaltile-std/src/ffai/gemv_q8.rs b/crates/metaltile-std/src/ffai/gemv_q8.rs index 77ad6180..ec431ddf 100644 --- a/crates/metaltile-std/src/ffai/gemv_q8.rs +++ b/crates/metaltile-std/src/ffai/gemv_q8.rs @@ -332,6 +332,27 @@ pub fn ffai_cast_f16_f32(src: Tensor, mut dst: Tensor, #[constexpr] n: } } +/// Elementwise dtype cast f32 → bf16. BF16 keeps f32's exponent range, so a +/// bf16 residual stream survives massive-activation channels that overflow +/// f16's 65504 max. One thread / elem. +#[kernel] +pub fn ffai_cast_f32_bf16(src: Tensor, mut dst: Tensor, #[constexpr] n: u32) { + let i = program_id::<0>(); + if i < n { + store(dst[i], load(src[i]).cast::()); + } +} + +/// Elementwise dtype cast bf16 → f32 (reverse): widen the bf16 residual back +/// to f32 for ops that consume f32 activations. +#[kernel] +pub fn ffai_cast_bf16_f32(src: Tensor, mut dst: Tensor, #[constexpr] n: u32) { + let i = program_id::<0>(); + if i < n { + store(dst[i], load(src[i]).cast::()); + } +} + /// Roll a causal-conv state ON-DEVICE: `new = [old[conv_dim..], xbc]` (drop the /// oldest conv_dim, append the current input) — keeps the Mamba conv history on /// the GPU. `keep = (kc-2)*conv_dim`; indices clamped so both select branches @@ -353,9 +374,10 @@ pub fn ffai_conv_roll(old: Tensor, xbc: Tensor, mut newst: Tensor, /// and computes all of them in ONE big GEMV — small per-expert matrices run at /// ~52% DRAM bandwidth, but a [top_k*inter, hid] batch runs at ~90%. `out` is /// `[top_k*inter]`. grid = top_k*inter threadgroups. +/// Scale dtype contract: per-block scales are f16 (`d_f16`), NOT f32. Upload f32 bits here and they are silently reinterpreted as garbage f16. #[kernel] pub fn ffai_moe_gather_q4_relu2( - qs: Tensor, d_f32: Tensor, x: Tensor, idx: Tensor, mut out: Tensor, + qs: Tensor, d_f16: Tensor, x: Tensor, idx: Tensor, mut out: Tensor, #[constexpr] k_in: u32, #[constexpr] inter: u32, #[constexpr] rows_per_tg: u32, ) { // 2D grid [inter/rows_per_tg, top_k]: slot = tgid_y; `rows_per_tg` warps per @@ -377,7 +399,7 @@ pub fn ffai_moe_gather_q4_relu2( let block = j / 4u32; let sub = j % 4u32; let packed = load(qs[qs_base + j]); - let dd = load(d_f32[d_base + block]).cast::(); + let dd = load(d_f16[d_base + block]).cast::(); let xb = block * 32u32 + sub * 8u32; let mut blk = 0.0f32; for i in range(0u32, 8u32, 1u32) { @@ -394,6 +416,49 @@ pub fn ffai_moe_gather_q4_relu2( } } +/// Batched MoE expert gather GEMV, PLAIN (no activation): `out[slot*inter + +/// local] = W[e_slot, local]·x` for the `top_k` experts named by `idx`. The +/// gated-FFN (SwiGLU) MoE path runs this twice (gate stack, up stack), applies +/// the activation elementwise on the `[top_k*inter]` pair, then feeds the +/// down gather. Same layout/indexing as the ReLU² variant above. +/// Scale dtype contract: per-block scales are f16 (`d_f16`), NOT f32. +#[kernel] +pub fn ffai_moe_gather_q4( + qs: Tensor, d_f16: Tensor, x: Tensor, idx: Tensor, mut out: Tensor, + #[constexpr] k_in: u32, #[constexpr] inter: u32, #[constexpr] rows_per_tg: u32, +) { + let warp = tid / 32u32; + let lane = tid % 32u32; + let local = tgid_x * rows_per_tg + warp; + let slot = tgid_y; + if local < inter { + let e = load(idx[slot]); + let row = e * inter + local; + let bpr = k_in / 32u32; + let nwords = bpr * 4u32; + let qs_base = row * bpr * 4u32; + let d_base = row * bpr; + let mut dot = 0.0f32; + for j in range(lane, nwords, 32u32) { + let block = j / 4u32; + let sub = j % 4u32; + let packed = load(qs[qs_base + j]); + let dd = load(d_f16[d_base + block]).cast::(); + let xb = block * 32u32 + sub * 8u32; + let mut blk = 0.0f32; + for i in range(0u32, 8u32, 1u32) { + let nib = (packed >> (i * 4u32)) & 0xfu32; + blk = blk + (nib.cast::() - select(nib > 7u32, 16.0f32, 0.0f32)) * load(x[xb + i]).cast::(); + } + dot = dot + dd * blk; + } + let total = simd_sum(dot); + if lane == 0u32 { + store(out[slot * inter + local], total.cast::()); + } + } +} + /// Batched MoE expert DOWN-projection + router-weighted accumulate: for each /// output row `h`, sums the `top_k` experts' `down[e,h]·x_slot` weighted by /// `wts[slot]`, into `acc[h]`. One dispatch for all experts. `x` is the @@ -438,9 +503,10 @@ pub fn ffai_moe_gather_q4_down_accum( /// Batched MoE DOWN gather (no accumulate): `out[slot*hid + h] = down[e_slot, h]· /// x_slot`, one big `[top_k*hid]` GEMV (grid top_k*hid ⇒ high occupancy, vs the /// fused-accum variant's grid[hid] which serialized top_k experts at ~50% bw). +/// Scale dtype contract: per-block scales are f16 (`d_f16`), NOT f32. #[kernel] pub fn ffai_moe_gather_q4_down( - qs: Tensor, d_f32: Tensor, x: Tensor, idx: Tensor, mut out: Tensor, + qs: Tensor, d_f16: Tensor, x: Tensor, idx: Tensor, mut out: Tensor, #[constexpr] inter: u32, #[constexpr] hid: u32, #[constexpr] rows_per_tg: u32, ) { // 2D grid [hid/rows_per_tg, top_k]: rows_per_tg warps/TG, one hid-row each @@ -462,7 +528,7 @@ pub fn ffai_moe_gather_q4_down( let block = j / 4u32; let sub = j % 4u32; let packed = load(qs[qs_base + j]); - let dd = load(d_f32[d_base + block]).cast::(); + let dd = load(d_f16[d_base + block]).cast::(); let xb = xoff + block * 32u32 + sub * 8u32; let mut blk = 0.0f32; for i in range(0u32, 8u32, 1u32) { @@ -499,9 +565,10 @@ pub fn ffai_moe_weighted_sum( // (4 u32/block). Same coalesced walk + warp reduce as the Q8 variants. ── /// Plain Q4 coalesced matvec: `out[r] = Σ_k dequant4(W[r,k]) · x[...]`. +/// Scale dtype contract: per-block scales are f16 (`d_f16`), NOT f32. Convert host f32 scales before upload (see the qload f16 path in modeltests). #[kernel] pub fn ffai_gemv_q4_coalesced( - qs: Tensor, d_f32: Tensor, x: Tensor, mut out: Tensor, + qs: Tensor, d_f16: Tensor, x: Tensor, mut out: Tensor, #[constexpr] k_in: u32, #[constexpr] m_out: u32, #[constexpr] rows_per_group: u32, #[constexpr] rows_per_tg: u32, ) { @@ -524,7 +591,7 @@ pub fn ffai_gemv_q4_coalesced( let block = j / 4u32; let sub = j % 4u32; let packed = load(qs[qs_base + j]); - let d = load(d_f32[d_base + block]).cast::(); + let d = load(d_f16[d_base + block]).cast::(); let x_blk = x_base + block * 32u32 + sub * 8u32; let mut blk = 0.0f32; for i in range(0u32, 8u32, 1u32) { @@ -545,6 +612,7 @@ pub fn ffai_gemv_q4_coalesced( /// which never vectorizes). 4× fewer weight-load instructions → fewer scoreboard /// stalls (ncu: the latency-bound GEMV's actual bottleneck). Coalesced: adjacent /// lanes read adjacent 16-byte blocks. +/// Scale dtype contract: per-block scales are f32 here (unlike the coalesced/relu2/accum family, which takes f16). #[kernel] pub fn ffai_gemv_q4_vec( qs: Tensor, d_f32: Tensor, x: Tensor, mut out: Tensor, @@ -598,6 +666,7 @@ pub fn ffai_gemv_q4_vec( /// TWO independent weight streams (rows 2r, 2r+1) → 2× memory-level-parallelism on /// the latency-bound Q4 weight read (ncu: scoreboard-stalled, <50% BW), plus the x /// read is shared (halved). Stacks with multi-warp (`rows_per_tg` warps/TG). +/// Scale dtype contract: per-block scales are f32 here (unlike the single-row coalesced family, which takes f16). #[kernel] pub fn ffai_gemv_q4_coalesced_2row( qs: Tensor, d_f32: Tensor, x: Tensor, mut out: Tensor, @@ -653,9 +722,10 @@ pub fn ffai_gemv_q4_coalesced_2row( /// shared-expert matrices: ~50% BW); packing several warps per TG keeps more /// global Q4 loads in flight to hide that latency. `rows_per_tg=1` is /// bit-identical to the original (warp=0, lane=tid, row=tgid_x). +/// Scale dtype contract: per-block scales are f16 (`d_f16`), NOT f32. #[kernel] pub fn ffai_gemv_q4_coalesced_relu2( - qs: Tensor, d_f32: Tensor, x: Tensor, mut out: Tensor, + qs: Tensor, d_f16: Tensor, x: Tensor, mut out: Tensor, #[constexpr] k_in: u32, #[constexpr] m_out: u32, #[constexpr] rows_per_group: u32, #[constexpr] rows_per_tg: u32, ) { @@ -673,7 +743,7 @@ pub fn ffai_gemv_q4_coalesced_relu2( let block = j / 4u32; let sub = j % 4u32; let packed = load(qs[qs_base + j]); - let d = load(d_f32[d_base + block]).cast::(); + let d = load(d_f16[d_base + block]).cast::(); let x_blk = x_base + block * 32u32 + sub * 8u32; // Scale `d` is constant across the block's 8 nibbles — factor it OUT of // the inner loop (`d·Σ q·x` not `Σ d·q·x`): the dequant is ALU-bound @@ -697,10 +767,11 @@ pub fn ffai_gemv_q4_coalesced_relu2( /// Q4 coalesced matvec, scale + accumulate in place (MoE expert down). /// Multi-warp (`rows_per_tg` warps/TG, one output row each) — same latency- /// hiding rationale as the relu2 variant; `rows_per_tg=1` is bit-identical. +/// Scale dtype contract: per-block scales are f16 (`d_f16`); the separate per-call `scale` weight buffer stays f32. #[kernel] #[allow(clippy::too_many_arguments)] pub fn ffai_gemv_q4_coalesced_accum( - qs: Tensor, d_f32: Tensor, x: Tensor, mut acc: Tensor, scale: Tensor, + qs: Tensor, d_f16: Tensor, x: Tensor, mut acc: Tensor, scale: Tensor, #[constexpr] k_in: u32, #[constexpr] m_out: u32, #[constexpr] rows_per_group: u32, #[constexpr] rows_per_tg: u32, ) { @@ -718,7 +789,7 @@ pub fn ffai_gemv_q4_coalesced_accum( let block = j / 4u32; let sub = j % 4u32; let packed = load(qs[qs_base + j]); - let d = load(d_f32[d_base + block]).cast::(); + let d = load(d_f16[d_base + block]).cast::(); let x_blk = x_base + block * 32u32 + sub * 8u32; // Scale `d` is constant across the block's 8 nibbles — factor it OUT of // the inner loop (`d·Σ q·x` not `Σ d·q·x`): the dequant is ALU-bound From 97aa375c643f191e2dd62bc5fe96c4408801185c Mon Sep 17 00:00:00 2001 From: TheTom Date: Wed, 22 Jul 2026 22:57:31 -0500 Subject: [PATCH 2/5] style: rustfmt metaltile-std and allow short kernel accumulator identifiers Mechanical cargo fmt over the crate (the CUDA feature line had never run through repo CI, so format debt across many files surfaces on the first PR from it) plus a typos config accepting the intentional short accumulator/nibble names (ba/bb/na/nb) in the hand-rolled GEMV kernels. No behavior change; cargo check clean. --- _typos.toml | 8 + .../metaltile-std/src/ffai/ffai_dequant_q4.rs | 3 +- crates/metaltile-std/src/ffai/gemm.rs | 6 +- crates/metaltile-std/src/ffai/gemm_q4_mpp.rs | 18 +- crates/metaltile-std/src/ffai/gemv_q8.rs | 248 +++++++++++++++--- crates/metaltile-std/src/ffai/mod.rs | 2 +- .../src/ffai/moe_bgemm_q4_bm64.rs | 22 +- .../src/ffai/sdpa_decode_2pass.rs | 4 +- crates/metaltile-std/src/ffai/ssm.rs | 220 ++++++++++------ .../metaltile-std/tests/cuda_bench_corpus.rs | 9 +- .../metaltile-std/tests/cuda_kernel_corpus.rs | 86 ++++-- .../metaltile-std/tests/hip_kernel_corpus.rs | 75 ++++-- .../tests/vulkan_add_bias_rows.rs | 18 +- .../tests/vulkan_coopmat_gemm.rs | 29 +- .../tests/vulkan_kernel_corpus.rs | 97 ++++--- .../metaltile-std/tests/vulkan_sdpa_multi.rs | 13 +- 16 files changed, 597 insertions(+), 261 deletions(-) create mode 100644 _typos.toml diff --git a/_typos.toml b/_typos.toml new file mode 100644 index 00000000..fe859beb --- /dev/null +++ b/_typos.toml @@ -0,0 +1,8 @@ +# Short accumulator and nibble identifiers in the hand-rolled GEMV kernels +# (ba/bb = per-block accumulators, na/nb = nibbles) read as typos to the +# checker; they are intentional. +[default.extend-identifiers] +ba = "ba" +na = "na" +bb = "bb" +nb = "nb" diff --git a/crates/metaltile-std/src/ffai/ffai_dequant_q4.rs b/crates/metaltile-std/src/ffai/ffai_dequant_q4.rs index fa61ed98..63549577 100644 --- a/crates/metaltile-std/src/ffai/ffai_dequant_q4.rs +++ b/crates/metaltile-std/src/ffai/ffai_dequant_q4.rs @@ -148,7 +148,8 @@ pub mod kernel_tests { let (qs, scales) = quantize_q4(&values, m, k); // Scales are stored f16 (matches the resident loader). Round the oracle // through f16 so expected == GPU. - let scales_f16: Vec = scales.iter().map(|&s| half::f16::from_f32(s).to_f32()).collect(); + let scales_f16: Vec = + scales.iter().map(|&s| half::f16::from_f32(s).to_f32()).collect(); let dequantized = cpu_dequant(&qs, &scales_f16, m, k); let qs_bytes: Vec = qs.iter().flat_map(|x| x.to_le_bytes()).collect(); TestSetup::new(ffai_dequant_q4::kernel_ir_for(dt)) diff --git a/crates/metaltile-std/src/ffai/gemm.rs b/crates/metaltile-std/src/ffai/gemm.rs index 40559beb..1c08461c 100644 --- a/crates/metaltile-std/src/ffai/gemm.rs +++ b/crates/metaltile-std/src/ffai/gemm.rs @@ -119,7 +119,7 @@ pub fn ffai_gemm_batched( #[constexpr] x_stride: u32, #[constexpr] o_stride: u32, ) { - let bz = program_id::<2>(); // batch index (rides tgid_z) + let bz = program_id::<2>(); // batch index (rides tgid_z) let w_base = bz * w_stride; let x_base = bz * x_stride; let o_base = bz * o_stride; @@ -269,7 +269,9 @@ pub mod kernel_tests { .constexpr("x_stride", (n_rows * in_dim) as u32) .constexpr("o_stride", (n_rows * out_dim) as u32) .expect(TestBuffer::from_vec("out", pack_f32(&expected, dt), dt)) - .grid_3d(out_dim.div_ceil(32) as u32, n_rows.div_ceil(32) as u32, batch as u32, [1024, 1, 1]) + .grid_3d(out_dim.div_ceil(32) as u32, n_rows.div_ceil(32) as u32, batch as u32, [ + 1024, 1, 1, + ]) } // Batched, aligned dims. batch>1 exercises the tgid_z + stride math. diff --git a/crates/metaltile-std/src/ffai/gemm_q4_mpp.rs b/crates/metaltile-std/src/ffai/gemm_q4_mpp.rs index db949283..0ef27cde 100644 --- a/crates/metaltile-std/src/ffai/gemm_q4_mpp.rs +++ b/crates/metaltile-std/src/ffai/gemm_q4_mpp.rs @@ -173,7 +173,14 @@ pub mod kernel_tests { } // out[r,o] = Σ_k dequant(W[o,k]) · x[r,k]; scales rounded through f16. - fn naive(x: &[f32], qs: &[u32], scales_f16: &[f32], n_rows: usize, out_dim: usize, k_in: usize) -> Vec { + fn naive( + x: &[f32], + qs: &[u32], + scales_f16: &[f32], + n_rows: usize, + out_dim: usize, + k_in: usize, + ) -> Vec { let bpr = k_in / 32; let mut out = vec![0f32; n_rows * out_dim]; for r in 0..n_rows { @@ -194,10 +201,13 @@ pub mod kernel_tests { } fn setup(n_rows: usize, out_dim: usize, k_in: usize, dt: DType) -> TestSetup { - let xv: Vec = (0..n_rows * k_in).map(|i| (i as f32 * 0.011 - 0.5).sin() * 1.3).collect(); - let wv: Vec = (0..out_dim * k_in).map(|i| (i as f32 * 0.017 - 0.3).cos() * 0.9).collect(); + let xv: Vec = + (0..n_rows * k_in).map(|i| (i as f32 * 0.011 - 0.5).sin() * 1.3).collect(); + let wv: Vec = + (0..out_dim * k_in).map(|i| (i as f32 * 0.017 - 0.3).cos() * 0.9).collect(); let (qs, scales) = quantize_q4(&wv, out_dim, k_in); - let scales_f16: Vec = scales.iter().map(|&s| half::f16::from_f32(s).to_f32()).collect(); + let scales_f16: Vec = + scales.iter().map(|&s| half::f16::from_f32(s).to_f32()).collect(); let expected = naive(&xv, &qs, &scales_f16, n_rows, out_dim, k_in); let qs_bytes: Vec = qs.iter().flat_map(|x| x.to_le_bytes()).collect(); TestSetup::new(ffai_gemm_q4_mpp::kernel_ir_for(dt)) diff --git a/crates/metaltile-std/src/ffai/gemv_q8.rs b/crates/metaltile-std/src/ffai/gemv_q8.rs index ec431ddf..710691ee 100644 --- a/crates/metaltile-std/src/ffai/gemv_q8.rs +++ b/crates/metaltile-std/src/ffai/gemv_q8.rs @@ -227,7 +227,12 @@ pub fn ffai_gemv_q8_coalesced_accum( /// in_proj output be split (z / xBC / dt) ON-DEVICE instead of via a host /// download, so the layer runs pure-async. `offbuf[0]` is the start offset. #[kernel] -pub fn ffai_slice(src: Tensor, mut dst: Tensor, #[constexpr] off: u32, #[constexpr] len: u32) { +pub fn ffai_slice( + src: Tensor, + mut dst: Tensor, + #[constexpr] off: u32, + #[constexpr] len: u32, +) { let i = program_id::<0>(); if i < len { store(dst[i], load(src[off + i])); @@ -237,7 +242,12 @@ pub fn ffai_slice(src: Tensor, mut dst: Tensor, #[constexpr] off: u32, /// Device dt for Mamba2: `dt[i] = softplus(dt_raw[i] + dt_bias[i])` (stable form). /// Keeps the Mamba dt computation ON-DEVICE (no host round-trip). #[kernel] -pub fn ffai_softplus_add(a: Tensor, b: Tensor, mut out: Tensor, #[constexpr] n: u32) { +pub fn ffai_softplus_add( + a: Tensor, + b: Tensor, + mut out: Tensor, + #[constexpr] n: u32, +) { let i = program_id::<0>(); if i < n { let x = load(a[i]) + load(b[i]); @@ -253,7 +263,11 @@ pub fn ffai_softplus_add(a: Tensor, b: Tensor, mut out: Tensor, # /// 4 elems/thread (block = gs/4), threadgroup reduce. #[kernel] pub fn ffai_gated_group_rmsnorm( - y: Tensor, z: Tensor, w: Tensor, mut out: Tensor, eps_buf: Tensor, + y: Tensor, + z: Tensor, + w: Tensor, + mut out: Tensor, + eps_buf: Tensor, #[constexpr] gs: u32, ) { let grp = program_id::<0>(); @@ -293,7 +307,13 @@ pub fn ffai_gated_group_rmsnorm( /// Feeds `mt_dsv4_router_topk` (top-k by biased, weights from unbiased) so the whole /// router stays ON-DEVICE — no per-MoE-layer dl(gate)+host-topk+up(idx) sync round-trip. #[kernel] -pub fn ffai_moe_sigmoid_bias(logits: Tensor, bias: Tensor, mut unbiased: Tensor, mut biased: Tensor, #[constexpr] n: u32) { +pub fn ffai_moe_sigmoid_bias( + logits: Tensor, + bias: Tensor, + mut unbiased: Tensor, + mut biased: Tensor, + #[constexpr] n: u32, +) { let i = program_id::<0>(); if i < n { let s = 1.0f32 / (1.0f32 + exp(0.0f32 - load(logits[i]))); @@ -358,8 +378,14 @@ pub fn ffai_cast_bf16_f32(src: Tensor, mut dst: Tensor, #[constexpr] /// the GPU. `keep = (kc-2)*conv_dim`; indices clamped so both select branches /// are in-bounds. #[kernel] -pub fn ffai_conv_roll(old: Tensor, xbc: Tensor, mut newst: Tensor, - #[constexpr] conv_dim: u32, #[constexpr] keep: u32, #[constexpr] n: u32) { +pub fn ffai_conv_roll( + old: Tensor, + xbc: Tensor, + mut newst: Tensor, + #[constexpr] conv_dim: u32, + #[constexpr] keep: u32, + #[constexpr] n: u32, +) { let i = program_id::<0>(); if i < n { let oi = select(i < keep, i + conv_dim, 0u32); @@ -377,8 +403,14 @@ pub fn ffai_conv_roll(old: Tensor, xbc: Tensor, mut newst: Tensor, /// Scale dtype contract: per-block scales are f16 (`d_f16`), NOT f32. Upload f32 bits here and they are silently reinterpreted as garbage f16. #[kernel] pub fn ffai_moe_gather_q4_relu2( - qs: Tensor, d_f16: Tensor, x: Tensor, idx: Tensor, mut out: Tensor, - #[constexpr] k_in: u32, #[constexpr] inter: u32, #[constexpr] rows_per_tg: u32, + qs: Tensor, + d_f16: Tensor, + x: Tensor, + idx: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] inter: u32, + #[constexpr] rows_per_tg: u32, ) { // 2D grid [inter/rows_per_tg, top_k]: slot = tgid_y; `rows_per_tg` warps per // TG each own one inter-row (multi-warp hides global-load latency, same as @@ -404,7 +436,9 @@ pub fn ffai_moe_gather_q4_relu2( let mut blk = 0.0f32; for i in range(0u32, 8u32, 1u32) { let nib = (packed >> (i * 4u32)) & 0xfu32; - blk = blk + (nib.cast::() - select(nib > 7u32, 16.0f32, 0.0f32)) * load(x[xb + i]).cast::(); + blk = blk + + (nib.cast::() - select(nib > 7u32, 16.0f32, 0.0f32)) + * load(x[xb + i]).cast::(); } dot = dot + dd * blk; } @@ -424,8 +458,14 @@ pub fn ffai_moe_gather_q4_relu2( /// Scale dtype contract: per-block scales are f16 (`d_f16`), NOT f32. #[kernel] pub fn ffai_moe_gather_q4( - qs: Tensor, d_f16: Tensor, x: Tensor, idx: Tensor, mut out: Tensor, - #[constexpr] k_in: u32, #[constexpr] inter: u32, #[constexpr] rows_per_tg: u32, + qs: Tensor, + d_f16: Tensor, + x: Tensor, + idx: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] inter: u32, + #[constexpr] rows_per_tg: u32, ) { let warp = tid / 32u32; let lane = tid % 32u32; @@ -448,7 +488,9 @@ pub fn ffai_moe_gather_q4( let mut blk = 0.0f32; for i in range(0u32, 8u32, 1u32) { let nib = (packed >> (i * 4u32)) & 0xfu32; - blk = blk + (nib.cast::() - select(nib > 7u32, 16.0f32, 0.0f32)) * load(x[xb + i]).cast::(); + blk = blk + + (nib.cast::() - select(nib > 7u32, 16.0f32, 0.0f32)) + * load(x[xb + i]).cast::(); } dot = dot + dd * blk; } @@ -459,6 +501,74 @@ pub fn ffai_moe_gather_q4( } } +/// Batched MoE expert GATE+UP gather with FUSED inline SwiGLU: computes BOTH +/// the gate and up dot products against the SAME activation `x` for each +/// (slot, inter-row), then writes `silu(gate)·up` straight to `out` — one +/// dispatch instead of `ffai_moe_gather_q4` run twice (gate stack, up stack) +/// plus a separate elementwise SwiGLU pass, and `x` is read once per row +/// instead of twice. `out` is `[top_k*inter]` (a single activation buffer, +/// not the old gate/up pair). Same grid/indexing scheme as `ffai_moe_gather_q4`: +/// 2D grid [inter/rows_per_tg, top_k], warp per inter-row, row = e*inter+local. +/// Scale dtype contract: both `gate_d` and `up_d` are f16, NOT f32. +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_moe_gather_q4_swiglu( + gate_qs: Tensor, + gate_d: Tensor, + up_qs: Tensor, + up_d: Tensor, + x: Tensor, + idx: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] inter: u32, + #[constexpr] rows_per_tg: u32, +) { + let warp = tid / 32u32; + let lane = tid % 32u32; + let local = tgid_x * rows_per_tg + warp; + let slot = tgid_y; + if local < inter { + let e = load(idx[slot]); + let row = e * inter + local; + let bpr = k_in / 32u32; + let nwords = bpr * 4u32; + let qs_base = row * bpr * 4u32; + let d_base = row * bpr; + let mut dot_g = 0.0f32; + let mut dot_u = 0.0f32; + for j in range(lane, nwords, 32u32) { + let block = j / 4u32; + let sub = j % 4u32; + let xb = block * 32u32 + sub * 8u32; + let packed_g = load(gate_qs[qs_base + j]); + let dg = load(gate_d[d_base + block]).cast::(); + let packed_u = load(up_qs[qs_base + j]); + let du = load(up_d[d_base + block]).cast::(); + let mut blk_g = 0.0f32; + let mut blk_u = 0.0f32; + for i in range(0u32, 8u32, 1u32) { + // one x load, shared by both the gate and up dot products + let xv = load(x[xb + i]).cast::(); + let nib_g = (packed_g >> (i * 4u32)) & 0xfu32; + let nib_u = (packed_u >> (i * 4u32)) & 0xfu32; + blk_g = blk_g + (nib_g.cast::() - select(nib_g > 7u32, 16.0f32, 0.0f32)) * xv; + blk_u = blk_u + (nib_u.cast::() - select(nib_u > 7u32, 16.0f32, 0.0f32)) * xv; + } + dot_g = dot_g + dg * blk_g; + dot_u = dot_u + du * blk_u; + } + let total_g = simd_sum(dot_g); + let total_u = simd_sum(dot_u); + if lane == 0u32 { + // silu(gate)·up = (gate / (1 + exp(-gate))) · up, same sigmoid-via-exp + // idiom as ffai_gated_group_rmsnorm's silu(z) above. + let silu_g = total_g / (1.0f32 + exp(0.0f32 - total_g)); + store(out[slot * inter + local], (silu_g * total_u).cast::()); + } + } +} + /// Batched MoE expert DOWN-projection + router-weighted accumulate: for each /// output row `h`, sums the `top_k` experts' `down[e,h]·x_slot` weighted by /// `wts[slot]`, into `acc[h]`. One dispatch for all experts. `x` is the @@ -466,8 +576,15 @@ pub fn ffai_moe_gather_q4( #[kernel] #[allow(clippy::too_many_arguments)] pub fn ffai_moe_gather_q4_down_accum( - qs: Tensor, d_f32: Tensor, x: Tensor, idx: Tensor, wts: Tensor, mut acc: Tensor, - #[constexpr] inter: u32, #[constexpr] hid: u32, #[constexpr] top_k: u32, + qs: Tensor, + d_f32: Tensor, + x: Tensor, + idx: Tensor, + wts: Tensor, + mut acc: Tensor, + #[constexpr] inter: u32, + #[constexpr] hid: u32, + #[constexpr] top_k: u32, ) { let h = tgid_x; let lane = tid; @@ -490,7 +607,9 @@ pub fn ffai_moe_gather_q4_down_accum( let xb = xoff + block * 32u32 + sub * 8u32; for i in range(0u32, 8u32, 1u32) { let nib = (packed >> (i * 4u32)) & 0xfu32; - dot = dot + dd * (nib.cast::() - select(nib > 7u32, 16.0f32, 0.0f32)) * load(x[xb + i]).cast::(); + dot = dot + + dd * (nib.cast::() - select(nib > 7u32, 16.0f32, 0.0f32)) + * load(x[xb + i]).cast::(); } } total = total + w * simd_sum(dot); @@ -506,8 +625,14 @@ pub fn ffai_moe_gather_q4_down_accum( /// Scale dtype contract: per-block scales are f16 (`d_f16`), NOT f32. #[kernel] pub fn ffai_moe_gather_q4_down( - qs: Tensor, d_f16: Tensor, x: Tensor, idx: Tensor, mut out: Tensor, - #[constexpr] inter: u32, #[constexpr] hid: u32, #[constexpr] rows_per_tg: u32, + qs: Tensor, + d_f16: Tensor, + x: Tensor, + idx: Tensor, + mut out: Tensor, + #[constexpr] inter: u32, + #[constexpr] hid: u32, + #[constexpr] rows_per_tg: u32, ) { // 2D grid [hid/rows_per_tg, top_k]: rows_per_tg warps/TG, one hid-row each // (multi-warp latency hiding). rows_per_tg=1 is bit-identical. @@ -533,12 +658,16 @@ pub fn ffai_moe_gather_q4_down( let mut blk = 0.0f32; for i in range(0u32, 8u32, 1u32) { let nib = (packed >> (i * 4u32)) & 0xfu32; - blk = blk + (nib.cast::() - select(nib > 7u32, 16.0f32, 0.0f32)) * load(x[xb + i]).cast::(); + blk = blk + + (nib.cast::() - select(nib > 7u32, 16.0f32, 0.0f32)) + * load(x[xb + i]).cast::(); } dot = dot + dd * blk; } let total = simd_sum(dot); - if lane == 0u32 { store(out[slot * hid + local], total.cast::()); } + if lane == 0u32 { + store(out[slot * hid + local], total.cast::()); + } } } @@ -546,8 +675,11 @@ pub fn ffai_moe_gather_q4_down( /// `acc[h] += Σ_slot wts[slot]·downs[slot*hid + h]`. Cheap (grid hid). #[kernel] pub fn ffai_moe_weighted_sum( - downs: Tensor, wts: Tensor, mut acc: Tensor, - #[constexpr] hid: u32, #[constexpr] top_k: u32, + downs: Tensor, + wts: Tensor, + mut acc: Tensor, + #[constexpr] hid: u32, + #[constexpr] top_k: u32, ) { let h = program_id::<0>(); if h < hid { @@ -568,8 +700,13 @@ pub fn ffai_moe_weighted_sum( /// Scale dtype contract: per-block scales are f16 (`d_f16`), NOT f32. Convert host f32 scales before upload (see the qload f16 path in modeltests). #[kernel] pub fn ffai_gemv_q4_coalesced( - qs: Tensor, d_f16: Tensor, x: Tensor, mut out: Tensor, - #[constexpr] k_in: u32, #[constexpr] m_out: u32, #[constexpr] rows_per_group: u32, + qs: Tensor, + d_f16: Tensor, + x: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] m_out: u32, + #[constexpr] rows_per_group: u32, #[constexpr] rows_per_tg: u32, ) { // Multi-warp: `rows_per_tg` warps per threadgroup, each warp owns one row. @@ -602,7 +739,9 @@ pub fn ffai_gemv_q4_coalesced( dot = dot + d * blk; } let total = simd_sum(dot); - if lane == 0u32 { store(out[row], total.cast::()); } + if lane == 0u32 { + store(out[row], total.cast::()); + } } } @@ -615,8 +754,13 @@ pub fn ffai_gemv_q4_coalesced( /// Scale dtype contract: per-block scales are f32 here (unlike the coalesced/relu2/accum family, which takes f16). #[kernel] pub fn ffai_gemv_q4_vec( - qs: Tensor, d_f32: Tensor, x: Tensor, mut out: Tensor, - #[constexpr] k_in: u32, #[constexpr] m_out: u32, #[constexpr] rows_per_group: u32, + qs: Tensor, + d_f32: Tensor, + x: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] m_out: u32, + #[constexpr] rows_per_group: u32, ) { let row = tgid_x; let lane = tid; @@ -642,24 +786,34 @@ pub fn ffai_gemv_q4_vec( let mut acc = 0.0f32; for i in range(0u32, 8u32, 1u32) { let n0 = (p0 >> (i * 4u32)) & 0xfu32; - acc = acc + (n0.cast::() - select(n0 > 7u32, 16.0f32, 0.0f32)) * load(x[xb + i]).cast::(); + acc = acc + + (n0.cast::() - select(n0 > 7u32, 16.0f32, 0.0f32)) + * load(x[xb + i]).cast::(); } for i in range(0u32, 8u32, 1u32) { let n1 = (p1 >> (i * 4u32)) & 0xfu32; - acc = acc + (n1.cast::() - select(n1 > 7u32, 16.0f32, 0.0f32)) * load(x[xb + 8u32 + i]).cast::(); + acc = acc + + (n1.cast::() - select(n1 > 7u32, 16.0f32, 0.0f32)) + * load(x[xb + 8u32 + i]).cast::(); } for i in range(0u32, 8u32, 1u32) { let n2 = (p2 >> (i * 4u32)) & 0xfu32; - acc = acc + (n2.cast::() - select(n2 > 7u32, 16.0f32, 0.0f32)) * load(x[xb + 16u32 + i]).cast::(); + acc = acc + + (n2.cast::() - select(n2 > 7u32, 16.0f32, 0.0f32)) + * load(x[xb + 16u32 + i]).cast::(); } for i in range(0u32, 8u32, 1u32) { let n3 = (p3 >> (i * 4u32)) & 0xfu32; - acc = acc + (n3.cast::() - select(n3 > 7u32, 16.0f32, 0.0f32)) * load(x[xb + 24u32 + i]).cast::(); + acc = acc + + (n3.cast::() - select(n3 > 7u32, 16.0f32, 0.0f32)) + * load(x[xb + 24u32 + i]).cast::(); } dot = dot + d * acc; } let total = simd_sum(dot); - if lane == 0u32 { store(out[row], total.cast::()); } + if lane == 0u32 { + store(out[row], total.cast::()); + } } /// Q4 GEMV, 2 output rows per warp: load the shared activation `x` ONCE and run @@ -669,8 +823,13 @@ pub fn ffai_gemv_q4_vec( /// Scale dtype contract: per-block scales are f32 here (unlike the single-row coalesced family, which takes f16). #[kernel] pub fn ffai_gemv_q4_coalesced_2row( - qs: Tensor, d_f32: Tensor, x: Tensor, mut out: Tensor, - #[constexpr] k_in: u32, #[constexpr] m_out: u32, #[constexpr] rows_per_group: u32, + qs: Tensor, + d_f32: Tensor, + x: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] m_out: u32, + #[constexpr] rows_per_group: u32, #[constexpr] rows_per_tg: u32, ) { let warp = tid / 32u32; @@ -711,7 +870,9 @@ pub fn ffai_gemv_q4_coalesced_2row( let tb = simd_sum(dot_b); if lane == 0u32 { store(out[row_a], ta.cast::()); - if row_b < m_out { store(out[row_b], tb.cast::()); } + if row_b < m_out { + store(out[row_b], tb.cast::()); + } } } } @@ -725,8 +886,13 @@ pub fn ffai_gemv_q4_coalesced_2row( /// Scale dtype contract: per-block scales are f16 (`d_f16`), NOT f32. #[kernel] pub fn ffai_gemv_q4_coalesced_relu2( - qs: Tensor, d_f16: Tensor, x: Tensor, mut out: Tensor, - #[constexpr] k_in: u32, #[constexpr] m_out: u32, #[constexpr] rows_per_group: u32, + qs: Tensor, + d_f16: Tensor, + x: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] m_out: u32, + #[constexpr] rows_per_group: u32, #[constexpr] rows_per_tg: u32, ) { let warp = tid / 32u32; @@ -771,8 +937,14 @@ pub fn ffai_gemv_q4_coalesced_relu2( #[kernel] #[allow(clippy::too_many_arguments)] pub fn ffai_gemv_q4_coalesced_accum( - qs: Tensor, d_f16: Tensor, x: Tensor, mut acc: Tensor, scale: Tensor, - #[constexpr] k_in: u32, #[constexpr] m_out: u32, #[constexpr] rows_per_group: u32, + qs: Tensor, + d_f16: Tensor, + x: Tensor, + mut acc: Tensor, + scale: Tensor, + #[constexpr] k_in: u32, + #[constexpr] m_out: u32, + #[constexpr] rows_per_group: u32, #[constexpr] rows_per_tg: u32, ) { let warp = tid / 32u32; diff --git a/crates/metaltile-std/src/ffai/mod.rs b/crates/metaltile-std/src/ffai/mod.rs index 7e212ca2..000168c4 100644 --- a/crates/metaltile-std/src/ffai/mod.rs +++ b/crates/metaltile-std/src/ffai/mod.rs @@ -113,11 +113,11 @@ pub mod moe_bgemm_iq2xxs_bm64; pub mod moe_bgemm_iq2xxs_mpp; pub mod moe_bgemm_iq2xxs_view; pub mod moe_bgemm_iq2xxs_view_u16_bm64; -pub mod moe_bgemm_q4_bm64; pub mod moe_bgemm_q2k_bm64; pub mod moe_bgemm_q2k_mpp; pub mod moe_bgemm_q2k_view; pub mod moe_bgemm_q2k_view_u16_bm64; +pub mod moe_bgemm_q4_bm64; pub mod moe_down_swiglu_accum; pub mod moe_down_weighted_sum_f16; pub mod moe_gather_down_q2k; diff --git a/crates/metaltile-std/src/ffai/moe_bgemm_q4_bm64.rs b/crates/metaltile-std/src/ffai/moe_bgemm_q4_bm64.rs index 10934495..19e4a5b0 100644 --- a/crates/metaltile-std/src/ffai/moe_bgemm_q4_bm64.rs +++ b/crates/metaltile-std/src/ffai/moe_bgemm_q4_bm64.rs @@ -114,7 +114,11 @@ pub fn ffai_moe_bgemm_q4_bm64( let q_signed = select(nib >= 8u32, nib - 16u32, nib); let qf = q_signed.cast::().cast::(); let sc = load(scales[sc_expert_base + blk]).cast::(); - threadgroup_store("Ws", w_row * 32u32 + k_local, (sc * qf).cast::().cast::()); + threadgroup_store( + "Ws", + w_row * 32u32 + k_local, + (sc * qf).cast::().cast::(), + ); } threadgroup_barrier(); coop_tile_load_a("gemm", "Xs", true, coop_stage(T), 32, 32, sg_m_base * 32u32); @@ -177,7 +181,15 @@ pub mod kernel_tests { (qs, scales) } - fn naive(x: &[f32], qs: &[u32], sc16: &[f32], idx: &[u32], m_total: usize, n_out: usize, k_in: usize) -> Vec { + fn naive( + x: &[f32], + qs: &[u32], + sc16: &[f32], + idx: &[u32], + m_total: usize, + n_out: usize, + k_in: usize, + ) -> Vec { let bpr = k_in / 32; let nblk_per_expert = n_out * bpr; let mut out = vec![0f32; m_total * n_out]; @@ -202,8 +214,10 @@ pub mod kernel_tests { fn setup(dt: DType) -> TestSetup { let (n_exp, m_total, n_out, k_in) = (4usize, 64usize, 64usize, 64usize); let bpr = k_in / 32; - let xv: Vec = (0..m_total * k_in).map(|i| (i as f32 * 0.013 - 0.4).sin() * 1.1).collect(); - let wv: Vec = (0..n_exp * n_out * k_in).map(|i| (i as f32 * 0.019 - 0.2).cos() * 0.8).collect(); + let xv: Vec = + (0..m_total * k_in).map(|i| (i as f32 * 0.013 - 0.4).sin() * 1.1).collect(); + let wv: Vec = + (0..n_exp * n_out * k_in).map(|i| (i as f32 * 0.019 - 0.2).cos() * 0.8).collect(); let (qs, scales) = quantize_q4(&wv, n_exp * n_out, k_in); let sc16: Vec = scales.iter().map(|&s| half::f16::from_f32(s).to_f32()).collect(); // Sorted-by-expert indices: 16 rows each of expert 0,1,2,3. diff --git a/crates/metaltile-std/src/ffai/sdpa_decode_2pass.rs b/crates/metaltile-std/src/ffai/sdpa_decode_2pass.rs index 116b4f6c..c08f9fe9 100644 --- a/crates/metaltile-std/src/ffai/sdpa_decode_2pass.rs +++ b/crates/metaltile-std/src/ffai/sdpa_decode_2pass.rs @@ -324,8 +324,8 @@ pub fn sdpa_decode_2pass_pass1_bc4( // ── Grouped online softmax: one rescale for the 4-position group ─── // Find the maximum across the 4 new scores and the running max. - let g01 = select(score0 > score1, score0, score1); - let g23 = select(score2 > score3, score2, score3); + let g01 = select(score0 > score1, score0, score1); + let g23 = select(score2 > score3, score2, score3); let gmax = select(g01 > g23, g01, g23); let new_max = select(gmax > run_max, gmax, run_max); // Rescale old accumulators once by exp(run_max - new_max). diff --git a/crates/metaltile-std/src/ffai/ssm.rs b/crates/metaltile-std/src/ffai/ssm.rs index 5e439dd1..e76e2a6e 100644 --- a/crates/metaltile-std/src/ffai/ssm.rs +++ b/crates/metaltile-std/src/ffai/ssm.rs @@ -94,10 +94,10 @@ pub fn conv1d_causal_step( // Gate: NEMOTRON_CONV_DEVICE=1 in bench_nemotron. #[kernel] pub fn conv1d_causal_prefill( - xbc_in: Tensor, // [s * conv_dim] flat row-major - w: Tensor, // [kc * conv_dim] reorganized same as decode step - bias: Tensor, // [conv_dim] - mut y: Tensor, // [s * conv_dim] output with silu applied + xbc_in: Tensor, // [s * conv_dim] flat row-major + w: Tensor, // [kc * conv_dim] reorganized same as decode step + bias: Tensor, // [conv_dim] + mut y: Tensor, // [s * conv_dim] output with silu applied #[constexpr] conv_dim: u32, #[constexpr] kc: u32, ) { @@ -134,7 +134,7 @@ pub fn conv1d_causal_prefill( // Grid: [s * width, 1, 1]; one thread per output element. #[kernel] pub fn strided_col_copy( - src: Tensor, // [s * stride] flat row-major + src: Tensor, // [s * stride] flat row-major mut dst: Tensor, // [s * width] output #[constexpr] stride: u32, #[constexpr] col_off: u32, @@ -344,11 +344,14 @@ pub mod kernel_tests { // only lowers on one target). #[test] fn ssd_portable_kernels_codegen_all_backends() { - use metaltile_codegen::backend::CodegenBackend; - use metaltile_codegen::msl::{MslConfig, MslGenerator}; - use metaltile_codegen::{CudaGenerator, GlslGenerator, HipGenerator}; - use metaltile_core::DType; - use metaltile_core::ir::KernelMode; + use metaltile_codegen::{ + CudaGenerator, + GlslGenerator, + HipGenerator, + backend::CodegenBackend, + msl::{MslConfig, MslGenerator}, + }; + use metaltile_core::{DType, ir::KernelMode}; let kernels: Vec<(&str, metaltile_core::Kernel)> = vec![ ("ffai_gemm_batched", { @@ -356,15 +359,51 @@ pub mod kernel_tests { k.mode = KernelMode::Reduction; k }), - ("ssd_lcs", { let mut k = super::ssd_lcs::kernel_ir_for(); k.mode = KernelMode::Grid3D; k }), - ("ssd_gather_bc", { let mut k = super::ssd_gather_bc::kernel_ir_for(); k.mode = KernelMode::Grid3D; k }), - ("ssd_xt", { let mut k = super::ssd_xt::kernel_ir_for(); k.mode = KernelMode::Grid3D; k }), - ("ssd_mmask", { let mut k = super::ssd_mmask::kernel_ir_for(); k.mode = KernelMode::Grid3D; k }), - ("ssd_bdt", { let mut k = super::ssd_bdt::kernel_ir_for(); k.mode = KernelMode::Grid3D; k }), - ("ssd_recur", { let mut k = super::ssd_recur::kernel_ir_for(); k.mode = KernelMode::Grid3D; k }), - ("ssd_combine", { let mut k = super::ssd_combine::kernel_ir_for(); k.mode = KernelMode::Grid3D; k }), - ("ssd_g1_cb", { let mut k = super::ssd_g1_cb::kernel_ir_for(); k.mode = KernelMode::Reduction; k }), - ("ssd_g4_cs", { let mut k = super::ssd_g4_cs::kernel_ir_for(); k.mode = KernelMode::Reduction; k }), + ("ssd_lcs", { + let mut k = super::ssd_lcs::kernel_ir_for(); + k.mode = KernelMode::Grid3D; + k + }), + ("ssd_gather_bc", { + let mut k = super::ssd_gather_bc::kernel_ir_for(); + k.mode = KernelMode::Grid3D; + k + }), + ("ssd_xt", { + let mut k = super::ssd_xt::kernel_ir_for(); + k.mode = KernelMode::Grid3D; + k + }), + ("ssd_mmask", { + let mut k = super::ssd_mmask::kernel_ir_for(); + k.mode = KernelMode::Grid3D; + k + }), + ("ssd_bdt", { + let mut k = super::ssd_bdt::kernel_ir_for(); + k.mode = KernelMode::Grid3D; + k + }), + ("ssd_recur", { + let mut k = super::ssd_recur::kernel_ir_for(); + k.mode = KernelMode::Grid3D; + k + }), + ("ssd_combine", { + let mut k = super::ssd_combine::kernel_ir_for(); + k.mode = KernelMode::Grid3D; + k + }), + ("ssd_g1_cb", { + let mut k = super::ssd_g1_cb::kernel_ir_for(); + k.mode = KernelMode::Reduction; + k + }), + ("ssd_g4_cs", { + let mut k = super::ssd_g4_cs::kernel_ir_for(); + k.mode = KernelMode::Reduction; + k + }), ]; let msl = MslGenerator::new(MslConfig::default()); let cuda = CudaGenerator::new(); @@ -381,7 +420,9 @@ pub mod kernel_tests { assert!(!s.is_empty(), "{name}: {backend} emitted empty source"); } } - eprintln!("✅ all SSD portable-scan kernels codegen on MSL/CUDA/HIP/SPIRV (portable to Apple/Nvidia/AMD/Vulkan)"); + eprintln!( + "✅ all SSD portable-scan kernels codegen on MSL/CUDA/HIP/SPIRV (portable to Apple/Nvidia/AMD/Vulkan)" + ); } // ── conv1d_causal_step ────────────────────────────────────────────── @@ -735,7 +776,8 @@ pub mod kernel_tests { fn conv1d_causal_prefill_setup(s: usize, conv_dim: usize, kc: usize) -> TestSetup { let dt = DType::F32; let xbc: Vec = (0..s * conv_dim).map(|i| ((i as f32) * 0.011).sin() * 0.5).collect(); - let w: Vec = (0..kc * conv_dim).map(|i| 0.1 + ((i as f32) * 0.019).cos() * 0.2).collect(); + let w: Vec = + (0..kc * conv_dim).map(|i| 0.1 + ((i as f32) * 0.019).cos() * 0.2).collect(); let bias: Vec = (0..conv_dim).map(|i| (i as f32) * 0.001 - 0.05).collect(); let y_exp = conv1d_causal_prefill_oracle(&xbc, &w, &bias, s, conv_dim, kc); use super::conv1d_causal_prefill; @@ -756,7 +798,13 @@ pub mod kernel_tests { // ── strided_col_copy ────────────────────────────────────────────────── - fn strided_col_copy_oracle(src: &[f32], s: usize, stride: usize, col_off: usize, width: usize) -> Vec { + fn strided_col_copy_oracle( + src: &[f32], + s: usize, + stride: usize, + col_off: usize, + width: usize, + ) -> Vec { (0..s).flat_map(|ti| (0..width).map(move |ci| src[ti * stride + col_off + ci])).collect() } @@ -783,10 +831,14 @@ pub mod kernel_tests { fn softplus_add_rows_oracle(src: &[f32], bias: &[f32], n: usize) -> Vec { let s = src.len() / n; - (0..s).flat_map(|ti| (0..n).map(move |hi| { - let raw = src[ti * n + hi] + bias[hi]; - if raw > 20.0 { raw } else { (1.0f32 + raw.exp()).ln() } - })).collect() + (0..s) + .flat_map(|ti| { + (0..n).map(move |hi| { + let raw = src[ti * n + hi] + bias[hi]; + if raw > 20.0 { raw } else { (1.0f32 + raw.exp()).ln() } + }) + }) + .collect() } fn softplus_add_rows_setup(s: usize, n: usize) -> TestSetup { @@ -918,10 +970,10 @@ pub mod kernel_benches { // Each thread identifies which output slice it belongs to and writes there. #[kernel] pub fn mamba_split_proj( - proj: Tensor, // [s * in_proj_out] flat row-major - mut z_out: Tensor, // [s * di] + proj: Tensor, // [s * in_proj_out] flat row-major + mut z_out: Tensor, // [s * di] mut xbc_out: Tensor, // [s * conv_dim] - mut dt_out: Tensor, // [s * m_nh] + mut dt_out: Tensor, // [s * m_nh] #[constexpr] in_proj_out: u32, #[constexpr] di: u32, #[constexpr] conv_dim: u32, @@ -953,13 +1005,13 @@ pub fn mamba_split_proj( // Grid: [s * conv_dim, 1, 1]; one thread per source element. #[kernel] pub fn mamba_split_conv( - yc: Tensor, // [s * conv_dim] flat row-major - mut x_out: Tensor, // [s * di] - mut b_out: Tensor, // [s * ng_ds] - mut c_out: Tensor, // [s * ng_ds] + yc: Tensor, // [s * conv_dim] flat row-major + mut x_out: Tensor, // [s * di] + mut b_out: Tensor, // [s * ng_ds] + mut c_out: Tensor, // [s * ng_ds] #[constexpr] conv_dim: u32, #[constexpr] di: u32, - #[constexpr] ng_ds: u32, // ng * ds + #[constexpr] ng_ds: u32, // ng * ds ) { let idx = program_id::<0>(); let ti = idx / conv_dim; @@ -990,19 +1042,19 @@ pub fn mamba_split_conv( // Each thread in the block handles 4 consecutive elements. #[kernel] pub fn gated_group_rmsnorm_batched( - y: Tensor, // [s * di] flat - z: Tensor, // [s * di] flat - w: Tensor, // [di] norm weights (shared across tokens) + y: Tensor, // [s * di] flat + z: Tensor, // [s * di] flat + w: Tensor, // [di] norm weights (shared across tokens) mut out: Tensor, // [s * di] eps_buf: Tensor, // [1] - #[constexpr] gs: u32, // group size (512 for Nemotron) - #[constexpr] ng: u32, // number of groups per token (8 for Nemotron) + #[constexpr] gs: u32, // group size (512 for Nemotron) + #[constexpr] ng: u32, // number of groups per token (8 for Nemotron) ) { // program_id::<0>() = token * ng + group let tg = program_id::<0>(); let grp = tg - (tg / ng) * ng; - let ti = tg / ng; - let rs = ti * ng * gs + grp * gs; // start offset in [s * di] + let ti = tg / ng; + let rs = ti * ng * gs + grp * gs; // start offset in [s * di] let col = tid * 4u32; let in_bounds = col + 3u32 < gs; let safe_col = select(in_bounds, col, 0u32); @@ -1026,10 +1078,10 @@ pub fn gated_group_rmsnorm_batched( let rms = rsqrt(ssq / (gs.cast::()) + eps); if in_bounds { let base = rs + col; - store(out[base], (g0 * rms * load(w[grp * gs + col]))); - store(out[base + 1u32], (g1 * rms * load(w[grp * gs + col + 1u32]))); - store(out[base + 2u32], (g2 * rms * load(w[grp * gs + col + 2u32]))); - store(out[base + 3u32], (g3 * rms * load(w[grp * gs + col + 3u32]))); + store(out[base], (g0 * rms * load(w[grp * gs + col]))); + store(out[base + 1u32], (g1 * rms * load(w[grp * gs + col + 1u32]))); + store(out[base + 2u32], (g2 * rms * load(w[grp * gs + col + 2u32]))); + store(out[base + 3u32], (g3 * rms * load(w[grp * gs + col + 3u32]))); } } @@ -1051,15 +1103,15 @@ pub fn gated_group_rmsnorm_batched( // dt layout [T, H]; lcs layout [nc*H, L]. Grid: [nc*H, 1, 1]. #[kernel] pub fn ssd_lcs( - dt: Tensor, // [T, H] - a_log: Tensor, // [H] - mut lcs: Tensor, // [nc*H, L] + dt: Tensor, // [T, H] + a_log: Tensor, // [H] + mut lcs: Tensor, // [nc*H, L] #[constexpr] t_total: u32, #[constexpr] n_heads: u32, #[constexpr] l: u32, #[constexpr] nc: u32, ) { - let idx = program_id::<0>(); // bh = c*H + h + let idx = program_id::<0>(); // bh = c*H + h let total = nc * n_heads; if idx < total { let c = idx / n_heads; @@ -1079,8 +1131,8 @@ pub fn ssd_lcs( // h/hpg). One thread per output element. Grid: [nc*H*L*ds, 1, 1]. #[kernel] pub fn ssd_gather_bc( - b_mat: Tensor, // [T, G, ds] - c_mat: Tensor, // [T, G, ds] + b_mat: Tensor, // [T, G, ds] + c_mat: Tensor, // [T, G, ds] mut b_out: Tensor, // [nc*H, L, ds] mut c_out: Tensor, // [nc*H, L, ds] #[constexpr] t_total: u32, @@ -1115,8 +1167,8 @@ pub fn ssd_gather_bc( // Grid: [nc*H*dh*L, 1, 1]. #[kernel] pub fn ssd_xt( - x: Tensor, // [T, H, dh] - mut xt: Tensor, // [nc*H, dh, L] + x: Tensor, // [T, H, dh] + mut xt: Tensor, // [nc*H, dh, L] #[constexpr] t_total: u32, #[constexpr] n_heads: u32, #[constexpr] dh: u32, @@ -1126,7 +1178,7 @@ pub fn ssd_xt( let e = program_id::<0>(); let n = nc * n_heads * dh * l; if e < n { - let i = e - (e / l) * l; // position within chunk + let i = e - (e / l) * l; // position within chunk let p = (e / l) - (e / (l * dh)) * dh; // head_dim index let bh = e / (l * dh); let c = bh / n_heads; @@ -1144,9 +1196,9 @@ pub fn ssd_xt( // Grid: [nc*H*L*L, 1, 1]. #[kernel] pub fn ssd_mmask( - cb: Tensor, // [nc*H, L, L] = C·Bᵀ - lcs: Tensor, // [nc*H, L] - dt: Tensor, // [T, H] + cb: Tensor, // [nc*H, L, L] = C·Bᵀ + lcs: Tensor, // [nc*H, L] + dt: Tensor, // [T, H] mut m_out: Tensor, // [nc*H, L, L] #[constexpr] t_total: u32, #[constexpr] n_heads: u32, @@ -1178,10 +1230,10 @@ pub fn ssd_mmask( // Grid: [nc*H*ds*L, 1, 1]. #[kernel] pub fn ssd_bdt( - b_mat: Tensor, // [T, G, ds] - lcs: Tensor, // [nc*H, L] - dt: Tensor, // [T, H] - mut bdt: Tensor, // [nc*H, ds, L] + b_mat: Tensor, // [T, G, ds] + lcs: Tensor, // [nc*H, L] + dt: Tensor, // [T, H] + mut bdt: Tensor, // [nc*H, ds, L] #[constexpr] t_total: u32, #[constexpr] n_heads: u32, #[constexpr] n_groups: u32, @@ -1193,7 +1245,7 @@ pub fn ssd_bdt( let e = program_id::<0>(); let n = nc * n_heads * ds * l; if e < n { - let j = e - (e / l) * l; // position + let j = e - (e / l) * l; // position let s = (e / l) - (e / (l * ds)) * ds; // state index let bh = e / (l * ds); let c = bh / n_heads; @@ -1217,10 +1269,10 @@ pub fn ssd_bdt( // Grid: [H*ds*dh, 1, 1]; one thread per (head, s, p), loops nc serially. #[kernel] pub fn ssd_recur( - s_chunk: Tensor, // [nc*H, ds, dh] - lcs: Tensor, // [nc*H, L] - state_in: Tensor, // [H, dh, ds] - mut sin_t: Tensor, // [nc*H, dh, ds] (S_inᵀ per chunk) + s_chunk: Tensor, // [nc*H, ds, dh] + lcs: Tensor, // [nc*H, L] + state_in: Tensor, // [H, dh, ds] + mut sin_t: Tensor, // [nc*H, dh, ds] (S_inᵀ per chunk) mut state_out: Tensor, // [H, dh, ds] #[constexpr] n_heads: u32, #[constexpr] dh: u32, @@ -1231,7 +1283,7 @@ pub fn ssd_recur( let idx = program_id::<0>(); let tot = n_heads * ds * dh; if idx < tot { - let p = idx - (idx / dh) * dh; // head_dim + let p = idx - (idx / dh) * dh; // head_dim let s = (idx / dh) - (idx / (dh * ds)) * ds; // state let h = idx / (dh * ds); let mut st = load(state_in[(h * dh + p) * ds + s]); @@ -1253,12 +1305,12 @@ pub fn ssd_recur( // Grid: [nc*H*L*dh, 1, 1]. #[kernel] pub fn ssd_combine( - y_intra: Tensor, // [nc*H, L, dh] - cs: Tensor, // [nc*H, L, dh] - lcs: Tensor, // [nc*H, L] - x: Tensor, // [T, H, dh] - d_skip: Tensor, // [H] - mut y: Tensor, // [T, H, dh] + y_intra: Tensor, // [nc*H, L, dh] + cs: Tensor, // [nc*H, L, dh] + lcs: Tensor, // [nc*H, L] + x: Tensor, // [T, H, dh] + d_skip: Tensor, // [H] + mut y: Tensor, // [T, H, dh] #[constexpr] t_total: u32, #[constexpr] n_heads: u32, #[constexpr] dh: u32, @@ -1311,11 +1363,11 @@ pub fn ssd_combine( // (b_g, c_g) but without ever materializing them or CB. #[kernel] pub fn ssd_g1_cb( - b_mat: Tensor, // [T, G, ds] (weight role) - c_mat: Tensor, // [T, G, ds] (input role) - lcs: Tensor, // [nc*H, L] - dt: Tensor, // [T, H] - mut out: Tensor, // [nc*H, L, L] (M = CB ⊙ decay-mask) + b_mat: Tensor, // [T, G, ds] (weight role) + c_mat: Tensor, // [T, G, ds] (input role) + lcs: Tensor, // [nc*H, L] + dt: Tensor, // [T, H] + mut out: Tensor, // [nc*H, L, L] (M = CB ⊙ decay-mask) #[constexpr] t_total: u32, #[constexpr] n_heads: u32, #[constexpr] n_groups: u32, @@ -1323,7 +1375,7 @@ pub fn ssd_g1_cb( #[constexpr] l: u32, #[constexpr] ds: u32, ) { - let bz = program_id::<2>(); // batch = c*H + h + let bz = program_id::<2>(); // batch = c*H + h let c = bz / n_heads; let h = bz - c * n_heads; let g = h / hpg; @@ -1384,9 +1436,9 @@ pub fn ssd_g1_cb( // k = ds. sin_t is already per-batch [nc*H, dh, ds]; only C is broadcast. #[kernel] pub fn ssd_g4_cs( - sin_t: Tensor, // [nc*H, dh, ds] (weight role, per-batch) - c_mat: Tensor, // [T, G, ds] (input role, broadcast) - mut out: Tensor, // [nc*H, L, dh] + sin_t: Tensor, // [nc*H, dh, ds] (weight role, per-batch) + c_mat: Tensor, // [T, G, ds] (input role, broadcast) + mut out: Tensor, // [nc*H, L, dh] #[constexpr] t_total: u32, #[constexpr] n_heads: u32, #[constexpr] n_groups: u32, @@ -1395,11 +1447,11 @@ pub fn ssd_g4_cs( #[constexpr] ds: u32, #[constexpr] dh: u32, ) { - let bz = program_id::<2>(); // batch = c*H + h + let bz = program_id::<2>(); // batch = c*H + h let c = bz / n_heads; let h = bz - c * n_heads; let g = h / hpg; - let w_base = bz * dh * ds; // sin_t[bz] base + let w_base = bz * dh * ds; // sin_t[bz] base let tid = simd_id * 32u32 + simd_lane; let lr = tid / 32u32; // output row within tile (i, 0..31) let lo = tid % 32u32; // output col within tile (p, 0..31) diff --git a/crates/metaltile-std/tests/cuda_bench_corpus.rs b/crates/metaltile-std/tests/cuda_bench_corpus.rs index 9efd1b40..a628376b 100644 --- a/crates/metaltile-std/tests/cuda_bench_corpus.rs +++ b/crates/metaltile-std/tests/cuda_bench_corpus.rs @@ -33,10 +33,11 @@ //! kernels bench honestly. #![cfg(feature = "cuda")] +use std::collections::BTreeMap; + use metaltile::runner::BenchStats; use metaltile_core::dtype::DType; use metaltile_runtime::CudaDevice; -use std::collections::BTreeMap; /// Warmup launches (JIT/cache/clock ramp) before timing. const WARMUP: u32 = 5; @@ -114,7 +115,7 @@ fn bench_corpus_on_cuda() { // Effective bandwidth: bytes moved ÷ steady-state (min) latency. let gbps = moved as f64 / (stats.min_us * 1_000.0); rows.push(Row { name: t.name().to_string(), dt: dt_label(dt), stats, gbps }); - } + }, // Unsupported / codegen-gap kernels just don't get a number. Err(_) => skipped += 1, } @@ -123,7 +124,9 @@ fn bench_corpus_on_cuda() { // Sort heaviest-first by steady-state latency — the meaningful signal at // corpus sizes (GB/s is launch-bound noise; see the module docs). - rows.sort_by(|a, b| b.stats.min_us.partial_cmp(&a.stats.min_us).unwrap_or(std::cmp::Ordering::Equal)); + rows.sort_by(|a, b| { + b.stats.min_us.partial_cmp(&a.stats.min_us).unwrap_or(std::cmp::Ordering::Equal) + }); eprintln!( "{:<34} {:>5} {:>10} {:>10} {:>10} {:>8} {:>9}", diff --git a/crates/metaltile-std/tests/cuda_kernel_corpus.rs b/crates/metaltile-std/tests/cuda_kernel_corpus.rs index 47ad0dc0..3d8e27e0 100644 --- a/crates/metaltile-std/tests/cuda_kernel_corpus.rs +++ b/crates/metaltile-std/tests/cuda_kernel_corpus.rs @@ -14,26 +14,44 @@ //! Runs only with `--features cuda` on a CUDA host (the GX10 / sm_121). #![cfg(feature = "cuda")] +use std::collections::BTreeMap; + use metaltile_core::dtype::DType; use metaltile_runtime::CudaDevice; -use std::collections::BTreeMap; fn read_raw_f32(bytes: &[u8], dt: DType, n: usize) -> Vec { match dt { - DType::F32 => bytes.chunks_exact(4).take(n) - .map(|b| f32::from_le_bytes(b.try_into().unwrap())).collect(), - DType::F16 => bytes.chunks_exact(2).take(n).map(|b| { - let bits = u16::from_le_bytes(b.try_into().unwrap()); - half::f16::from_bits(bits).to_f32() - }).collect(), - DType::BF16 => bytes.chunks_exact(2).take(n).map(|b| { - let bits = u16::from_le_bytes(b.try_into().unwrap()); - half::bf16::from_bits(bits).to_f32() - }).collect(), - DType::I32 => bytes.chunks_exact(4).take(n) - .map(|b| i32::from_le_bytes(b.try_into().unwrap()) as f32).collect(), - DType::U32 => bytes.chunks_exact(4).take(n) - .map(|b| u32::from_le_bytes(b.try_into().unwrap()) as f32).collect(), + DType::F32 => bytes + .chunks_exact(4) + .take(n) + .map(|b| f32::from_le_bytes(b.try_into().unwrap())) + .collect(), + DType::F16 => bytes + .chunks_exact(2) + .take(n) + .map(|b| { + let bits = u16::from_le_bytes(b.try_into().unwrap()); + half::f16::from_bits(bits).to_f32() + }) + .collect(), + DType::BF16 => bytes + .chunks_exact(2) + .take(n) + .map(|b| { + let bits = u16::from_le_bytes(b.try_into().unwrap()); + half::bf16::from_bits(bits).to_f32() + }) + .collect(), + DType::I32 => bytes + .chunks_exact(4) + .take(n) + .map(|b| i32::from_le_bytes(b.try_into().unwrap()) as f32) + .collect(), + DType::U32 => bytes + .chunks_exact(4) + .take(n) + .map(|b| u32::from_le_bytes(b.try_into().unwrap()) as f32) + .collect(), DType::I8 => bytes.iter().take(n).map(|&b| b as i8 as f32).collect(), DType::U8 => bytes.iter().take(n).map(|&b| b as f32).collect(), _ => vec![0.0; n], @@ -53,31 +71,40 @@ const KNOWN_HARD: &[(&str, &str)] = &[ // (flash-attention accumulation over K=head_dim). d64 + all f16/bf16 pass. ]; -fn known_hard(name: &str) -> bool { - KNOWN_HARD.iter().any(|(k, _)| name.contains(k)) -} +fn known_hard(name: &str) -> bool { KNOWN_HARD.iter().any(|(k, _)| name.contains(k)) } fn is_unsupported(msg: &str) -> bool { let m = msg.to_lowercase(); [ // Codegen coverage gaps (kernel not wired yet on CUDA). - "phase 1", "phase 2", "not supported", "not yet implemented", "strided", - "kernelmode", "multi-dimensional", "transform", "secondary", + "phase 1", + "phase 2", + "not supported", + "not yet implemented", + "strided", + "kernelmode", + "multi-dimensional", + "transform", + "secondary", // Device-capability limits (mirrors the Vulkan harness's device-cap // bucket): a kernel the codegen *does* cover but the target arch // physically cannot run. These reflect bit-accuracy on what the arch // CAN run, so classify as UNSUPPORTED rather than a hard ERROR. // - >48KB dynamic shared memory on pre-Volta (sm_5x/6x, e.g. Pascal): // typed MetalTileError::DeviceCapability surfaced before launch. - "device capability", "unsupported on this device", ">48kb", + "device capability", + "unsupported on this device", + ">48kb", "dynamic shared memory", // - raw driver rejection if it ever escapes pre-launch validation. "culaunchkernel: invalid argument", // - tensor-core MMA / WMMA paths needing sm_70+ on older arches. - "requires sm_70", "tensor core", "mma is not supported", + "requires sm_70", + "tensor core", + "mma is not supported", ] - .iter() - .any(|p| m.contains(p)) + .iter() + .any(|p| m.contains(p)) } #[test] @@ -149,9 +176,10 @@ fn run_corpus_on_cuda() { known += 1; } else { mismatch += 1; - hard_failures.push(format!("MISMATCH {label}: max|Δ|={worst:.3e} > {tol:.3e}")); + hard_failures + .push(format!("MISMATCH {label}: max|Δ|={worst:.3e} > {tol:.3e}")); } - } + }, Err(e) => { let msg = e.to_string(); if known_hard(t.name()) { @@ -173,13 +201,15 @@ fn run_corpus_on_cuda() { error += 1; hard_failures.push(format!("ERROR {label}: {msg}")); } - } + }, } } } eprintln!("\n=== CUDA corpus result ==="); - eprintln!("PASS={pass} KNOWN_HARD={known} MISMATCH={mismatch} UNSUPPORTED={unsupported} ERROR={error}"); + eprintln!( + "PASS={pass} KNOWN_HARD={known} MISMATCH={mismatch} UNSUPPORTED={unsupported} ERROR={error}" + ); eprintln!("--- unsupported reasons (top buckets) ---"); let mut reasons: Vec<_> = unsup_reasons.iter().collect(); reasons.sort_by(|a, b| b.1.cmp(a.1)); diff --git a/crates/metaltile-std/tests/hip_kernel_corpus.rs b/crates/metaltile-std/tests/hip_kernel_corpus.rs index bd4ced22..b11de1de 100644 --- a/crates/metaltile-std/tests/hip_kernel_corpus.rs +++ b/crates/metaltile-std/tests/hip_kernel_corpus.rs @@ -17,26 +17,44 @@ //! Runs only with `--features hip`. #![cfg(feature = "hip")] +use std::collections::BTreeMap; + use metaltile_core::dtype::DType; use metaltile_runtime::HipDevice; -use std::collections::BTreeMap; fn read_raw_f32(bytes: &[u8], dt: DType, n: usize) -> Vec { match dt { - DType::F32 => bytes.chunks_exact(4).take(n) - .map(|b| f32::from_le_bytes(b.try_into().unwrap())).collect(), - DType::F16 => bytes.chunks_exact(2).take(n).map(|b| { - let bits = u16::from_le_bytes(b.try_into().unwrap()); - half::f16::from_bits(bits).to_f32() - }).collect(), - DType::BF16 => bytes.chunks_exact(2).take(n).map(|b| { - let bits = u16::from_le_bytes(b.try_into().unwrap()); - half::bf16::from_bits(bits).to_f32() - }).collect(), - DType::I32 => bytes.chunks_exact(4).take(n) - .map(|b| i32::from_le_bytes(b.try_into().unwrap()) as f32).collect(), - DType::U32 => bytes.chunks_exact(4).take(n) - .map(|b| u32::from_le_bytes(b.try_into().unwrap()) as f32).collect(), + DType::F32 => bytes + .chunks_exact(4) + .take(n) + .map(|b| f32::from_le_bytes(b.try_into().unwrap())) + .collect(), + DType::F16 => bytes + .chunks_exact(2) + .take(n) + .map(|b| { + let bits = u16::from_le_bytes(b.try_into().unwrap()); + half::f16::from_bits(bits).to_f32() + }) + .collect(), + DType::BF16 => bytes + .chunks_exact(2) + .take(n) + .map(|b| { + let bits = u16::from_le_bytes(b.try_into().unwrap()); + half::bf16::from_bits(bits).to_f32() + }) + .collect(), + DType::I32 => bytes + .chunks_exact(4) + .take(n) + .map(|b| i32::from_le_bytes(b.try_into().unwrap()) as f32) + .collect(), + DType::U32 => bytes + .chunks_exact(4) + .take(n) + .map(|b| u32::from_le_bytes(b.try_into().unwrap()) as f32) + .collect(), DType::I8 => bytes.iter().take(n).map(|&b| b as i8 as f32).collect(), DType::U8 => bytes.iter().take(n).map(|&b| b as f32).collect(), _ => vec![0.0; n], @@ -60,23 +78,28 @@ fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 { // bit-accurate to the per-kernel band. const KNOWN_HARD: &[(&str, &str)] = &[]; -fn known_hard(name: &str) -> bool { - KNOWN_HARD.iter().any(|(k, _)| name.contains(k)) -} +fn known_hard(name: &str) -> bool { KNOWN_HARD.iter().any(|(k, _)| name.contains(k)) } fn is_unsupported(msg: &str) -> bool { let m = msg.to_lowercase(); [ - "phase 1", "phase 2", "not supported", "not yet implemented", "strided", - "kernelmode", "multi-dimensional", "transform", "secondary", + "phase 1", + "phase 2", + "not supported", + "not yet implemented", + "strided", + "kernelmode", + "multi-dimensional", + "transform", + "secondary", // HIP-specific compile failures we treat as "kernel uses a CUDA // construct HIP doesn't accept" — counted as UNSUPPORTED so the // corpus result tracks what's bit-accurate, not what'd hit a // future textual-transform extension. "hiprtc", ] - .iter() - .any(|p| m.contains(p)) + .iter() + .any(|p| m.contains(p)) } #[test] @@ -155,7 +178,7 @@ fn run_corpus_on_hip() { hard_failures .push(format!("MISMATCH {label}: max|Δ|={worst:.3e} > {tol:.3e}")); } - } + }, Err(e) => { let msg = e.to_string(); if known_hard(t.name()) { @@ -176,13 +199,15 @@ fn run_corpus_on_hip() { error += 1; hard_failures.push(format!("ERROR {label}: {msg}")); } - } + }, } } } eprintln!("\n=== HIP corpus result ==="); - eprintln!("PASS={pass} KNOWN_HARD={known} MISMATCH={mismatch} UNSUPPORTED={unsupported} ERROR={error}"); + eprintln!( + "PASS={pass} KNOWN_HARD={known} MISMATCH={mismatch} UNSUPPORTED={unsupported} ERROR={error}" + ); eprintln!("--- unsupported reasons (top buckets) ---"); let mut reasons: Vec<_> = unsup_reasons.iter().collect(); reasons.sort_by(|a, b| b.1.cmp(a.1)); diff --git a/crates/metaltile-std/tests/vulkan_add_bias_rows.rs b/crates/metaltile-std/tests/vulkan_add_bias_rows.rs index 90f8478d..286d46b6 100644 --- a/crates/metaltile-std/tests/vulkan_add_bias_rows.rs +++ b/crates/metaltile-std/tests/vulkan_add_bias_rows.rs @@ -18,9 +18,11 @@ use std::collections::BTreeMap; -use metaltile_core::dtype::DType; -use metaltile_core::ir::{BinOpKind, IndexExpr, Kernel, Op, Param, ParamKind, ValueId}; -use metaltile_core::shape::Shape; +use metaltile_core::{ + dtype::DType, + ir::{BinOpKind, IndexExpr, Kernel, Op, Param, ParamKind, ValueId}, + shape::Shape, +}; use metaltile_runtime::VulkanDevice; fn p(name: &str, is_output: bool) -> Param { @@ -59,14 +61,14 @@ fn bias_rows_kernel(n: u32, index_mode: &str) -> Kernel { match index_mode { "mod" => { k.body.push_op(Op::BinOp { op: BinOpKind::Mod, lhs: i, rhs: n_const }, col); - } + }, "divmul" => { let q = nid(); k.body.push_op(Op::BinOp { op: BinOpKind::Div, lhs: i, rhs: n_const }, q); let qn = nid(); k.body.push_op(Op::BinOp { op: BinOpKind::Mul, lhs: q, rhs: n_const }, qn); k.body.push_op(Op::BinOp { op: BinOpKind::Sub, lhs: i, rhs: qn }, col); - } + }, _ => unreachable!(), } @@ -105,11 +107,7 @@ fn f32_bytes(v: &[f32]) -> Vec { } fn read_f32(bytes: &[u8], n: usize) -> Vec { - bytes - .chunks_exact(4) - .take(n) - .map(|b| f32::from_le_bytes(b.try_into().unwrap())) - .collect() + bytes.chunks_exact(4).take(n).map(|b| f32::from_le_bytes(b.try_into().unwrap())).collect() } fn run_case(dev: &VulkanDevice, n_rows: usize, n: usize, index_mode: &str) { diff --git a/crates/metaltile-std/tests/vulkan_coopmat_gemm.rs b/crates/metaltile-std/tests/vulkan_coopmat_gemm.rs index 8af55636..2d167bb7 100644 --- a/crates/metaltile-std/tests/vulkan_coopmat_gemm.rs +++ b/crates/metaltile-std/tests/vulkan_coopmat_gemm.rs @@ -15,8 +15,7 @@ //! # coopmat: MT_VK_COOPMAT=1 (same line) #![cfg(feature = "vulkan")] -use std::collections::BTreeMap; -use std::time::Instant; +use std::{collections::BTreeMap, time::Instant}; use metaltile_core::dtype::DType; use metaltile_runtime::VulkanDevice; @@ -31,9 +30,7 @@ fn xorshift(s: &mut u32) -> u32 { x } -fn f16_round(v: f32) -> f32 { - half::f16::from_f32(v).to_f32() -} +fn f16_round(v: f32) -> f32 { half::f16::from_f32(v).to_f32() } #[test] fn coopmat_gemm_q8_mpp_correct_and_fast() { @@ -90,12 +87,8 @@ fn coopmat_gemm_q8_mpp_correct_and_fast() { let pack_f16 = |v: &[f32]| -> Vec { v.iter().flat_map(|&f| half::f16::from_f32(f).to_bits().to_le_bytes()).collect() }; - let pack_u32 = |v: &[u32]| -> Vec { - v.iter().flat_map(|&u| u.to_le_bytes()).collect() - }; - let pack_f32 = |v: &[f32]| -> Vec { - v.iter().flat_map(|&f| f.to_le_bytes()).collect() - }; + let pack_u32 = |v: &[u32]| -> Vec { v.iter().flat_map(|&u| u.to_le_bytes()).collect() }; + let pack_f32 = |v: &[f32]| -> Vec { v.iter().flat_map(|&f| f.to_le_bytes()).collect() }; let mut buffers: BTreeMap> = BTreeMap::new(); buffers.insert("x".into(), pack_f16(&x)); @@ -110,9 +103,7 @@ fn coopmat_gemm_q8_mpp_correct_and_fast() { let tpg = [128u32, 1, 1]; // Correctness. - let outputs = dev - .run_kernel(&kernel, &buffers, grid, tpg) - .expect("run_kernel"); + let outputs = dev.run_kernel(&kernel, &buffers, grid, tpg).expect("run_kernel"); let got_bytes = outputs.get("out").expect("out buffer"); let got: Vec = got_bytes .chunks_exact(2) @@ -161,7 +152,9 @@ fn coopmat_gemm_q8_mpp_correct_and_fast() { maxd_ab = maxd_ab.max(d); max_rel_ab = max_rel_ab.max(d / refv[i].abs().max(8.0)); } - eprintln!(" coopmat vs scalar (informational): maxAbsDiff={maxd_ab:.4} maxRelDiff={max_rel_ab:.4}"); + eprintln!( + " coopmat vs scalar (informational): maxAbsDiff={maxd_ab:.4} maxRelDiff={max_rel_ab:.4}" + ); // The two GPU paths differ in staging precision (coopmat: f16 // shared tiles + f32 fragment accumulate; scalar: f32 shared // tiles). Both round to f16 output. The gap is bounded by the @@ -185,6 +178,10 @@ fn coopmat_gemm_q8_mpp_correct_and_fast() { let tflops = flops / per / 1e12; eprintln!( " shape {}x{}x{} {:.3} ms/dispatch {:.2} TFLOP/s (incl. run_kernel overhead)", - n_rows, out_dim, k_in, per * 1e3, tflops + n_rows, + out_dim, + k_in, + per * 1e3, + tflops ); } diff --git a/crates/metaltile-std/tests/vulkan_kernel_corpus.rs b/crates/metaltile-std/tests/vulkan_kernel_corpus.rs index 06d25564..10a7eec4 100644 --- a/crates/metaltile-std/tests/vulkan_kernel_corpus.rs +++ b/crates/metaltile-std/tests/vulkan_kernel_corpus.rs @@ -15,26 +15,44 @@ //! Runs only with `--features vulkan`. #![cfg(feature = "vulkan")] +use std::collections::BTreeMap; + use metaltile_core::dtype::DType; use metaltile_runtime::VulkanDevice; -use std::collections::BTreeMap; fn read_raw_f32(bytes: &[u8], dt: DType, n: usize) -> Vec { match dt { - DType::F32 => bytes.chunks_exact(4).take(n) - .map(|b| f32::from_le_bytes(b.try_into().unwrap())).collect(), - DType::F16 => bytes.chunks_exact(2).take(n).map(|b| { - let bits = u16::from_le_bytes(b.try_into().unwrap()); - half::f16::from_bits(bits).to_f32() - }).collect(), - DType::BF16 => bytes.chunks_exact(2).take(n).map(|b| { - let bits = u16::from_le_bytes(b.try_into().unwrap()); - half::bf16::from_bits(bits).to_f32() - }).collect(), - DType::I32 => bytes.chunks_exact(4).take(n) - .map(|b| i32::from_le_bytes(b.try_into().unwrap()) as f32).collect(), - DType::U32 => bytes.chunks_exact(4).take(n) - .map(|b| u32::from_le_bytes(b.try_into().unwrap()) as f32).collect(), + DType::F32 => bytes + .chunks_exact(4) + .take(n) + .map(|b| f32::from_le_bytes(b.try_into().unwrap())) + .collect(), + DType::F16 => bytes + .chunks_exact(2) + .take(n) + .map(|b| { + let bits = u16::from_le_bytes(b.try_into().unwrap()); + half::f16::from_bits(bits).to_f32() + }) + .collect(), + DType::BF16 => bytes + .chunks_exact(2) + .take(n) + .map(|b| { + let bits = u16::from_le_bytes(b.try_into().unwrap()); + half::bf16::from_bits(bits).to_f32() + }) + .collect(), + DType::I32 => bytes + .chunks_exact(4) + .take(n) + .map(|b| i32::from_le_bytes(b.try_into().unwrap()) as f32) + .collect(), + DType::U32 => bytes + .chunks_exact(4) + .take(n) + .map(|b| u32::from_le_bytes(b.try_into().unwrap()) as f32) + .collect(), DType::I8 => bytes.iter().take(n).map(|&b| b as i8 as f32).collect(), DType::U8 => bytes.iter().take(n).map(|&b| b as f32).collect(), _ => vec![0.0; n], @@ -51,28 +69,42 @@ fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 { // 4164/4164 = 100% bit-accurate. const KNOWN_HARD: &[(&str, &str)] = &[]; -fn known_hard(name: &str) -> bool { - KNOWN_HARD.iter().any(|(k, _)| name.contains(k)) -} +fn known_hard(name: &str) -> bool { KNOWN_HARD.iter().any(|(k, _)| name.contains(k)) } fn is_unsupported(msg: &str) -> bool { let m = msg.to_lowercase(); [ - "phase 1", "phase 2", "phase 3", "phase 4", - "not supported", "not yet implemented", "not yet supported", - "strided", "kernelmode", "multi-dimensional", "transform", "secondary", - "dtype", "f16", "bf16", "i8", + "phase 1", + "phase 2", + "phase 3", + "phase 4", + "not supported", + "not yet implemented", + "not yet supported", + "strided", + "kernelmode", + "multi-dimensional", + "transform", + "secondary", + "dtype", + "f16", + "bf16", + "i8", // Shaderc compile failures we treat as UNSUPPORTED so the corpus // result reflects bit-accuracy on what's actually wired, not // shader-language gaps. - "shaderc_compile", "spirv:", "spirv ", "decode-", + "shaderc_compile", + "spirv:", + "spirv ", + "decode-", // Vulkan device limits — workgroup size cap, descriptor count, // push-constant size — these are dtype-orthogonal device caps // rather than codegen bugs. - "vkresult=", "no memory type", + "vkresult=", + "no memory type", ] - .iter() - .any(|p| m.contains(p)) + .iter() + .any(|p| m.contains(p)) } #[test] @@ -147,7 +179,7 @@ fn run_corpus_on_vulkan() { hard_failures .push(format!("MISMATCH {label}: max|Δ|={worst:.3e} > {tol:.3e}")); } - } + }, Err(e) => { let msg = e.to_string(); if known_hard(t.name()) { @@ -168,13 +200,15 @@ fn run_corpus_on_vulkan() { error += 1; hard_failures.push(format!("ERROR {label}: {msg}")); } - } + }, } } } eprintln!("\n=== Vulkan corpus result ==="); - eprintln!("PASS={pass} KNOWN_HARD={known} MISMATCH={mismatch} UNSUPPORTED={unsupported} ERROR={error}"); + eprintln!( + "PASS={pass} KNOWN_HARD={known} MISMATCH={mismatch} UNSUPPORTED={unsupported} ERROR={error}" + ); eprintln!("--- unsupported reasons (top buckets) ---"); let mut reasons: Vec<_> = unsup_reasons.iter().collect(); reasons.sort_by(|a, b| b.1.cmp(a.1)); @@ -205,10 +239,7 @@ fn run_corpus_on_vulkan() { s.split(" [").next().unwrap_or("").to_string() }) .collect(); - eprintln!( - "--- unique MISMATCH kernel-base count: {} ---", - unique_mm_kernels.len() - ); + eprintln!("--- unique MISMATCH kernel-base count: {} ---", unique_mm_kernels.len()); for k in &unique_mm_kernels { eprintln!(" · {k}"); } diff --git a/crates/metaltile-std/tests/vulkan_sdpa_multi.rs b/crates/metaltile-std/tests/vulkan_sdpa_multi.rs index a208a4a6..babf4541 100644 --- a/crates/metaltile-std/tests/vulkan_sdpa_multi.rs +++ b/crates/metaltile-std/tests/vulkan_sdpa_multi.rs @@ -35,9 +35,7 @@ fn xorshift(s: &mut u32) -> u32 { x } -fn rnd(s: &mut u32) -> f32 { - ((xorshift(s) % 2000) as f32 / 1000.0) - 1.0 -} +fn rnd(s: &mut u32) -> f32 { ((xorshift(s) % 2000) as f32 / 1000.0) - 1.0 } #[allow(clippy::too_many_arguments)] fn naive_sdpa_multi( @@ -98,11 +96,7 @@ fn f32_bytes(v: &[f32]) -> Vec { } fn read_f32(bytes: &[u8], n: usize) -> Vec { - bytes - .chunks_exact(4) - .take(n) - .map(|b| f32::from_le_bytes(b.try_into().unwrap())) - .collect() + bytes.chunks_exact(4).take(n).map(|b| f32::from_le_bytes(b.try_into().unwrap())).collect() } #[allow(clippy::too_many_arguments)] @@ -143,8 +137,7 @@ fn run_case( buffers.insert("base_kv".into(), (base_kv as u32).to_le_bytes().to_vec()); buffers.insert("n_query".into(), (n_query as u32).to_le_bytes().to_vec()); buffers.insert("kv_stride".into(), (kv_stride as u32).to_le_bytes().to_vec()); - buffers - .insert("heads_per_group".into(), (heads_per_group as u32).to_le_bytes().to_vec()); + buffers.insert("heads_per_group".into(), (heads_per_group as u32).to_le_bytes().to_vec()); buffers.insert("causal".into(), (causal as u32).to_le_bytes().to_vec()); buffers.insert("scale".into(), scale.to_le_bytes().to_vec()); From e79266b0f86bda4bf03c6df48894fa243951b3b6 Mon Sep 17 00:00:00 2001 From: TheTom Date: Thu, 23 Jul 2026 06:45:00 -0500 Subject: [PATCH 3/5] chore: workspace clippy + fmt clean, wire std smoke-test dev-deps Repo-wide cargo fmt, clippy fixes across the codegen crates (collapsed ifs, dropped useless format! calls, allow on the 9-arg spirv emit_op), typos config for the intentional optin identifiers, and the missing metaltile-core/metaltile-codegen dev-dependencies that the in-module smoke tests in metaltile-std reference (this also un-breaks the full test-target build, which previously failed with unresolved-crate errors before any test ran). No behavior change. --- _typos.toml | 2 + crates/metaltile-cli/src/cmd/update.rs | 3 +- crates/metaltile-codegen/src/backend.rs | 156 +++- crates/metaltile-codegen/src/cuda/mod.rs | 439 +++++++---- crates/metaltile-codegen/src/hip/mod.rs | 63 +- crates/metaltile-codegen/src/spirv/mod.rs | 719 ++++++++++-------- crates/metaltile-runtime/build.rs | 37 +- .../metaltile-runtime/src/device/cuda/ffi.rs | 74 +- .../metaltile-runtime/src/device/cuda/mod.rs | 622 ++++++++++----- .../metaltile-runtime/src/device/hip/ffi.rs | 24 +- .../metaltile-runtime/src/device/hip/mod.rs | 95 +-- .../src/device/vulkan/ffi.rs | 20 +- .../src/device/vulkan/mod.rs | 191 ++--- crates/metaltile-runtime/src/lib.rs | 14 +- crates/metaltile-runtime/tests/cuda_smoke.rs | 72 +- crates/metaltile-runtime/tests/hip_smoke.rs | 122 +-- .../metaltile-runtime/tests/vulkan_smoke.rs | 132 ++-- crates/metaltile-std/Cargo.toml | 3 + 18 files changed, 1718 insertions(+), 1070 deletions(-) diff --git a/_typos.toml b/_typos.toml index fe859beb..1589097e 100644 --- a/_typos.toml +++ b/_typos.toml @@ -6,3 +6,5 @@ ba = "ba" na = "na" bb = "bb" nb = "nb" +optin = "optin" +OPTIN = "OPTIN" diff --git a/crates/metaltile-cli/src/cmd/update.rs b/crates/metaltile-cli/src/cmd/update.rs index b005a06f..53060ddd 100644 --- a/crates/metaltile-cli/src/cmd/update.rs +++ b/crates/metaltile-cli/src/cmd/update.rs @@ -337,7 +337,8 @@ fn install_binary(src: &PathBuf, dest: &PathBuf) -> Result<(), crate::CliError> #[cfg(unix)] { use std::os::unix::fs::PermissionsExt as _; - fs::set_permissions(&tmp, fs::Permissions::from_mode(0o755)).map_err(crate::CliError::Io)?; + fs::set_permissions(&tmp, fs::Permissions::from_mode(0o755)) + .map_err(crate::CliError::Io)?; } fs::rename(&tmp, dest).map_err(|e| { diff --git a/crates/metaltile-codegen/src/backend.rs b/crates/metaltile-codegen/src/backend.rs index 59f1255e..01666e3b 100644 --- a/crates/metaltile-codegen/src/backend.rs +++ b/crates/metaltile-codegen/src/backend.rs @@ -301,16 +301,35 @@ impl TargetProfile { use UnaryOpKind::*; match self.target { Target::Metal => match op { - Exp => "exp", Exp2 => "exp2", Expm1 => "expm1", - Log => "log", Log2 => "log2", Log10 => "log10", - Sqrt => "sqrt", Rsqrt => "rsqrt", Recip => "1.0/", - Abs => "abs", Neg => "-", Ceil => "ceil", Floor => "floor", - Round => "round", Trunc => "trunc", Sign => "sign", - Sin => "sin", Cos => "cos", Tan => "tan", - Asin => "asin", Acos => "acos", Atan => "atan", - Sinh => "sinh", Cosh => "cosh", - Asinh => "asinh", Acosh => "acosh", Atanh => "atanh", - Erf => "erf", ErfInv => "erfinv", + Exp => "exp", + Exp2 => "exp2", + Expm1 => "expm1", + Log => "log", + Log2 => "log2", + Log10 => "log10", + Sqrt => "sqrt", + Rsqrt => "rsqrt", + Recip => "1.0/", + Abs => "abs", + Neg => "-", + Ceil => "ceil", + Floor => "floor", + Round => "round", + Trunc => "trunc", + Sign => "sign", + Sin => "sin", + Cos => "cos", + Tan => "tan", + Asin => "asin", + Acos => "acos", + Atan => "atan", + Sinh => "sinh", + Cosh => "cosh", + Asinh => "asinh", + Acosh => "acosh", + Atanh => "atanh", + Erf => "erf", + ErfInv => "erfinv", }, // CUDA *precise* device math (single-precision `f` suffix) — NOT // the `__`-prefixed fast-math intrinsics: the CPU oracle is IEEE, @@ -327,28 +346,66 @@ impl TargetProfile { // with ~2 ULP error vs ~1 ULP. Safe for attention softmax. // HIP: keep `expf` (hipRTC may not have `__expf`). Target::Cuda => match op { - Exp => "__expf", Exp2 => "exp2f", Expm1 => "expm1f", - Log => "logf", Log2 => "log2f", Log10 => "log10f", - Sqrt => "sqrtf", Rsqrt => "rsqrtf", Recip => "__frcp_rn", - Abs => "fabsf", Neg => "-", Ceil => "ceilf", Floor => "floorf", - Round => "roundf", Trunc => "truncf", Sign => "copysignf", - Sin => "sinf", Cos => "cosf", Tan => "tanf", - Asin => "asinf", Acos => "acosf", Atan => "atanf", - Sinh => "sinhf", Cosh => "coshf", - Asinh => "asinhf", Acosh => "acoshf", Atanh => "atanhf", - Erf => "erff", ErfInv => "erfinvf", + Exp => "__expf", + Exp2 => "exp2f", + Expm1 => "expm1f", + Log => "logf", + Log2 => "log2f", + Log10 => "log10f", + Sqrt => "sqrtf", + Rsqrt => "rsqrtf", + Recip => "__frcp_rn", + Abs => "fabsf", + Neg => "-", + Ceil => "ceilf", + Floor => "floorf", + Round => "roundf", + Trunc => "truncf", + Sign => "copysignf", + Sin => "sinf", + Cos => "cosf", + Tan => "tanf", + Asin => "asinf", + Acos => "acosf", + Atan => "atanf", + Sinh => "sinhf", + Cosh => "coshf", + Asinh => "asinhf", + Acosh => "acoshf", + Atanh => "atanhf", + Erf => "erff", + ErfInv => "erfinvf", }, Target::Hip => match op { - Exp => "expf", Exp2 => "exp2f", Expm1 => "expm1f", - Log => "logf", Log2 => "log2f", Log10 => "log10f", - Sqrt => "sqrtf", Rsqrt => "rsqrtf", Recip => "__frcp_rn", - Abs => "fabsf", Neg => "-", Ceil => "ceilf", Floor => "floorf", - Round => "roundf", Trunc => "truncf", Sign => "copysignf", - Sin => "sinf", Cos => "cosf", Tan => "tanf", - Asin => "asinf", Acos => "acosf", Atan => "atanf", - Sinh => "sinhf", Cosh => "coshf", - Asinh => "asinhf", Acosh => "acoshf", Atanh => "atanhf", - Erf => "erff", ErfInv => "erfinvf", + Exp => "expf", + Exp2 => "exp2f", + Expm1 => "expm1f", + Log => "logf", + Log2 => "log2f", + Log10 => "log10f", + Sqrt => "sqrtf", + Rsqrt => "rsqrtf", + Recip => "__frcp_rn", + Abs => "fabsf", + Neg => "-", + Ceil => "ceilf", + Floor => "floorf", + Round => "roundf", + Trunc => "truncf", + Sign => "copysignf", + Sin => "sinf", + Cos => "cosf", + Tan => "tanf", + Asin => "asinf", + Acos => "acosf", + Atan => "atanf", + Sinh => "sinhf", + Cosh => "coshf", + Asinh => "asinhf", + Acosh => "acoshf", + Atanh => "atanhf", + Erf => "erff", + ErfInv => "erfinvf", }, // GLSL.std.450 names (the SPIR-V extended-instruction set used // by Vulkan compute). `shaderc` lowers these via the standard @@ -356,17 +413,36 @@ impl TargetProfile { // rsqrt. `erf`/`erfinv` aren't in GLSL.std.450 — kernels using // them fall back to a software approximation in the preamble. Target::Spirv => match op { - Exp => "exp", Exp2 => "exp2", Expm1 => "mt_expm1", - Log => "log", Log2 => "log2", Log10 => "mt_log10", - Sqrt => "sqrt", Rsqrt => "inversesqrt", Recip => "1.0/", - Abs => "abs", Neg => "-", Ceil => "ceil", Floor => "floor", - Round => "roundEven", Trunc => "trunc", Sign => "sign", - Sin => "sin", Cos => "cos", Tan => "tan", - Asin => "asin", Acos => "acos", Atan => "atan", - Sinh => "sinh", Cosh => "cosh", - Asinh => "asinh", Acosh => "acosh", Atanh => "atanh", + Exp => "exp", + Exp2 => "exp2", + Expm1 => "mt_expm1", + Log => "log", + Log2 => "log2", + Log10 => "mt_log10", + Sqrt => "sqrt", + Rsqrt => "inversesqrt", + Recip => "1.0/", + Abs => "abs", + Neg => "-", + Ceil => "ceil", + Floor => "floor", + Round => "roundEven", + Trunc => "trunc", + Sign => "sign", + Sin => "sin", + Cos => "cos", + Tan => "tan", + Asin => "asin", + Acos => "acos", + Atan => "atan", + Sinh => "sinh", + Cosh => "cosh", + Asinh => "asinh", + Acosh => "acosh", + Atanh => "atanh", // Software approximations live in the GLSL preamble. - Erf => "mt_erf", ErfInv => "mt_erfinv", + Erf => "mt_erf", + ErfInv => "mt_erfinv", }, } } diff --git a/crates/metaltile-codegen/src/cuda/mod.rs b/crates/metaltile-codegen/src/cuda/mod.rs index 2c093cf7..1543ced1 100644 --- a/crates/metaltile-codegen/src/cuda/mod.rs +++ b/crates/metaltile-codegen/src/cuda/mod.rs @@ -30,14 +30,26 @@ //! - A trailing `unsigned int _n_elems` param + an `if (gtid >= _n_elems) //! return;` guard replace Metal's non-uniform-threadgroup bounds model. -use std::collections::BTreeMap; -use std::fmt::Write as _; +use std::{collections::BTreeMap, fmt::Write as _}; use metaltile_core::{ dtype::DType, ir::{ - ActKind, AtomicKind, BinOpKind, Block, CoopTileAccMode, CoopTileScope, IndexExpr, Kernel, - KernelMode, Op, Param, ParamKind, ReduceKind, UnaryOpKind, ValueId, + ActKind, + AtomicKind, + BinOpKind, + Block, + CoopTileAccMode, + CoopTileScope, + IndexExpr, + Kernel, + KernelMode, + Op, + Param, + ParamKind, + ReduceKind, + UnaryOpKind, + ValueId, }, }; @@ -160,9 +172,7 @@ impl Default for CudaGenerator { } impl CudaGenerator { - pub fn new() -> Self { - CudaGenerator { profile: TargetProfile::cuda() } - } + pub fn new() -> Self { CudaGenerator { profile: TargetProfile::cuda() } } /// Build a generator pinned to a specific profile (e.g. a Blackwell /// profile with `MmaStrategy::Tcgen05` for Phase 4). @@ -189,7 +199,7 @@ impl CudaGenerator { Some(hint) => format!("v_{hint}_{}", v.as_u32()), None => format!("v{}", v.as_u32()), } - } + }, None => "/**/".to_string(), } } @@ -216,7 +226,9 @@ impl CudaGenerator { let terms: Vec = many .iter() .enumerate() - .map(|(d, ix)| format!("({}) * {src}_strides[{d}]", self.idx_term(ix, block, ov))) + .map(|(d, ix)| { + format!("({}) * {src}_strides[{d}]", self.idx_term(ix, block, ov)) + }) .collect(); return Ok(terms.join(" + ")); } @@ -240,13 +252,13 @@ impl CudaGenerator { return Err(Error::UnsupportedOp(format!( "cuda: multi-dim index needs static dims on `{src}`" ))); - } + }, } } terms.push(format!("({}) * {stride}u", self.idx_term(ix, block, ov))); } Ok(terms.join(" + ")) - } + }, } } @@ -293,7 +305,7 @@ impl CudaGenerator { } else { format!(" const {ty}* __restrict__ {}", p.name) }) - } + }, ParamKind::Scalar => Ok(format!(" {} {}", cuda_type_name(p.dtype), p.name)), // Strided tensor: data ptr + companion shape/strides buffers // (referenced in the IR via Load{src:"{name}_shape"/"_strides"}). @@ -304,7 +316,7 @@ impl CudaGenerator { " {q}{ty}* __restrict__ {0},\n const unsigned int* __restrict__ {0}_shape,\n const unsigned int* __restrict__ {0}_strides", p.name )) - } + }, } } @@ -323,7 +335,7 @@ impl CudaGenerator { // Elementwise mode it is the global linear index. writeln!(out, " const unsigned int tid = _gtid;").ok(); self.emit_simd_aliases(out); - } + }, KernelMode::Reduction => { // The `n_simd==0 → return` freeze guard is only for the // threadgroup-reduce tree (needs ≥32 threads). Kernels that @@ -331,16 +343,28 @@ impl CudaGenerator { // return — gate the guard on Op::Reduce presence. let needs_guard = kernel_has_reduce(kernel); self.emit_reduction_preamble(out, needs_guard); - } + }, KernelMode::Grid3D => { // 3-axis global thread id; ProgramId(axis) → gid_{x,y,z}. // Same launch geometry as Metal's dispatchThreadgroups, so // no extra bounds guard is needed (parity with Metal). - writeln!(out, " const unsigned int gid_x = blockIdx.x * blockDim.x + threadIdx.x;").ok(); - writeln!(out, " const unsigned int gid_y = blockIdx.y * blockDim.y + threadIdx.y;").ok(); - writeln!(out, " const unsigned int gid_z = blockIdx.z * blockDim.z + threadIdx.z;").ok(); + writeln!( + out, + " const unsigned int gid_x = blockIdx.x * blockDim.x + threadIdx.x;" + ) + .ok(); + writeln!( + out, + " const unsigned int gid_y = blockIdx.y * blockDim.y + threadIdx.y;" + ) + .ok(); + writeln!( + out, + " const unsigned int gid_z = blockIdx.z * blockDim.z + threadIdx.z;" + ) + .ok(); self.emit_simd_aliases(out); - } + }, KernelMode::SimdGroup2D => { // tid = threadgroup_position_in_grid (→ blockIdx), // lid = thread_position_in_threadgroup (→ threadIdx). @@ -351,20 +375,28 @@ impl CudaGenerator { writeln!(out, " const unsigned int lid_y = threadIdx.y;").ok(); writeln!(out, " const unsigned int lid_z = threadIdx.z;").ok(); self.emit_simd_aliases(out); - } + }, KernelMode::Tile2D => { return Err(Error::UnsupportedOp( "cuda: KernelMode Tile2D (Op::Dot tiled-matmul path → wmma, Phase 3)".into(), )); - } + }, } // Apple simdgroup_matrix lane→element coords (from the // mma_layout_probe): elem 0 at (fm, fn0), elem 1 at (fm, fn1). if uses_simdgroup(kernel) { writeln!(out, " const unsigned int _sg_qid = simd_lane / 4u;").ok(); - writeln!(out, " const unsigned int _sg_fm = (_sg_qid & 4u) + ((simd_lane / 2u) % 4u);").ok(); - writeln!(out, " const unsigned int _sg_fn0 = (_sg_qid & 2u) * 2u + (simd_lane % 2u) * 2u;").ok(); + writeln!( + out, + " const unsigned int _sg_fm = (_sg_qid & 4u) + ((simd_lane / 2u) % 4u);" + ) + .ok(); + writeln!( + out, + " const unsigned int _sg_fn0 = (_sg_qid & 2u) * 2u + (simd_lane % 2u) * 2u;" + ) + .ok(); writeln!(out, " const unsigned int _sg_fn1 = _sg_fn0 + 1u;").ok(); } @@ -447,13 +479,18 @@ impl CudaGenerator { for (i, op) in blk.ops.iter().enumerate() { match op { Op::ThreadgroupAlloc { dtype, size, name } => { - v.push((name.clone(), cuda_type_name(*dtype), *size, false, dtype.size_bytes())); - } - Op::SimdgroupAlloc { .. } => { + v.push(( + name.clone(), + cuda_type_name(*dtype), + *size, + false, + dtype.size_bytes(), + )); + }, + Op::SimdgroupAlloc { .. } => if let Some(Some(vid)) = blk.results.get(i) { v.push((sgm_name(*vid), "float", 64, true, 4)); - } - } + }, Op::CoopTileSetup { name, m, n, k, exec_scope, .. } => { let pw = matches!(exec_scope, CoopTileScope::SimdGroup); let nm = ct_ident(name); @@ -473,8 +510,8 @@ impl CudaGenerator { v.push((format!("_CTC_{cnm}"), "float", m * n, pw, 4)); } } - } - _ => {} + }, + _ => {}, } } } @@ -565,14 +602,22 @@ impl CudaGenerator { // Elementwise: the single linear thread id. KernelMode::Elementwise => "_gtid".to_string(), // Grid3D: per-axis global thread id. - KernelMode::Grid3D => { - match axis { 0 => "gid_x", 1 => "gid_y", _ => "gid_z" }.to_string() + KernelMode::Grid3D => match axis { + 0 => "gid_x", + 1 => "gid_y", + _ => "gid_z", } + .to_string(), // Reduction & SimdGroup2D: which threadgroup (block) this is. - _ => match axis { 0 => "tgid_x", 1 => "tgid_y", _ => "tgid_z" }.to_string(), + _ => match axis { + 0 => "tgid_x", + 1 => "tgid_y", + _ => "tgid_z", + } + .to_string(), }; writeln!(out, "{pad}unsigned int {v} = {src};").ok(); - } + }, Op::Const { value } => { let v = self.vname(vid, block, ov); if *value >= 0 { @@ -580,7 +625,7 @@ impl CudaGenerator { } else { writeln!(out, "{pad}int {v} = {value};").ok(); } - } + }, Op::Load { src, indices, .. } => { let v = self.vname(vid, block, ov); // DSL builtin aliases. `simd_id` → warp index. In SimdGroup2D, @@ -602,9 +647,10 @@ impl CudaGenerator { // coalesced accesses. Requires __restrict__ on the parameter // (added in emit_param). Only apply to Tensor params (not // builtins like simd_id, tid_x, …). - let ro_param = kernel.params.iter().find(|p| { - p.name == *src && p.kind == ParamKind::Tensor && !p.is_output - }); + let ro_param = kernel + .params + .iter() + .find(|p| p.name == *src && p.kind == ParamKind::Tensor && !p.is_output); if let Some(p) = ro_param { // Cache-streaming (`__ldcs`, `.cs` hint) for the BIG read- // ONCE Q4 weights (always U32-packed): they're streamed @@ -613,7 +659,8 @@ impl CudaGenerator { // `.cs` marks the line evict-first so the weight stream // doesn't thrash L2. Gated by MT_LDCS_WEIGHTS (default off); // activations/scales (T/f16) stay __ldg (reused across lanes). - let ldcs = p.dtype == DType::U32 && std::env::var("MT_LDCS_WEIGHTS").is_ok(); + let ldcs = + p.dtype == DType::U32 && std::env::var("MT_LDCS_WEIGHTS").is_ok(); if ldcs { writeln!(out, "{pad}auto {v} = __ldcs(&{src}[{idx}]);").ok(); } else { @@ -623,12 +670,12 @@ impl CudaGenerator { writeln!(out, "{pad}auto {v} = {src}[{idx}];").ok(); } } - } + }, Op::Store { dst, indices, value, .. } => { let val = self.vname(Some(*value), block, ov); let idx = self.emit_idx(indices, block, ov, kernel, dst)?; writeln!(out, "{pad}{dst}[{idx}] = {val};").ok(); - } + }, Op::BinOp { op: bop, lhs, rhs } => { let v = self.vname(vid, block, ov); let l = self.vname(Some(*lhs), block, ov); @@ -644,43 +691,43 @@ impl CudaGenerator { } else { writeln!(out, "{pad}auto {v} = {};", cuda_binop(*bop, &l, &r)).ok(); } - } + }, Op::Fma { a, b, c } => { let v = self.vname(vid, block, ov); let av = self.vname(Some(*a), block, ov); let bv = self.vname(Some(*b), block, ov); let cv = self.vname(Some(*c), block, ov); writeln!(out, "{pad}auto {v} = fmaf({av}, {bv}, {cv});").ok(); - } + }, Op::UnaryOp { op: uop, value } => { let v = self.vname(vid, block, ov); let rv = self.vname(Some(*value), block, ov); writeln!(out, "{pad}auto {v} = {};", self.cuda_unary(*uop, &rv)).ok(); - } + }, Op::Cast { value, dtype } => { let v = self.vname(vid, block, ov); let rv = self.vname(Some(*value), block, ov); let ty = cuda_type_name(*dtype); // CUDA C-style cast; source type comes from `auto` inference. writeln!(out, "{pad}{ty} {v} = ({ty})({rv});").ok(); - } + }, Op::Select { cond, on_true, on_false } => { let v = self.vname(vid, block, ov); let c = self.vname(Some(*cond), block, ov); let a = self.vname(Some(*on_true), block, ov); let b = self.vname(Some(*on_false), block, ov); writeln!(out, "{pad}auto {v} = ({c}) ? ({a}) : ({b});").ok(); - } + }, // Mutable locals: the macro lowers `let mut x` to DeclareLocal and // reads to Load{src:"__ml_x"} (body.rs), so the names line up. Op::DeclareLocal { name, value } => { let rv = self.vname(Some(*value), block, ov); writeln!(out, "{pad}auto __ml_{name} = {rv};").ok(); - } + }, Op::SetLocal { name, value } => { let rv = self.vname(Some(*value), block, ov); writeln!(out, "{pad}__ml_{name} = {rv};").ok(); - } + }, Op::Activation { kind, value } => { let v = self.vname(vid, block, ov); let rv = self.vname(Some(*value), block, ov); @@ -692,7 +739,7 @@ impl CudaGenerator { ActKind::Tanh => format!("tanhf((float)({rv}))"), }; writeln!(out, "{pad}auto {v} = {expr};").ok(); - } + }, // Single-warp SIMD reduction → xor butterfly so ALL lanes get // the result (matches Metal `simd_*` broadcast semantics). // @@ -717,7 +764,8 @@ impl CudaGenerator { // different rounding path than the CPU oracle uses). writeln!(out, "{pad} volatile float _s = 0.0f;").ok(); writeln!(out, "{pad} for (int _i = 0; _i < {lw}; _i++) {{").ok(); - writeln!(out, "{pad} _s = _s + __shfl_sync(0xffffffffu, _src, _i);").ok(); + writeln!(out, "{pad} _s = _s + __shfl_sync(0xffffffffu, _src, _i);") + .ok(); writeln!(out, "{pad} }}").ok(); if *rk == ReduceKind::Mean { writeln!(out, "{pad} {v} = _s / {lw}.0f;").ok(); @@ -731,7 +779,8 @@ impl CudaGenerator { writeln!(out, "{pad}{{").ok(); writeln!(out, "{pad} float _s = (float)({rv});").ok(); writeln!(out, "{pad} for (int _o = {half}; _o > 0; _o >>= 1) {{").ok(); - writeln!(out, "{pad} float _x = __shfl_xor_sync(0xffffffffu, _s, _o);").ok(); + writeln!(out, "{pad} float _x = __shfl_xor_sync(0xffffffffu, _s, _o);") + .ok(); writeln!(out, "{pad} _s = {};", reduce_combine(*rk, "_s", "_x")).ok(); writeln!(out, "{pad} }}").ok(); if *rk == ReduceKind::Mean { @@ -741,7 +790,7 @@ impl CudaGenerator { } writeln!(out, "{pad}}}").ok(); } - } + }, // Warp prefix scan (Hillis-Steele via __shfl_up_sync), matching // Metal's simd_prefix_{inclusive,exclusive}_*. Op::SimdScan { value, op: rk, exclusive } => { @@ -753,33 +802,51 @@ impl CudaGenerator { writeln!(out, "{pad} float _s = (float)({rv});").ok(); writeln!(out, "{pad} for (int _o = 1; _o < {lw}; _o <<= 1) {{").ok(); writeln!(out, "{pad} float _t = __shfl_up_sync(0xffffffffu, _s, _o);").ok(); - writeln!(out, "{pad} if (simd_lane >= (unsigned int)_o) _s = {};", reduce_combine(*rk, "_s", "_t")).ok(); + writeln!( + out, + "{pad} if (simd_lane >= (unsigned int)_o) _s = {};", + reduce_combine(*rk, "_s", "_t") + ) + .ok(); writeln!(out, "{pad} }}").ok(); if *exclusive { // Shift right one lane; lane 0 gets the identity. writeln!(out, "{pad} float _x = __shfl_up_sync(0xffffffffu, _s, 1);").ok(); - writeln!(out, "{pad} {v} = (simd_lane == 0u) ? ({}) : _x;", reduce_init(*rk)).ok(); + writeln!( + out, + "{pad} {v} = (simd_lane == 0u) ? ({}) : _x;", + reduce_init(*rk) + ) + .ok(); } else { writeln!(out, "{pad} {v} = _s;").ok(); } writeln!(out, "{pad}}}").ok(); - } + }, // Scalar (rank-0) splat / zeros. Tile (array) shapes need the // threadgroup/stack allocation path (cooperative kernels, P3+). Op::Zeros { shape, .. } if shape.rank() == 0 => { let v = self.vname(vid, block, ov); writeln!(out, "{pad}float {v} = 0.0f;").ok(); - } + }, Op::Splat { value, shape, .. } if shape.rank() == 0 => { let v = self.vname(vid, block, ov); writeln!(out, "{pad}float {v} = {value}f;").ok(); - } + }, // Per-thread grid-stride accumulation over a row (Phase 2). // Each thread sums src[offset+tid], src[offset+tid+lsize], … // Correctness-first: the 4-wide vectorized form the MSL path uses // is a Phase-3 perf retune (SCOPE §6). Op::StrideReduce { - src, offset, stride, end, op: rk, transform, secondary_src, secondary_base, .. + src, + offset, + stride, + end, + op: rk, + transform, + secondary_src, + secondary_base, + .. } => { let v = self.vname(vid, block, ov); let off = self.vname(Some(*offset), block, ov); @@ -803,7 +870,7 @@ impl CudaGenerator { Some(sec) => { let bv = self.vname(*secondary_base, block, ov); format!("(float)({src}[_i]) * (float)({sec}[_i - {bv}])") - } + }, None => format!("(float)({src}[_i])"), }; // Per-element transform chain (e.g. square for sum-of-squares, @@ -824,7 +891,7 @@ impl CudaGenerator { }, Op::Cast { dtype, .. } => { format!("({})({e})", cuda_type_name(*dtype)) - } + }, Op::BinOp { op, rhs, .. } => { let rv = self.vname(Some(*rhs), block, ov); match op { @@ -834,23 +901,20 @@ impl CudaGenerator { BinOpKind::Div => format!("(({e}) / (float)({rv}))"), _ => e, } - } + }, _ => e, }; } e - } + }, }; writeln!(out, "{pad}float {v} = {};", reduce_init(*rk)).ok(); - writeln!( - out, - "{pad}for (unsigned int _i = {start}; _i < {en}; _i += {step}) {{" - ) - .ok(); + writeln!(out, "{pad}for (unsigned int _i = {start}; _i < {en}; _i += {step}) {{") + .ok(); writeln!(out, "{pad} float _e = {elem_expr};").ok(); writeln!(out, "{pad} {v} = {};", reduce_combine(*rk, &v, "_e")).ok(); writeln!(out, "{pad}}}").ok(); - } + }, // Threadgroup-scope reduction: warp shuffle + shared-mem tree // (the CUDA analog of msl::reduce two-level simd_sum). Phase 2. Op::Reduce { value, axis, op: rk } => { @@ -862,43 +926,42 @@ impl CudaGenerator { let v = self.vname(vid, block, ov); let input = self.vname(Some(*value), block, ov); self.emit_reduce_tree(&v, &input, *rk, out); - } + }, // ── Threadgroup / warp primitives (non-MMA) ──────────────── // Allocs are hoisted by emit_allocs; the ops themselves are no-ops. - Op::ThreadgroupAlloc { .. } | Op::StackAlloc { .. } => {} + Op::ThreadgroupAlloc { .. } | Op::StackAlloc { .. } => {}, Op::ThreadgroupLoad { name, index } | Op::StackLoad { name, index } => { let v = self.vname(vid, block, ov); let iv = self.vname(Some(*index), block, ov); writeln!(out, "{pad}auto {v} = {name}[{iv}];").ok(); - } - Op::ThreadgroupStore { name, index, value } - | Op::StackStore { name, index, value } => { + }, + Op::ThreadgroupStore { name, index, value } | Op::StackStore { name, index, value } => { let iv = self.vname(Some(*index), block, ov); let rv = self.vname(Some(*value), block, ov); writeln!(out, "{pad}{name}[{iv}] = {rv};").ok(); - } + }, Op::SimdLaneId => { let v = self.vname(vid, block, ov); writeln!(out, "{pad}unsigned int {v} = simd_lane;").ok(); - } + }, Op::SimdGroupId => { let v = self.vname(vid, block, ov); writeln!(out, "{pad}unsigned int {v} = simd_group;").ok(); - } + }, Op::SimdBroadcast { value, lane } => { let v = self.vname(vid, block, ov); let rv = self.vname(Some(*value), block, ov); let rl = self.vname(Some(*lane), block, ov); writeln!(out, "{pad}auto {v} = __shfl_sync(0xffffffffu, {rv}, {rl});").ok(); - } + }, Op::SimdShuffleXor { value, mask } => { let v = self.vname(vid, block, ov); let rv = self.vname(Some(*value), block, ov); writeln!(out, "{pad}auto {v} = __shfl_xor_sync(0xffffffffu, {rv}, {mask}u);").ok(); - } + }, Op::Barrier => { writeln!(out, "{pad}__syncthreads();").ok(); - } + }, Op::Atomic { op: ak, dst, index, value, .. } => { let iv = self.vname(Some(*index), block, ov); let rv = self.vname(Some(*value), block, ov); @@ -911,27 +974,27 @@ impl CudaGenerator { AtomicKind::Xor => "atomicXor", }; writeln!(out, "{pad}{f}(&{dst}[{iv}], {rv});").ok(); - } + }, Op::SimdgroupBarrier => { writeln!(out, "{pad}__syncwarp();").ok(); - } + }, // ── simdgroup_matrix (software emulation) ───────── // Each fragment is a per-warp 8×8 row-major shared tile; lanes // map to elements via the Apple layout (_sg_fm/_sg_fn0/_sg_fn1). // Cooperative across the 32-lane warp = 32-lane simdgroup. - Op::SimdgroupAlloc { .. } => {} // declared+zeroed in emit_allocs + Op::SimdgroupAlloc { .. } => {}, // declared+zeroed in emit_allocs Op::SimdgroupElemLoad { value, index } => { let v = self.vname(vid, block, ov); let m = sgm_name(*value); let fnk = if *index == 0 { "_sg_fn0" } else { "_sg_fn1" }; writeln!(out, "{pad}float {v} = {m}[simd_group * 64u + _sg_fm * 8u + {fnk}];").ok(); - } + }, Op::SimdgroupElemStore { value, index, data } => { let m = sgm_name(*value); let dv = self.vname(Some(*data), block, ov); let fnk = if *index == 0 { "_sg_fn0" } else { "_sg_fn1" }; writeln!(out, "{pad}{m}[simd_group * 64u + _sg_fm * 8u + {fnk}] = {dv};").ok(); - } + }, // Cooperative fragment load from a threadgroup array, one row per // (fm,fn): dest[fm][fn] = tg[off + (transpose ? fn*stride+fm : fm*stride+fn)]. Op::SimdgroupLoad { dest, tg, offset, stride, transpose } => { @@ -947,10 +1010,11 @@ impl CudaGenerator { writeln!( out, "{pad} {m}[_base + _sg_fm * 8u + {fnk}] = (float)({tg}[{off} + {idx}]);" - ).ok(); + ) + .ok(); } writeln!(out, "{pad}}}").ok(); - } + }, // C += A * B (8×8×8), cooperative across the warp. Op::SimdgroupMatMul { a, b, c } => { let (ma, mb, mc) = (sgm_name(*a), sgm_name(*b), sgm_name(*c)); @@ -967,20 +1031,19 @@ impl CudaGenerator { writeln!(out, "{pad} {mc}[_bs + _sg_fm * 8u + _sg_fn0] = _acc0;").ok(); writeln!(out, "{pad} {mc}[_bs + _sg_fm * 8u + _sg_fn1] = _acc1;").ok(); writeln!(out, "{pad}}}").ok(); - } + }, // ── CoopTile cooperative GEMM (MPP/NAX → software emulation) ── // Opaque tiles (Load→Run→Store, no element access) → staged in // shared memory with a self-consistent row-major layout. C // persists across K-block Run calls (accumulate mode). - Op::CoopTileSetup { .. } => {} // tiles declared in emit_allocs + Op::CoopTileSetup { .. } => {}, // tiles declared in emit_allocs Op::CoopTileZero { name } => { - let (m, n, _, _, _, _, _, simd) = self.coop_cfg(kernel, name) - .ok_or_else(|| Error::UnsupportedOp(format!("cuda: CoopTile `{name}` no Setup")))?; + let (m, n, _, _, _, _, _, simd) = self.coop_cfg(kernel, name).ok_or_else(|| { + Error::UnsupportedOp(format!("cuda: CoopTile `{name}` no Setup")) + })?; let cnm = ct_ident(&coop_c_name(name)); - let local_c = matches!( - self.profile.mma, - crate::backend::MmaStrategy::SoftwareLocalC - ); + let local_c = + matches!(self.profile.mma, crate::backend::MmaStrategy::SoftwareLocalC); if local_c && simd { // Lane-local C: each lane zeros its own slots. let lw = self.profile.lane_width; @@ -995,18 +1058,22 @@ impl CudaGenerator { let base = coop_base(simd, m * n); writeln!(out, "{pad}for (unsigned int _e = {sid}; _e < {}u; _e += {ssize}) _CTC_{cnm}[{base} + _e] = 0.0f;", m * n).ok(); } - } + }, // Per Apple MPP headers: metal::extents is {inner, // outer}; with tensor_inline the leading-dim stride ld == EI (the // inner extent), and the *descriptor flag* (ta/tb) — not the // extents — chooses transpose. So stride = ei always. Op::CoopTileLoadA { name, ptr_name, ptr_offset, dtype: _, ei, .. } => { - let (m, _, k, ta, _, _, _, simd) = self.coop_cfg(kernel, name) - .ok_or_else(|| Error::UnsupportedOp(format!("cuda: CoopTile `{name}` no Setup")))?; + let (m, _, k, ta, _, _, _, simd) = + self.coop_cfg(kernel, name).ok_or_else(|| { + Error::UnsupportedOp(format!("cuda: CoopTile `{name}` no Setup")) + })?; let nm = ct_ident(name); let (sid, ssize, _) = coop_scope(simd); let base = coop_base(simd, m * k); - let off = ptr_offset.map(|o| self.vname(Some(o), block, ov)).unwrap_or_else(|| "0".into()); + let off = ptr_offset + .map(|o| self.vname(Some(o), block, ov)) + .unwrap_or_else(|| "0".into()); // A is m×k, leading dim ei. ta=false: ptr[i*ei+j]; ta=true: // memory is k×m → ptr[j*ei+i]. (i=_e/k row, j=_e%k col.) let src = if ta { @@ -1016,14 +1083,18 @@ impl CudaGenerator { }; writeln!(out, "{pad}__syncthreads();").ok(); writeln!(out, "{pad}for (unsigned int _e = {sid}; _e < {}u; _e += {ssize}) _CTA_{nm}[{base} + _e] = (float)({ptr_name}[{off} + {src}]);", m * k).ok(); - } + }, Op::CoopTileLoadB { name, ptr_name, ptr_offset, dtype: _, ei, .. } => { - let (_, n, k, _, tb, _, _, simd) = self.coop_cfg(kernel, name) - .ok_or_else(|| Error::UnsupportedOp(format!("cuda: CoopTile `{name}` no Setup")))?; + let (_, n, k, _, tb, _, _, simd) = + self.coop_cfg(kernel, name).ok_or_else(|| { + Error::UnsupportedOp(format!("cuda: CoopTile `{name}` no Setup")) + })?; let nm = ct_ident(name); let (sid, ssize, _) = coop_scope(simd); let base = coop_base(simd, k * n); - let off = ptr_offset.map(|o| self.vname(Some(o), block, ov)).unwrap_or_else(|| "0".into()); + let off = ptr_offset + .map(|o| self.vname(Some(o), block, ov)) + .unwrap_or_else(|| "0".into()); // B is k×n, leading dim ei. tb=false: ptr[i*ei+j]; tb=true: // memory is n×k → ptr[j*ei+i]. (i=_e/n row, j=_e%n col.) let src = if tb { @@ -1032,16 +1103,16 @@ impl CudaGenerator { format!("(_e / {n}u) * {ei}u + (_e % {n}u)") }; writeln!(out, "{pad}for (unsigned int _e = {sid}; _e < {}u; _e += {ssize}) _CTB_{nm}[{base} + _e] = (float)({ptr_name}[{off} + {src}]);", k * n).ok(); - } + }, Op::CoopTileRun { name, .. } => { - let (m, n, k, _, _, _, accum, simd) = self.coop_cfg(kernel, name) - .ok_or_else(|| Error::UnsupportedOp(format!("cuda: CoopTile `{name}` no Setup")))?; + let (m, n, k, _, _, _, accum, simd) = + self.coop_cfg(kernel, name).ok_or_else(|| { + Error::UnsupportedOp(format!("cuda: CoopTile `{name}` no Setup")) + })?; let nm = ct_ident(name); let cnm = ct_ident(&coop_c_name(name)); - let local_c = matches!( - self.profile.mma, - crate::backend::MmaStrategy::SoftwareLocalC - ); + let local_c = + matches!(self.profile.mma, crate::backend::MmaStrategy::SoftwareLocalC); if local_c && simd { // SoftwareLocalC: per-warp lane-local C. Each lane // walks its own `m*n/lw` slots, computing the same @@ -1052,15 +1123,13 @@ impl CudaGenerator { let ba = coop_base(true, m * k); let bb = coop_base(true, k * n); writeln!(out, "{pad}__syncthreads();").ok(); - writeln!(out, "{pad}for (unsigned int _li = 0u; _li < {per_lane}u; _li++) {{").ok(); + writeln!(out, "{pad}for (unsigned int _li = 0u; _li < {per_lane}u; _li++) {{") + .ok(); writeln!(out, "{pad} unsigned int _e = simd_lane + _li * {lw}u;").ok(); writeln!(out, "{pad} if (_e >= {}u) break;", m * n).ok(); writeln!(out, "{pad} unsigned int _i = _e / {n}u, _j = _e % {n}u;").ok(); - let init = if accum { - format!("_CTC_lane_{cnm}[_li]") - } else { - "0.0f".to_string() - }; + let init = + if accum { format!("_CTC_lane_{cnm}[_li]") } else { "0.0f".to_string() }; writeln!(out, "{pad} float _acc = {init};").ok(); writeln!(out, "{pad} for (unsigned int _l = 0u; _l < {k}u; _l++) _acc += _CTA_{nm}[{ba} + _i * {k}u + _l] * _CTB_{nm}[{bb} + _l * {n}u + _j];").ok(); writeln!(out, "{pad} _CTC_lane_{cnm}[_li] = _acc;").ok(); @@ -1071,7 +1140,12 @@ impl CudaGenerator { let (ba, bb, bc) = (coop_base(simd, m * k), coop_base(simd, k * n), coop_base(simd, m * n)); writeln!(out, "{pad}__syncthreads();").ok(); - writeln!(out, "{pad}for (unsigned int _e = {sid}; _e < {}u; _e += {ssize}) {{", m * n).ok(); + writeln!( + out, + "{pad}for (unsigned int _e = {sid}; _e < {}u; _e += {ssize}) {{", + m * n + ) + .ok(); writeln!(out, "{pad} unsigned int _i = _e / {n}u, _j = _e % {n}u;").ok(); let init = if accum { format!("_CTC_{cnm}[{bc} + _e]") } else { "0.0f".into() }; writeln!(out, "{pad} float _acc = {init};").ok(); @@ -1080,26 +1154,32 @@ impl CudaGenerator { writeln!(out, "{pad}}}").ok(); writeln!(out, "{pad}__syncthreads();").ok(); } - } + }, Op::CoopTileStoreC { name, ptr_name, ptr_offset, dtype, ei, .. } => { - let (m, n, _, _, _, _, _, simd) = self.coop_cfg(kernel, name) - .ok_or_else(|| Error::UnsupportedOp(format!("cuda: CoopTile `{name}` no Setup")))?; + let (m, n, _, _, _, _, _, simd) = self.coop_cfg(kernel, name).ok_or_else(|| { + Error::UnsupportedOp(format!("cuda: CoopTile `{name}` no Setup")) + })?; let cnm = ct_ident(&coop_c_name(name)); - let off = ptr_offset.map(|o| self.vname(Some(o), block, ov)).unwrap_or_else(|| "0".into()); + let off = ptr_offset + .map(|o| self.vname(Some(o), block, ov)) + .unwrap_or_else(|| "0".into()); let ty = cuda_type_name(*dtype); let dst = format!("(_e / {n}u) * {ei}u + (_e % {n}u)"); - let local_c = matches!( - self.profile.mma, - crate::backend::MmaStrategy::SoftwareLocalC - ); + let local_c = + matches!(self.profile.mma, crate::backend::MmaStrategy::SoftwareLocalC); if local_c && simd { let lw = self.profile.lane_width; let per_lane = (m * n).div_ceil(lw).max(1); writeln!(out, "{pad}__syncthreads();").ok(); - writeln!(out, "{pad}for (unsigned int _li = 0u; _li < {per_lane}u; _li++) {{").ok(); + writeln!(out, "{pad}for (unsigned int _li = 0u; _li < {per_lane}u; _li++) {{") + .ok(); writeln!(out, "{pad} unsigned int _e = simd_lane + _li * {lw}u;").ok(); writeln!(out, "{pad} if (_e >= {}u) break;", m * n).ok(); - writeln!(out, "{pad} {ptr_name}[{off} + {dst}] = ({ty})(_CTC_lane_{cnm}[_li]);").ok(); + writeln!( + out, + "{pad} {ptr_name}[{off} + {dst}] = ({ty})(_CTC_lane_{cnm}[_li]);" + ) + .ok(); writeln!(out, "{pad}}}").ok(); } else { let (sid, ssize, _) = coop_scope(simd); @@ -1108,14 +1188,15 @@ impl CudaGenerator { writeln!(out, "{pad}__syncthreads();").ok(); writeln!(out, "{pad}for (unsigned int _e = {sid}; _e < {}u; _e += {ssize}) {ptr_name}[{off} + {dst}] = ({ty})(_CTC_{cnm}[{base} + _e]);", m * n).ok(); } - } + }, // ── Control flow (nested-block recursion) ────────────────── Op::Loop { var, start, end, step, body } => { let s = self.vname(Some(*start), block, ov); let e = self.vname(Some(*end), block, ov); let st = self.vname(Some(*step), block, ov); let lv = format!("i_{}", var.as_u32()); - writeln!(out, "{pad}for (unsigned int {lv} = {s}; {lv} < {e}; {lv} += {st}) {{").ok(); + writeln!(out, "{pad}for (unsigned int {lv} = {s}; {lv} < {e}; {lv} += {st}) {{") + .ok(); if let Some(bb) = kernel.blocks.get(body) { let mut child = self.child_ov(block, ov); // The loop induction var is referenced inside the body by @@ -1125,7 +1206,7 @@ impl CudaGenerator { self.emit_ops(bb, kernel, &child, out)?; } writeln!(out, "{pad}}}").ok(); - } + }, Op::If { cond, then_block, else_block } => { let c = self.vname(Some(*cond), block, ov); writeln!(out, "{pad}if ({c}) {{").ok(); @@ -1141,13 +1222,13 @@ impl CudaGenerator { } } writeln!(out, "{pad}}}").ok(); - } + }, other => { return Err(Error::UnsupportedOp(format!( "cuda: op {} not supported yet", op_variant_name(other) ))); - } + }, } Ok(()) } @@ -1178,7 +1259,8 @@ impl CudaGenerator { writeln!(out, "{pad} __syncthreads();").ok(); // Phase 3: warp 0 reduces the per-warp totals and broadcasts via [0]. writeln!(out, "{pad} if (simd_group == 0u) {{").ok(); - writeln!(out, "{pad} float _wv = simd_lane < n_simd ? {buf}[simd_lane] : {init};").ok(); + writeln!(out, "{pad} float _wv = simd_lane < n_simd ? {buf}[simd_lane] : {init};") + .ok(); writeln!(out, "{pad} for (int _o = {half}; _o > 0; _o >>= 1) {{").ok(); writeln!(out, "{pad} float _ov = __shfl_down_sync(0xffffffffu, _wv, _o);").ok(); writeln!(out, "{pad} _wv = {};", reduce_combine(kind, "_wv", "_ov")).ok(); @@ -1207,14 +1289,18 @@ impl CudaGenerator { if let Op::CoopTileSetup { name: nm, m, n, k, ta, tb, tc, acc_mode, exec_scope, .. } = op + && nm == name { - if nm == name { - return Some(( - *m, *n, *k, *ta, *tb, *tc, - matches!(acc_mode, CoopTileAccMode::MultiplyAccumulate), - matches!(exec_scope, CoopTileScope::SimdGroup), - )); - } + return Some(( + *m, + *n, + *k, + *ta, + *tb, + *tc, + matches!(acc_mode, CoopTileAccMode::MultiplyAccumulate), + matches!(exec_scope, CoopTileScope::SimdGroup), + )); } } } @@ -1234,7 +1320,7 @@ impl CudaGenerator { Rsqrt => return format!("mt_rsqrt((float)({arg}))"), Exp => return format!("mt_expf((float)({arg}))"), Log => return format!("mt_logf((float)({arg}))"), - _ => {} + _ => {}, } } match op { @@ -1242,7 +1328,8 @@ impl CudaGenerator { Recip => format!("(1.0f / {arg})"), // sign(x): -1/0/+1 (matches Metal `sign`, incl. 0 → 0). // `copysignf` would map 0 → +1, so emit the explicit form. - Sign => format!("((float)({arg}) > 0.0f ? 1.0f : ((float)({arg}) < 0.0f ? -1.0f : 0.0f))"), + Sign => + format!("((float)({arg}) > 0.0f ? 1.0f : ((float)({arg}) < 0.0f ? -1.0f : 0.0f))"), // Block-scaled-quant decode helpers (preamble __device__ fns). _ => format!("{}({arg})", self.profile.unary_intrinsic(op)), } @@ -1259,10 +1346,9 @@ impl CodegenBackend for CudaGenerator { // inline pass. Non-call kernels are unaffected; on failure we fall // back to the raw IR. let mut inlined = kernel.clone(); - let k: &Kernel = match crate::passes::run_passes( - &mut inlined, - &[Box::new(crate::passes::kernel_inline::KernelInlinePass)], - ) { + let k: &Kernel = match crate::passes::run_passes(&mut inlined, &[Box::new( + crate::passes::kernel_inline::KernelInlinePass, + )]) { Ok(()) => &inlined, Err(_) => kernel, }; @@ -1361,9 +1447,7 @@ fn sgm_name(v: ValueId) -> String { format!("_SGM_{}", v.as_u32()) } /// The C-accumulator group name for a CoopTile setup: a `*_acc` /// (multiply-accumulate) setup shares its C tile with the base setup /// (e.g. `qk_acc` accumulates onto `qk`'s C across head_dim chunks). -fn coop_c_name(name: &str) -> String { - name.strip_suffix("_acc").unwrap_or(name).to_string() -} +fn coop_c_name(name: &str) -> String { name.strip_suffix("_acc").unwrap_or(name).to_string() } /// Sanitize a CoopTile name into a valid C identifier suffix. fn ct_ident(name: &str) -> String { @@ -1405,8 +1489,10 @@ fn op_variant_name(op: &Op) -> String { #[cfg(test)] mod tests { - use metaltile_core::ir::{BinOpKind, IndexExpr, Op}; - use metaltile_core::shape::Shape; + use metaltile_core::{ + ir::{BinOpKind, IndexExpr, Op}, + shape::Shape, + }; use super::*; @@ -1425,12 +1511,22 @@ mod tests { k.body.push_op(Op::ProgramId { axis: 0 }, ValueId::new(0)); k.body.name_value(ValueId::new(0), "idx"); k.body.push_op( - Op::Load { src: "a".into(), indices: vec![IndexExpr::Value(ValueId::new(0))], mask: None, other: None }, + Op::Load { + src: "a".into(), + indices: vec![IndexExpr::Value(ValueId::new(0))], + mask: None, + other: None, + }, ValueId::new(1), ); k.body.name_value(ValueId::new(1), "x"); k.body.push_op( - Op::Load { src: "b".into(), indices: vec![IndexExpr::Value(ValueId::new(0))], mask: None, other: None }, + Op::Load { + src: "b".into(), + indices: vec![IndexExpr::Value(ValueId::new(0))], + mask: None, + other: None, + }, ValueId::new(2), ); k.body.name_value(ValueId::new(2), "y"); @@ -1473,12 +1569,18 @@ mod tests { let mut k = Kernel::new("row_reduce_sum"); k.mode = KernelMode::Reduction; k.params.push(Param { - name: "inp".into(), dtype: DType::F32, shape: Shape::scalar(), - is_output: false, kind: ParamKind::Tensor, + name: "inp".into(), + dtype: DType::F32, + shape: Shape::scalar(), + is_output: false, + kind: ParamKind::Tensor, }); k.params.push(Param { - name: "out".into(), dtype: DType::F32, shape: Shape::scalar(), - is_output: true, kind: ParamKind::Tensor, + name: "out".into(), + dtype: DType::F32, + shape: Shape::scalar(), + is_output: true, + kind: ParamKind::Tensor, }); k.constexprs.push(metaltile_core::ir::ConstExprDecl { name: metaltile_core::constexpr::ConstExpr::new("n"), @@ -1486,8 +1588,12 @@ mod tests { value: None, }); let (row, nv, rs, re, acc, res) = ( - ValueId::new(0), ValueId::new(1), ValueId::new(2), - ValueId::new(3), ValueId::new(4), ValueId::new(5), + ValueId::new(0), + ValueId::new(1), + ValueId::new(2), + ValueId::new(3), + ValueId::new(4), + ValueId::new(5), ); k.body.push_op(Op::ProgramId { axis: 0 }, row); k.body.name_value(row, "row"); @@ -1498,9 +1604,15 @@ mod tests { k.body.name_value(re, "re"); k.body.push_op( Op::StrideReduce { - src: "inp".into(), offset: rs, stride: nv, end: re, - op: ReduceKind::Sum, dtype: DType::F32, - transform: None, secondary_src: None, secondary_base: None, + src: "inp".into(), + offset: rs, + stride: nv, + end: re, + op: ReduceKind::Sum, + dtype: DType::F32, + transform: None, + secondary_src: None, + secondary_base: None, }, acc, ); @@ -1508,7 +1620,10 @@ mod tests { k.body.push_op(Op::Reduce { value: acc, axis: 0, op: ReduceKind::Sum }, res); k.body.name_value(res, "result"); k.body.push_op_no_result(Op::Store { - dst: "out".into(), indices: vec![IndexExpr::Value(row)], value: res, mask: None, + dst: "out".into(), + indices: vec![IndexExpr::Value(row)], + value: res, + mask: None, }); k } diff --git a/crates/metaltile-codegen/src/hip/mod.rs b/crates/metaltile-codegen/src/hip/mod.rs index 47f9ada4..df5a0fee 100644 --- a/crates/metaltile-codegen/src/hip/mod.rs +++ b/crates/metaltile-codegen/src/hip/mod.rs @@ -51,9 +51,7 @@ impl Default for HipGenerator { } impl HipGenerator { - pub fn new() -> Self { - Self::with_profile(TargetProfile::hip()) - } + pub fn new() -> Self { Self::with_profile(TargetProfile::hip()) } /// Pin to a specific HIP profile. For [`TargetProfile::hip_wave64`] /// (CDNA — MI200/MI300/MI350), we *also* swap the inner CUDA @@ -172,12 +170,22 @@ mod tests { k.body.push_op(Op::ProgramId { axis: 0 }, ValueId::new(0)); k.body.name_value(ValueId::new(0), "idx"); k.body.push_op( - Op::Load { src: "a".into(), indices: vec![IndexExpr::Value(ValueId::new(0))], mask: None, other: None }, + Op::Load { + src: "a".into(), + indices: vec![IndexExpr::Value(ValueId::new(0))], + mask: None, + other: None, + }, ValueId::new(1), ); k.body.name_value(ValueId::new(1), "x"); k.body.push_op( - Op::Load { src: "b".into(), indices: vec![IndexExpr::Value(ValueId::new(0))], mask: None, other: None }, + Op::Load { + src: "b".into(), + indices: vec![IndexExpr::Value(ValueId::new(0))], + mask: None, + other: None, + }, ValueId::new(2), ); k.body.name_value(ValueId::new(2), "y"); @@ -234,17 +242,25 @@ mod tests { /// textually. We don't run it (no CDNA hardware available); we just /// inspect the source. fn row_reduce_sum_ir() -> Kernel { - use metaltile_core::constexpr::ConstExpr; - use metaltile_core::ir::{ConstExprDecl, KernelMode, ReduceKind}; + use metaltile_core::{ + constexpr::ConstExpr, + ir::{ConstExprDecl, KernelMode, ReduceKind}, + }; let mut k = Kernel::new("row_reduce_sum"); k.mode = KernelMode::Reduction; k.params.push(Param { - name: "inp".into(), dtype: DType::F32, shape: Shape::scalar(), - is_output: false, kind: ParamKind::Tensor, + name: "inp".into(), + dtype: DType::F32, + shape: Shape::scalar(), + is_output: false, + kind: ParamKind::Tensor, }); k.params.push(Param { - name: "out".into(), dtype: DType::F32, shape: Shape::scalar(), - is_output: true, kind: ParamKind::Tensor, + name: "out".into(), + dtype: DType::F32, + shape: Shape::scalar(), + is_output: true, + kind: ParamKind::Tensor, }); k.constexprs.push(ConstExprDecl { name: ConstExpr::new("n"), @@ -252,8 +268,12 @@ mod tests { value: None, }); let (row, nv, rs, re, acc, res) = ( - ValueId::new(0), ValueId::new(1), ValueId::new(2), - ValueId::new(3), ValueId::new(4), ValueId::new(5), + ValueId::new(0), + ValueId::new(1), + ValueId::new(2), + ValueId::new(3), + ValueId::new(4), + ValueId::new(5), ); k.body.push_op(Op::ProgramId { axis: 0 }, row); k.body.name_value(row, "row"); @@ -264,9 +284,15 @@ mod tests { k.body.name_value(re, "re"); k.body.push_op( Op::StrideReduce { - src: "inp".into(), offset: rs, stride: nv, end: re, - op: ReduceKind::Sum, dtype: DType::F32, - transform: None, secondary_src: None, secondary_base: None, + src: "inp".into(), + offset: rs, + stride: nv, + end: re, + op: ReduceKind::Sum, + dtype: DType::F32, + transform: None, + secondary_src: None, + secondary_base: None, }, acc, ); @@ -274,7 +300,10 @@ mod tests { k.body.push_op(Op::Reduce { value: acc, axis: 0, op: ReduceKind::Sum }, res); k.body.name_value(res, "result"); k.body.push_op_no_result(Op::Store { - dst: "out".into(), indices: vec![IndexExpr::Value(row)], value: res, mask: None, + dst: "out".into(), + indices: vec![IndexExpr::Value(row)], + value: res, + mask: None, }); k } diff --git a/crates/metaltile-codegen/src/spirv/mod.rs b/crates/metaltile-codegen/src/spirv/mod.rs index dff9a8be..bbb9dbe4 100644 --- a/crates/metaltile-codegen/src/spirv/mod.rs +++ b/crates/metaltile-codegen/src/spirv/mod.rs @@ -47,14 +47,23 @@ //! a small Phase-3 chore. //! - `Strided` params still unsupported. -use std::collections::BTreeMap; -use std::fmt::Write as _; +use std::{collections::BTreeMap, fmt::Write as _}; use metaltile_core::{ dtype::DType, ir::{ - ActKind, BinOpKind, Block, IndexExpr, Kernel, KernelMode, Op, Param, ParamKind, - ReduceKind, UnaryOpKind, ValueId, + ActKind, + BinOpKind, + Block, + IndexExpr, + Kernel, + KernelMode, + Op, + Param, + ParamKind, + ReduceKind, + UnaryOpKind, + ValueId, }, }; @@ -147,7 +156,7 @@ impl GlslGenerator { Some(hint) => format!("v_{hint}_{}", v.as_u32()), None => format!("v{}", v.as_u32()), } - } + }, None => "/**/".to_string(), } } @@ -158,7 +167,7 @@ impl GlslGenerator { IndexExpr::Const(c) => format!("uint({c})"), IndexExpr::Range(v, off) => { format!("({} + uint({off}))", self.vname(Some(*v), block, ov)) - } + }, } } @@ -207,16 +216,13 @@ impl GlslGenerator { return Err(Error::UnsupportedOp(format!( "spirv: multi-dim index needs static dims on `{src}`" ))); - } + }, } } - terms.push(format!( - "({}) * uint({stride})", - self.idx_term(ix, block, ov) - )); + terms.push(format!("({}) * uint({stride})", self.idx_term(ix, block, ov))); } Ok(terms.join(" + ")) - } + }, } } @@ -291,22 +297,29 @@ impl GlslGenerator { // Phase 2.1; landing the real decode unlocks all of them. writeln!(out, "float mt_decode_e2m1(uint code) {{").ok(); writeln!(out, " uint m = code & 7u;").ok(); - writeln!(out, " float mag = (m < 1u) ? 0.0 : (m < 2u) ? 0.5 : (m < 3u) ? 1.0 : (m < 4u) ? 1.5").ok(); - writeln!(out, " : (m < 5u) ? 2.0 : (m < 6u) ? 3.0 : (m < 7u) ? 4.0 : 6.0;").ok(); + writeln!( + out, + " float mag = (m < 1u) ? 0.0 : (m < 2u) ? 0.5 : (m < 3u) ? 1.0 : (m < 4u) ? 1.5" + ) + .ok(); + writeln!(out, " : (m < 5u) ? 2.0 : (m < 6u) ? 3.0 : (m < 7u) ? 4.0 : 6.0;") + .ok(); writeln!(out, " return ((code & 8u) != 0u) ? -mag : mag;").ok(); writeln!(out, "}}").ok(); writeln!(out, "float mt_decode_e4m3(uint bits) {{").ok(); writeln!(out, " uint e = (bits >> 3u) & 15u;").ok(); writeln!(out, " uint m = bits & 7u;").ok(); writeln!(out, " float mag = (e < 1u) ? (float(m) * 0.001953125)").ok(); - writeln!(out, " : ((1.0 + float(m) * 0.125) * exp2(float(int(e)) - 7.0));").ok(); + writeln!(out, " : ((1.0 + float(m) * 0.125) * exp2(float(int(e)) - 7.0));") + .ok(); writeln!(out, " return (((bits >> 7u) & 1u) != 0u) ? -mag : mag;").ok(); writeln!(out, "}}").ok(); writeln!(out, "float mt_decode_e5m2(uint bits) {{").ok(); writeln!(out, " uint e = (bits >> 2u) & 31u;").ok(); writeln!(out, " uint m = bits & 3u;").ok(); writeln!(out, " float mag = (e < 1u) ? (float(m) * 0.0000152587890625)").ok(); - writeln!(out, " : ((1.0 + float(m) * 0.25) * exp2(float(int(e)) - 15.0));").ok(); + writeln!(out, " : ((1.0 + float(m) * 0.25) * exp2(float(int(e)) - 15.0));") + .ok(); writeln!(out, " return (((bits >> 7u) & 1u) != 0u) ? -mag : mag;").ok(); writeln!(out, "}}").ok(); writeln!(out, "float mt_decode_int8(uint bits) {{").ok(); @@ -325,7 +338,11 @@ impl GlslGenerator { writeln!(out, "float mt_erf(float x) {{").ok(); writeln!(out, " float t = 1.0 / (1.0 + 0.3275911 * abs(x));").ok(); writeln!(out, " float y = 1.0 - (((((1.061405429 * t - 1.453152027) * t)").ok(); - writeln!(out, " + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-x*x);").ok(); + writeln!( + out, + " + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-x*x);" + ) + .ok(); writeln!(out, " return sign(x) * y;").ok(); writeln!(out, "}}").ok(); writeln!(out, "float mt_erfinv(float x) {{").ok(); @@ -550,12 +567,7 @@ impl GlslGenerator { && let Some(Some(vid)) = blk.results.get(i) { let count = 64 * n_warps; - writeln!( - out, - "shared float _SGM_{}[{count}];", - vid.as_u32() - ) - .ok(); + writeln!(out, "shared float _SGM_{}[{count}];", vid.as_u32()).ok(); } } } @@ -568,10 +580,7 @@ impl GlslGenerator { // private array (declared inside `main()`), saving // `n_warps * m*n * 4` bytes of shared. The decl itself is emitted // in `emit_body`. - let local_c = matches!( - self.profile.mma, - crate::backend::MmaStrategy::SoftwareLocalC - ); + let local_c = matches!(self.profile.mma, crate::backend::MmaStrategy::SoftwareLocalC); let mut seen_c: std::collections::BTreeSet = std::collections::BTreeSet::new(); for blk in std::iter::once(&kernel.body).chain(kernel.blocks.values()) { for op in &blk.ops { @@ -593,12 +602,7 @@ impl GlslGenerator { // subgroup-distributed registers (declared in // emit_body) so its shared slot is skipped too. if !((local_c || self.coopmat_on()) && pw) { - writeln!( - out, - "shared float _CTC_{cnm}[{}];", - mul * m * n - ) - .ok(); + writeln!(out, "shared float _CTC_{cnm}[{}];", mul * m * n).ok(); } } } @@ -609,11 +613,8 @@ impl GlslGenerator { fn emit_signature(&self, out: &mut String) { let [x, y, z] = self.local_size; - writeln!( - out, - "layout(local_size_x = {x}, local_size_y = {y}, local_size_z = {z}) in;", - ) - .ok(); + writeln!(out, "layout(local_size_x = {x}, local_size_y = {y}, local_size_z = {z}) in;",) + .ok(); writeln!(out).ok(); writeln!(out, "void main() {{").ok(); } @@ -625,13 +626,13 @@ impl GlslGenerator { writeln!(out, " if (_gtid >= pc._n_elems) return;").ok(); writeln!(out, " uint tid = _gtid;").ok(); self.emit_simd_aliases(out); - } + }, KernelMode::Grid3D => { writeln!(out, " uint gid_x = gl_GlobalInvocationID.x;").ok(); writeln!(out, " uint gid_y = gl_GlobalInvocationID.y;").ok(); writeln!(out, " uint gid_z = gl_GlobalInvocationID.z;").ok(); self.emit_simd_aliases(out); - } + }, KernelMode::SimdGroup2D => { // Apple `simdgroup`/`threadgroup_2d` mode — block IDs + // local thread IDs both accessible. Mirrors the CUDA @@ -644,7 +645,7 @@ impl GlslGenerator { writeln!(out, " uint lid_z = gl_LocalInvocationID.z;").ok(); writeln!(out, " uint tid = gl_LocalInvocationIndex;").ok(); self.emit_simd_aliases(out); - } + }, KernelMode::Reduction => { // Block-per-output-row reduction. `tid` is the local // 1-D thread index inside the workgroup; the per-axis @@ -666,12 +667,12 @@ impl GlslGenerator { // (see emit_simd_aliases). writeln!(out, " uint simd_lane = tid % {lw}u;").ok(); writeln!(out, " uint simd_group = tid / {lw}u;").ok(); - } + }, other => { return Err(Error::UnsupportedOp(format!( "spirv: KernelMode::{other:?} not yet supported" ))); - } + }, } // SoftwareLocalC: declare per-warp lane-local C arrays inside // main(). Each lane holds `m*n / lane_width` floats; they @@ -768,17 +769,15 @@ impl GlslGenerator { } fn child_ov(&self, parent: &Block, ov: &Names) -> Names { - let mut child: Names = parent - .names - .iter() - .map(|(&k, v)| (k, format!("v_{v}_{}", k.as_u32()))) - .collect(); + let mut child: Names = + parent.names.iter().map(|(&k, v)| (k, format!("v_{v}_{}", k.as_u32()))).collect(); for (&k, v) in ov { child.insert(k, v.clone()); } child } + #[allow(clippy::too_many_arguments)] fn emit_op( &self, op: &Op, @@ -796,9 +795,12 @@ impl GlslGenerator { let v = self.vname(vid, block, ov); let src = match kernel.mode { KernelMode::Elementwise => "_gtid".to_string(), - KernelMode::Grid3D => { - match axis { 0 => "gid_x", 1 => "gid_y", _ => "gid_z" }.to_string() + KernelMode::Grid3D => match axis { + 0 => "gid_x", + 1 => "gid_y", + _ => "gid_z", } + .to_string(), KernelMode::Reduction | KernelMode::SimdGroup2D => { // SimdGroup2D: each workgroup processes one // tile, so `program_id(axis)` returns the @@ -806,15 +808,20 @@ impl GlslGenerator { // TG saw tgid=0 and wrote on top of itself // (the `fused/gather/masked/segmented/splitk_*` // family cascade). - match axis { 0 => "tgid_x", 1 => "tgid_y", _ => "tgid_z" }.to_string() - } + match axis { + 0 => "tgid_x", + 1 => "tgid_y", + _ => "tgid_z", + } + .to_string() + }, _ => "0u".to_string(), }; writeln!(out, "{pad}uint {v} = {src};").ok(); if let Some(id) = vid { types.insert(id, "uint"); } - } + }, Op::Const { value } => { let v = self.vname(vid, block, ov); if *value >= 0 { @@ -828,7 +835,7 @@ impl GlslGenerator { types.insert(id, "int"); } } - } + }, Op::Load { src, indices, .. } => { let v = self.vname(vid, block, ov); let constexpr = kernel.constexprs.iter().find(|c| c.name.name() == src); @@ -868,7 +875,9 @@ impl GlslGenerator { _ => None, } { writeln!(out, "{pad}uint {v} = {b};").ok(); - if let Some(id) = vid { types.insert(id, "uint"); } + if let Some(id) = vid { + types.insert(id, "uint"); + } } else if let Some(c) = constexpr { // Constexpr at its declared dtype — integer // constexprs MUST stay integer or every kernel @@ -877,41 +886,49 @@ impl GlslGenerator { match c.dtype { DType::F32 => { writeln!(out, "{pad}float {v} = pc.{nm};").ok(); - if let Some(id) = vid { types.insert(id, "float"); } - } + if let Some(id) = vid { + types.insert(id, "float"); + } + }, DType::U32 | DType::U16 | DType::U8 => { writeln!(out, "{pad}uint {v} = pc.{nm};").ok(); - if let Some(id) = vid { types.insert(id, "uint"); } - } + if let Some(id) = vid { + types.insert(id, "uint"); + } + }, DType::I32 | DType::I8 | DType::I4 | DType::I64 => { writeln!(out, "{pad}int {v} = pc.{nm};").ok(); - if let Some(id) = vid { types.insert(id, "int"); } - } + if let Some(id) = vid { + types.insert(id, "int"); + } + }, _ => { writeln!(out, "{pad}float {v} = float(pc.{nm});").ok(); - if let Some(id) = vid { types.insert(id, "float"); } - } + if let Some(id) = vid { + types.insert(id, "float"); + } + }, } } else if let Some(local_name) = src.strip_prefix("__ml_") { // Mutable-local read. Use the declared type so // `rem / extent` (uint local) stays uint. let lt = local_types.get(local_name).copied().unwrap_or("float"); writeln!(out, "{pad}{lt} {v} = mt_loc_{local_name};").ok(); - if let Some(id) = vid { types.insert(id, lt); } + if let Some(id) = vid { + types.insert(id, lt); + } } else { // Other no-index loads (special const literals // like `-INFINITY`). Default to float. writeln!(out, "{pad}float {v} = float({});", safe_glsl_ident(src)).ok(); - if let Some(id) = vid { types.insert(id, "float"); } + if let Some(id) = vid { + types.insert(id, "float"); + } } } else { let idx = self.emit_idx(indices, block, ov, kernel, src)?; let arr = safe_glsl_ident(src); - let dtype = kernel - .params - .iter() - .find(|p| p.name == *src) - .map(|p| p.dtype); + let dtype = kernel.params.iter().find(|p| p.name == *src).map(|p| p.dtype); // Match the SSBO's natural type so bitwise ops on // u32 packs (`val >> 8 & 255`) don't go through a // float round-trip that loses bits. bf16 needs the @@ -920,24 +937,33 @@ impl GlslGenerator { // arithmetic stays in floating point. match dtype { Some(DType::BF16) => { - writeln!(out, "{pad}float {v} = mt_bf16_to_f32({arr}[uint({idx})]);").ok(); - if let Some(id) = vid { types.insert(id, "float"); } - } + writeln!(out, "{pad}float {v} = mt_bf16_to_f32({arr}[uint({idx})]);") + .ok(); + if let Some(id) = vid { + types.insert(id, "float"); + } + }, Some(DType::U32) | Some(DType::U16) | Some(DType::U8) => { writeln!(out, "{pad}uint {v} = uint({arr}[uint({idx})]);").ok(); - if let Some(id) = vid { types.insert(id, "uint"); } - } + if let Some(id) = vid { + types.insert(id, "uint"); + } + }, Some(DType::I32) | Some(DType::I8) | Some(DType::I4) | Some(DType::I64) => { writeln!(out, "{pad}int {v} = int({arr}[uint({idx})]);").ok(); - if let Some(id) = vid { types.insert(id, "int"); } - } + if let Some(id) = vid { + types.insert(id, "int"); + } + }, _ => { writeln!(out, "{pad}float {v} = float({arr}[uint({idx})]);").ok(); - if let Some(id) = vid { types.insert(id, "float"); } - } + if let Some(id) = vid { + types.insert(id, "float"); + } + }, } } - } + }, Op::Store { dst, indices, value, .. } => { let val = self.vname(Some(*value), block, ov); let idx = self.emit_idx(indices, block, ov, kernel, dst)?; @@ -950,11 +976,11 @@ impl GlslGenerator { Some(dt) => { let ty = Self::glsl_scalar_type(dt)?; format!("{ty}({val})") - } + }, None => format!("float({val})"), }; writeln!(out, "{pad}{arr}[uint({idx})] = {store_expr};").ok(); - } + }, Op::BinOp { op: bop, lhs, rhs } => { let v = self.vname(vid, block, ov); let l = self.vname(Some(*lhs), block, ov); @@ -1018,19 +1044,19 @@ impl GlslGenerator { if let Some(id) = vid { types.insert(id, ty); } - } + }, Op::Fma { a, b, c } => { let v = self.vname(vid, block, ov); let av = self.vname(Some(*a), block, ov); let bv = self.vname(Some(*b), block, ov); let cv = self.vname(Some(*c), block, ov); writeln!(out, "{pad}float {v} = fma({av}, {bv}, {cv});").ok(); - } + }, Op::UnaryOp { op: uop, value } => { let v = self.vname(vid, block, ov); let rv = self.vname(Some(*value), block, ov); writeln!(out, "{pad}float {v} = {};", self.glsl_unary(*uop, &rv)).ok(); - } + }, Op::Cast { value, dtype } => { let v = self.vname(vid, block, ov); let rv = self.vname(Some(*value), block, ov); @@ -1042,18 +1068,24 @@ impl GlslGenerator { match dtype { DType::U32 | DType::U16 | DType::U8 => { writeln!(out, "{pad}uint {v} = uint({rv});").ok(); - if let Some(id) = vid { types.insert(id, "uint"); } - } + if let Some(id) = vid { + types.insert(id, "uint"); + } + }, DType::I32 | DType::I8 | DType::I4 | DType::I64 => { writeln!(out, "{pad}int {v} = int({rv});").ok(); - if let Some(id) = vid { types.insert(id, "int"); } - } + if let Some(id) = vid { + types.insert(id, "int"); + } + }, _ => { writeln!(out, "{pad}float {v} = float({rv});").ok(); - if let Some(id) = vid { types.insert(id, "float"); } - } + if let Some(id) = vid { + types.insert(id, "float"); + } + }, } - } + }, Op::Select { cond, on_true, on_false } => { let v = self.vname(vid, block, ov); let c = self.vname(Some(*cond), block, ov); @@ -1071,8 +1103,10 @@ impl GlslGenerator { "float" }; writeln!(out, "{pad}{ty} {v} = ((bool({c})) ? ({a}) : ({b}));").ok(); - if let Some(id) = vid { types.insert(id, ty); } - } + if let Some(id) = vid { + types.insert(id, ty); + } + }, Op::DeclareLocal { name, value } => { let rv = self.vname(Some(*value), block, ov); // Type-aware declaration. `let mut rem = p` for uint @@ -1081,7 +1115,7 @@ impl GlslGenerator { let ty = types.get(value).copied().unwrap_or("float"); writeln!(out, "{pad}{ty} mt_loc_{name} = {rv};").ok(); local_types.insert(name.clone(), ty); - } + }, Op::SetLocal { name, value } => { let rv = self.vname(Some(*value), block, ov); let lt = local_types.get(name).copied().unwrap_or("float"); @@ -1091,7 +1125,7 @@ impl GlslGenerator { } else { writeln!(out, "{pad}mt_loc_{name} = {lt}({rv});").ok(); } - } + }, Op::Activation { kind, value } => { let v = self.vname(vid, block, ov); let rv = self.vname(Some(*value), block, ov); @@ -1103,12 +1137,20 @@ impl GlslGenerator { ActKind::Tanh => format!("tanh({rv})"), }; writeln!(out, "{pad}float {v} = {expr};").ok(); - } + }, // Per-thread grid-stride accumulation. For Reduction mode the // threadgroup cooperates (each thread strides by `lsize`); for // Grid3D each thread folds its own run (stride from the IR). Op::StrideReduce { - src, offset, stride, end, op: rk, transform, secondary_src, secondary_base, .. + src, + offset, + stride, + end, + op: rk, + transform, + secondary_src, + secondary_base, + .. } => { let v = self.vname(vid, block, ov); let off = self.vname(Some(*offset), block, ov); @@ -1140,18 +1182,15 @@ impl GlslGenerator { Some(sec) => { let bv = self.vname(*secondary_base, block, ov); let sec_arr = safe_glsl_ident(sec); - let sec_dtype = kernel - .params - .iter() - .find(|p| p.name == *sec) - .map(|p| p.dtype); + let sec_dtype = + kernel.params.iter().find(|p| p.name == *sec).map(|p| p.dtype); let load_sec = if matches!(sec_dtype, Some(DType::BF16)) { format!("mt_bf16_to_f32({sec_arr}[uint(_i) - uint({bv})])") } else { format!("float({sec_arr}[uint(_i) - uint({bv})])") }; format!("{} * {load_sec}", load_src("_i")) - } + }, None => load_src("_i"), }; let elem_expr = match transform.as_ref().map(|t| t.as_slice()) { @@ -1171,7 +1210,7 @@ impl GlslGenerator { Op::Cast { dtype, .. } => { let ty = Self::glsl_scalar_type(*dtype)?; format!("{ty}({e})") - } + }, Op::BinOp { op, rhs, .. } => { let rv = self.vname(Some(*rhs), block, ov); match op { @@ -1181,23 +1220,20 @@ impl GlslGenerator { BinOpKind::Div => format!("(({e}) / float({rv}))"), _ => e, } - } + }, _ => e, }; } e - } + }, }; writeln!(out, "{pad}float {v} = {};", reduce_init(*rk)).ok(); - writeln!( - out, - "{pad}for (uint _i = {start}; _i < uint({en}); _i += {step}) {{" - ) - .ok(); + writeln!(out, "{pad}for (uint _i = {start}; _i < uint({en}); _i += {step}) {{") + .ok(); writeln!(out, "{pad} float _e = {elem_expr};").ok(); writeln!(out, "{pad} {v} = {};", reduce_combine(*rk, &v, "_e")).ok(); writeln!(out, "{pad}}}").ok(); - } + }, // Workgroup-shared barrier-tree reduction — subgroup-width // agnostic (depends only on `local_size_x` + `barrier()`). // This is the portable path called out in @@ -1218,7 +1254,7 @@ impl GlslGenerator { return Err(Error::UnsupportedOp( "spirv: Reduce without a result value".into(), )); - } + }, }; let ls = self.local_size_total(); writeln!(out, "{pad}{red}[tid] = ({input});").ok(); @@ -1230,11 +1266,7 @@ impl GlslGenerator { // drops upper-partial lanes and loses ~1/3 of the reduction. // No-op for power-of-two ls (tid+_s 0u; _s >>= 1) {{" - ) - .ok(); + writeln!(out, "{pad}for (uint _s = {np}u / 2u; _s > 0u; _s >>= 1) {{").ok(); writeln!(out, "{pad} if (tid < _s && tid + _s < {ls}u) {{").ok(); let combine = reduce_combine(*rk, &format!("{red}[tid]"), &format!("{red}[tid + _s]")); @@ -1243,26 +1275,25 @@ impl GlslGenerator { writeln!(out, "{pad} barrier();").ok(); writeln!(out, "{pad}}}").ok(); if *rk == ReduceKind::Mean { - writeln!( - out, - "{pad}float {v} = {red}[0] / float({ls});" - ) - .ok(); + writeln!(out, "{pad}float {v} = {red}[0] / float({ls});").ok(); } else { writeln!(out, "{pad}float {v} = {red}[0];").ok(); } - } + }, // Threadgroup memory: declarations hoisted to file scope by // emit_shared_arrays; the ops themselves are loads/stores. - Op::ThreadgroupAlloc { .. } => {} // no-op (hoisted) + Op::ThreadgroupAlloc { .. } => {}, // no-op (hoisted) Op::ThreadgroupLoad { name, index } => { let v = self.vname(vid, block, ov); let iv = self.vname(Some(*index), block, ov); let n = safe_glsl_ident(name); - let (ty, expr) = dtype_load(tg_alloc_dtype(kernel, name), &format!("{n}[uint({iv})]")); + let (ty, expr) = + dtype_load(tg_alloc_dtype(kernel, name), &format!("{n}[uint({iv})]")); writeln!(out, "{pad}{ty} {v} = {expr};").ok(); - if let Some(id) = vid { types.insert(id, ty); } - } + if let Some(id) = vid { + types.insert(id, ty); + } + }, Op::ThreadgroupStore { name, index, value } => { let iv = self.vname(Some(*index), block, ov); let rv = self.vname(Some(*value), block, ov); @@ -1273,14 +1304,14 @@ impl GlslGenerator { let val = match dt { Some(DType::U32) | Some(DType::U16) | Some(DType::U8) => { format!("uint({rv})") - } + }, Some(DType::I32) | Some(DType::I8) | Some(DType::I4) | Some(DType::I64) => { format!("int({rv})") - } - _ => format!("{rv}"), + }, + _ => rv.to_string(), }; writeln!(out, "{pad}{n}[uint({iv})] = {val};").ok(); - } + }, // Per-thread (local) array — GLSL declares it in main. Op::StackAlloc { dtype, size, name } => { let ty = match dtype { @@ -1291,24 +1322,27 @@ impl GlslGenerator { }; let n = safe_glsl_ident(name); writeln!(out, "{pad}{ty} {n}[{size}];").ok(); - } + }, Op::StackLoad { name, index } => { let v = self.vname(vid, block, ov); let iv = self.vname(Some(*index), block, ov); let n = safe_glsl_ident(name); - let (ty, expr) = dtype_load(stack_alloc_dtype(kernel, name), &format!("{n}[uint({iv})]")); + let (ty, expr) = + dtype_load(stack_alloc_dtype(kernel, name), &format!("{n}[uint({iv})]")); writeln!(out, "{pad}{ty} {v} = {expr};").ok(); - if let Some(id) = vid { types.insert(id, ty); } - } + if let Some(id) = vid { + types.insert(id, ty); + } + }, Op::StackStore { name, index, value } => { let iv = self.vname(Some(*index), block, ov); let rv = self.vname(Some(*value), block, ov); let n = safe_glsl_ident(name); writeln!(out, "{pad}{n}[uint({iv})] = {rv};").ok(); - } + }, Op::Barrier => { writeln!(out, "{pad}barrier();").ok(); - } + }, // ── Subgroup (simdgroup / warp) primitives ──────────────── // Vulkan calls the warp-level group the "subgroup"; on RDNA // 4 wave32, subgroup size = 32 = profile.lane_width. We use @@ -1335,7 +1369,7 @@ impl GlslGenerator { // the sum across exactly 32 lanes. writeln!(out, "{pad}{v} = {v} / 32.0;").ok(); } - } + }, Op::SimdScan { value, op: rk, exclusive } => { let v = self.vname(vid, block, ov); let rv = self.vname(Some(*value), block, ov); @@ -1347,35 +1381,39 @@ impl GlslGenerator { ReduceKind::Product => format!("subgroup{prefix}Mul"), }; writeln!(out, "{pad}float {v} = {intrin}({rv});").ok(); - } + }, Op::SimdLaneId => { let v = self.vname(vid, block, ov); // Use the preamble's `simd_lane` (`tid % 32`) so this // matches the simdgroup-matrix lane→element mapping. writeln!(out, "{pad}uint {v} = simd_lane;").ok(); - if let Some(id) = vid { types.insert(id, "uint"); } - } + if let Some(id) = vid { + types.insert(id, "uint"); + } + }, Op::SimdGroupId => { let v = self.vname(vid, block, ov); // Same — `simd_group` from the preamble (`tid / 32`) // matches the shared simdgroup-matrix tile offset. writeln!(out, "{pad}uint {v} = simd_group;").ok(); - if let Some(id) = vid { types.insert(id, "uint"); } - } + if let Some(id) = vid { + types.insert(id, "uint"); + } + }, Op::SimdBroadcast { value, lane } => { let v = self.vname(vid, block, ov); let rv = self.vname(Some(*value), block, ov); let rl = self.vname(Some(*lane), block, ov); writeln!(out, "{pad}float {v} = subgroupBroadcast({rv}, uint({rl}));").ok(); - } + }, Op::SimdShuffleXor { value, mask } => { let v = self.vname(vid, block, ov); let rv = self.vname(Some(*value), block, ov); writeln!(out, "{pad}float {v} = subgroupShuffleXor({rv}, {mask}u);").ok(); - } + }, Op::SimdgroupBarrier => { writeln!(out, "{pad}subgroupMemoryBarrierShared(); subgroupBarrier();").ok(); - } + }, // Atomic ops on storage buffers — GLSL has direct intrinsics. // `Op::Atomic` is the IR's portable atomic-RMW; for f32 sums // we need `atomicAdd` on `float` which requires @@ -1407,10 +1445,10 @@ impl GlslGenerator { let val_cast = match dt { DType::U32 | DType::U16 | DType::U8 => format!("uint({rv})"), DType::I32 | DType::I8 | DType::I4 | DType::I64 => format!("int({rv})"), - _ => format!("{rv}"), + _ => rv.to_string(), }; writeln!(out, "{pad}{f}({arr}[uint({iv})], {val_cast});").ok(); - } + }, // ── simdgroup_matrix software emulation ───────────── // Apple's per-warp 8×8 float fragment. The CUDA emitter // already does the software emulation (each warp owns a @@ -1420,27 +1458,19 @@ impl GlslGenerator { // // The shared declaration is hoisted by `emit_shared_arrays` // (added below); per-op handlers only emit the access. - Op::SimdgroupAlloc { .. } => {} // declared at file scope + Op::SimdgroupAlloc { .. } => {}, // declared at file scope Op::SimdgroupElemLoad { value, index } => { let v = self.vname(vid, block, ov); let m = sgm_name(*value); let fnk = if *index == 0 { "_sg_fn0" } else { "_sg_fn1" }; - writeln!( - out, - "{pad}float {v} = {m}[simd_group * 64u + _sg_fm * 8u + {fnk}];" - ) - .ok(); - } + writeln!(out, "{pad}float {v} = {m}[simd_group * 64u + _sg_fm * 8u + {fnk}];").ok(); + }, Op::SimdgroupElemStore { value, index, data } => { let m = sgm_name(*value); let dv = self.vname(Some(*data), block, ov); let fnk = if *index == 0 { "_sg_fn0" } else { "_sg_fn1" }; - writeln!( - out, - "{pad}{m}[simd_group * 64u + _sg_fm * 8u + {fnk}] = {dv};" - ) - .ok(); - } + writeln!(out, "{pad}{m}[simd_group * 64u + _sg_fm * 8u + {fnk}] = {dv};").ok(); + }, Op::SimdgroupLoad { dest, tg, offset, stride, transpose } => { let m = sgm_name(*dest); let off = self.vname(Some(*offset), block, ov); @@ -1458,12 +1488,12 @@ impl GlslGenerator { .ok(); } writeln!(out, "{pad}}}").ok(); - } + }, // ── CoopTile cooperative GEMM (software emulation) ────────── // Mirrors the CUDA emitter's `mpp::matmul2d` software path, // adapted to GLSL. Per-warp A/B/C shared tiles are hoisted // by emit_shared_arrays; ops here are loads / runs / stores. - Op::CoopTileSetup { .. } => {} // declared at file scope + Op::CoopTileSetup { .. } => {}, // declared at file scope Op::CoopTileZero { name } => { let (m, n, _, _, _, _, _, simd) = self.coop_cfg(kernel, name).ok_or_else(|| { Error::UnsupportedOp(format!("spirv: CoopTile `{name}` no Setup")) @@ -1478,10 +1508,8 @@ impl GlslGenerator { writeln!(out, "{pad} _CMAcc_{cnm}[_mi][_ni] = coopmat(0.0);").ok(); return Ok(()); } - let local_c = matches!( - self.profile.mma, - crate::backend::MmaStrategy::SoftwareLocalC - ); + let local_c = + matches!(self.profile.mma, crate::backend::MmaStrategy::SoftwareLocalC); if local_c && simd { let lw = self.profile.lane_width; let per_lane = (m * n).div_ceil(lw).max(1); @@ -1500,7 +1528,7 @@ impl GlslGenerator { ) .ok(); } - } + }, Op::CoopTileLoadA { name, ptr_name, ptr_offset, dtype: _, ei, .. } => { let (m, _, k, ta, _, _, _, simd) = self.coop_cfg(kernel, name).ok_or_else(|| { @@ -1526,7 +1554,7 @@ impl GlslGenerator { m * k ) .ok(); - } + }, Op::CoopTileLoadB { name, ptr_name, ptr_offset, dtype: _, ei, .. } => { let (_, n, k, _, tb, _, _, simd) = self.coop_cfg(kernel, name).ok_or_else(|| { @@ -1551,7 +1579,7 @@ impl GlslGenerator { k * n ) .ok(); - } + }, Op::CoopTileRun { name, .. } => { let (m, n, k, _, _, _, accum, simd) = self.coop_cfg(kernel, name).ok_or_else(|| { @@ -1574,11 +1602,17 @@ impl GlslGenerator { let ba = coop_base_glsl(true, m * k); let bb = coop_base_glsl(true, k * n); writeln!(out, "{pad}barrier();").ok(); - writeln!(out, "{pad}[[unroll]] for (uint _kt = 0u; _kt < {kf}u; _kt++) {{").ok(); - writeln!(out, "{pad} [[unroll]] for (uint _mi = 0u; _mi < {mf}u; _mi++) {{").ok(); + writeln!(out, "{pad}[[unroll]] for (uint _kt = 0u; _kt < {kf}u; _kt++) {{") + .ok(); + writeln!(out, "{pad} [[unroll]] for (uint _mi = 0u; _mi < {mf}u; _mi++) {{") + .ok(); writeln!(out, "{pad} coopmat _ca;").ok(); writeln!(out, "{pad} coopMatLoad(_ca, _CTA_{nm}, {ba} + (_mi * 16u) * {k}u + _kt * 16u, {k}u, gl_CooperativeMatrixLayoutRowMajor);").ok(); - writeln!(out, "{pad} [[unroll]] for (uint _ni = 0u; _ni < {nf}u; _ni++) {{").ok(); + writeln!( + out, + "{pad} [[unroll]] for (uint _ni = 0u; _ni < {nf}u; _ni++) {{" + ) + .ok(); writeln!(out, "{pad} coopmat _cb;").ok(); writeln!(out, "{pad} coopMatLoad(_cb, _CTB_{nm}, {bb} + (_kt * 16u) * {n}u + _ni * 16u, {n}u, gl_CooperativeMatrixLayoutRowMajor);").ok(); writeln!(out, "{pad} _CMAcc_{cnm}[_mi][_ni] = coopMatMulAdd(_ca, _cb, _CMAcc_{cnm}[_mi][_ni]);").ok(); @@ -1588,29 +1622,20 @@ impl GlslGenerator { writeln!(out, "{pad}barrier();").ok(); return Ok(()); } - let local_c = matches!( - self.profile.mma, - crate::backend::MmaStrategy::SoftwareLocalC - ); + let local_c = + matches!(self.profile.mma, crate::backend::MmaStrategy::SoftwareLocalC); if local_c && simd { let lw = self.profile.lane_width; let per_lane = (m * n).div_ceil(lw).max(1); let ba = coop_base_glsl(true, m * k); let bb = coop_base_glsl(true, k * n); writeln!(out, "{pad}barrier();").ok(); - writeln!( - out, - "{pad}for (uint _li = 0u; _li < {per_lane}u; _li++) {{" - ) - .ok(); + writeln!(out, "{pad}for (uint _li = 0u; _li < {per_lane}u; _li++) {{").ok(); writeln!(out, "{pad} uint _e = simd_lane + _li * {lw}u;").ok(); writeln!(out, "{pad} if (_e >= {}u) break;", m * n).ok(); writeln!(out, "{pad} uint _i = _e / {n}u, _j = _e % {n}u;").ok(); - let init = if accum { - format!("_CTC_lane_{cnm}[_li]") - } else { - "0.0".to_string() - }; + let init = + if accum { format!("_CTC_lane_{cnm}[_li]") } else { "0.0".to_string() }; writeln!(out, "{pad} float _acc = {init};").ok(); writeln!(out, "{pad} for (uint _l = 0u; _l < {k}u; _l++) _acc += _CTA_{nm}[{ba} + _i * {k}u + _l] * _CTB_{nm}[{bb} + _l * {n}u + _j];").ok(); writeln!(out, "{pad} _CTC_lane_{cnm}[_li] = _acc;").ok(); @@ -1624,12 +1649,8 @@ impl GlslGenerator { coop_base_glsl(simd, m * n), ); writeln!(out, "{pad}barrier();").ok(); - writeln!( - out, - "{pad}for (uint _e = {sid}; _e < {}u; _e += {ssize}) {{", - m * n - ) - .ok(); + writeln!(out, "{pad}for (uint _e = {sid}; _e < {}u; _e += {ssize}) {{", m * n) + .ok(); writeln!(out, "{pad} uint _i = _e / {n}u, _j = _e % {n}u;").ok(); let init = if accum { format!("_CTC_{cnm}[{bc} + _e]") } else { "0.0".into() }; writeln!(out, "{pad} float _acc = {init};").ok(); @@ -1638,12 +1659,11 @@ impl GlslGenerator { writeln!(out, "{pad}}}").ok(); writeln!(out, "{pad}barrier();").ok(); } - } + }, Op::CoopTileStoreC { name, ptr_name, ptr_offset, dtype: _, ei, .. } => { - let (m, n, _, _, _, _, _, simd) = - self.coop_cfg(kernel, name).ok_or_else(|| { - Error::UnsupportedOp(format!("spirv: CoopTile `{name}` no Setup")) - })?; + let (m, n, _, _, _, _, _, simd) = self.coop_cfg(kernel, name).ok_or_else(|| { + Error::UnsupportedOp(format!("spirv: CoopTile `{name}` no Setup")) + })?; let cnm = ct_ident(&coop_c_name(name)); let off = ptr_offset .map(|o| self.vname(Some(o), block, ov)) @@ -1661,10 +1681,8 @@ impl GlslGenerator { writeln!(out, "{pad} coopMatStore(_CMAcc_{cnm}[_mi][_ni], {arr}, uint({off}) + (_mi * 16u) * {ei}u + _ni * 16u, {ei}u, gl_CooperativeMatrixLayoutRowMajor);").ok(); return Ok(()); } - let local_c = matches!( - self.profile.mma, - crate::backend::MmaStrategy::SoftwareLocalC - ); + let local_c = + matches!(self.profile.mma, crate::backend::MmaStrategy::SoftwareLocalC); if local_c && simd { let lw = self.profile.lane_width; let per_lane = (m * n).div_ceil(lw).max(1); @@ -1672,7 +1690,11 @@ impl GlslGenerator { writeln!(out, "{pad}for (uint _li = 0u; _li < {per_lane}u; _li++) {{").ok(); writeln!(out, "{pad} uint _e = simd_lane + _li * {lw}u;").ok(); writeln!(out, "{pad} if (_e >= {}u) break;", m * n).ok(); - writeln!(out, "{pad} {arr}[uint({off}) + {dst}] = float(_CTC_lane_{cnm}[_li]);").ok(); + writeln!( + out, + "{pad} {arr}[uint({off}) + {dst}] = float(_CTC_lane_{cnm}[_li]);" + ) + .ok(); writeln!(out, "{pad}}}").ok(); } else { let (sid, ssize, _) = coop_scope_glsl(simd, self.profile.lane_width); @@ -1685,7 +1707,7 @@ impl GlslGenerator { ) .ok(); } - } + }, Op::SimdgroupMatMul { a, b, c } => { let (ma, mb, mc) = (sgm_name(*a), sgm_name(*b), sgm_name(*c)); writeln!(out, "{pad}subgroupMemoryBarrierShared(); subgroupBarrier();").ok(); @@ -1701,7 +1723,7 @@ impl GlslGenerator { writeln!(out, "{pad} {mc}[_bs + _sg_fm * 8u + _sg_fn0] = _acc0;").ok(); writeln!(out, "{pad} {mc}[_bs + _sg_fm * 8u + _sg_fn1] = _acc1;").ok(); writeln!(out, "{pad}}}").ok(); - } + }, // Control flow — nested-block recursion mirroring the CUDA walker. Op::Loop { var, start, end, step, body } => { let s = self.vname(Some(*start), block, ov); @@ -1731,7 +1753,7 @@ impl GlslGenerator { self.emit_ops(bb, kernel, &child, types, local_types, out)?; } writeln!(out, "{pad}}}").ok(); - } + }, Op::If { cond, then_block, else_block } => { let c = self.vname(Some(*cond), block, ov); writeln!(out, "{pad}if (bool({c})) {{").ok(); @@ -1747,22 +1769,22 @@ impl GlslGenerator { } } writeln!(out, "{pad}}}").ok(); - } + }, // Scalar zero / splat. Op::Zeros { shape, .. } if shape.rank() == 0 => { let v = self.vname(vid, block, ov); writeln!(out, "{pad}float {v} = 0.0;").ok(); - } + }, Op::Splat { value, shape, .. } if shape.rank() == 0 => { let v = self.vname(vid, block, ov); writeln!(out, "{pad}float {v} = float({value});").ok(); - } + }, other => { return Err(Error::UnsupportedOp(format!( "spirv: op {} not yet supported", op_variant_name(other) ))); - } + }, } Ok(()) } @@ -1780,30 +1802,20 @@ impl GlslGenerator { for blk in std::iter::once(&kernel.body).chain(kernel.blocks.values()) { for op in &blk.ops { if let Op::CoopTileSetup { - name: nm, - m, - n, - k, - ta, - tb, - tc, - acc_mode, - exec_scope, - .. + name: nm, m, n, k, ta, tb, tc, acc_mode, exec_scope, .. } = op + && nm == name { - if nm == name { - return Some(( - *m, - *n, - *k, - *ta, - *tb, - *tc, - matches!(acc_mode, CoopTileAccMode::MultiplyAccumulate), - matches!(exec_scope, CoopTileScope::SimdGroup), - )); - } + return Some(( + *m, + *n, + *k, + *ta, + *tb, + *tc, + matches!(acc_mode, CoopTileAccMode::MultiplyAccumulate), + matches!(exec_scope, CoopTileScope::SimdGroup), + )); } } } @@ -1920,12 +1932,10 @@ fn stack_alloc_dtype(kernel: &Kernel, name: &str) -> Option { /// by ThreadgroupLoad / StackLoad. Returns `(glsl_type, expr_around_arr)`. fn dtype_load(dtype: Option, arr_idx: &str) -> (&'static str, String) { match dtype { - Some(DType::U32) | Some(DType::U16) | Some(DType::U8) => { - ("uint", format!("uint({arr_idx})")) - } - Some(DType::I32) | Some(DType::I8) | Some(DType::I4) | Some(DType::I64) => { - ("int", format!("int({arr_idx})")) - } + Some(DType::U32) | Some(DType::U16) | Some(DType::U8) => + ("uint", format!("uint({arr_idx})")), + Some(DType::I32) | Some(DType::I8) | Some(DType::I4) | Some(DType::I64) => + ("int", format!("int({arr_idx})")), _ => ("float", format!("float({arr_idx})")), } } @@ -1938,14 +1948,10 @@ fn sgm_name(v: ValueId) -> String { format!("_SGM_{}", v.as_u32()) } /// CoopTile name sanitization, matches the CUDA emitter so the shared /// arrays line up by spelling. fn ct_ident(name: &str) -> String { - name.chars() - .map(|c| if c.is_alphanumeric() || c == '_' { c } else { '_' }) - .collect() + name.chars().map(|c| if c.is_alphanumeric() || c == '_' { c } else { '_' }).collect() } -fn coop_c_name(name: &str) -> String { - name.strip_suffix("_acc").unwrap_or(name).to_string() -} +fn coop_c_name(name: &str) -> String { name.strip_suffix("_acc").unwrap_or(name).to_string() } /// `(thread-id, scope-size, barrier)` for a cooperative iteration in /// GLSL — same shape as the CUDA helper but with `subgroupBarrier()`. @@ -1958,11 +1964,7 @@ fn coop_scope_glsl(simd: bool, lane_width: u32) -> (&'static str, String, &'stat } fn coop_base_glsl(simd: bool, tile: u32) -> String { - if simd { - format!("simd_group * {tile}u") - } else { - "0u".to_string() - } + if simd { format!("simd_group * {tile}u") } else { "0u".to_string() } } /// Does the kernel use any simdgroup-matrix op (drives the lane-coord @@ -2005,39 +2007,126 @@ pub fn safe_glsl_ident(name: &str) -> String { // likely to hit. Not exhaustive — the compiler will catch anything we // miss, and we can grow this on demand. const RESERVED: &[&str] = &[ - "in", "out", "inout", "uniform", "buffer", "shared", - "const", "void", "bool", "int", "uint", "float", "double", - "vec2", "vec3", "vec4", "mat2", "mat3", "mat4", - "true", "false", "if", "else", "for", "while", "do", "switch", "case", - "return", "break", "continue", "discard", - "layout", "precision", "highp", "mediump", "lowp", - "attribute", "varying", - "centroid", "flat", "smooth", "noperspective", - "coherent", "volatile", "restrict", "readonly", "writeonly", - "sampler1D", "sampler2D", "sampler3D", "samplerCube", - "image1D", "image2D", "image3D", + "in", + "out", + "inout", + "uniform", + "buffer", + "shared", + "const", + "void", + "bool", + "int", + "uint", + "float", + "double", + "vec2", + "vec3", + "vec4", + "mat2", + "mat3", + "mat4", + "true", + "false", + "if", + "else", + "for", + "while", + "do", + "switch", + "case", + "return", + "break", + "continue", + "discard", + "layout", + "precision", + "highp", + "mediump", + "lowp", + "attribute", + "varying", + "centroid", + "flat", + "smooth", + "noperspective", + "coherent", + "volatile", + "restrict", + "readonly", + "writeonly", + "sampler1D", + "sampler2D", + "sampler3D", + "samplerCube", + "image1D", + "image2D", + "image3D", // GLSL reserves these for future use, even though they aren't // currently used; shaderc rejects them. The corpus hit them on // conv* / depthwise_conv2d / aura_value_int4 / ffai_gemm kernels. - "input", "output", "texture", "image", "sampler", - "active", "partition", "common", "filter", - "row_major", "column_major", "packed", - "asm", "class", "template", "this", "namespace", "interface", - "public", "static", "extern", "external", - "long", "short", "half", "fixed", "unsigned", - "hvec2", "hvec3", "hvec4", "fvec2", "fvec3", "fvec4", - "dvec2", "dvec3", "dvec4", "ivec2", "ivec3", "ivec4", - "uvec2", "uvec3", "uvec4", "bvec2", "bvec3", "bvec4", - "snorm", "unorm", "snorm8", "unorm8", + "input", + "output", + "texture", + "image", + "sampler", + "active", + "partition", + "common", + "filter", + "row_major", + "column_major", + "packed", + "asm", + "class", + "template", + "this", + "namespace", + "interface", + "public", + "static", + "extern", + "external", + "long", + "short", + "half", + "fixed", + "unsigned", + "hvec2", + "hvec3", + "hvec4", + "fvec2", + "fvec3", + "fvec4", + "dvec2", + "dvec3", + "dvec4", + "ivec2", + "ivec3", + "ivec4", + "uvec2", + "uvec3", + "uvec4", + "bvec2", + "bvec3", + "bvec4", + "snorm", + "unorm", + "snorm8", + "unorm8", // GLSL built-in functions / common names that surface as params: - "length", "distance", "dot", "cross", "normalize", "reflect", - "refract", "transpose", "determinant", "inverse", + "length", + "distance", + "dot", + "cross", + "normalize", + "reflect", + "refract", + "transpose", + "determinant", + "inverse", ]; - if RESERVED.contains(&name) { - format!("_b_{name}") - } else { - name.to_string() - } + if RESERVED.contains(&name) { format!("_b_{name}") } else { name.to_string() } } impl CodegenBackend for GlslGenerator { @@ -2049,10 +2138,9 @@ impl CodegenBackend for GlslGenerator { // Run the same backend-neutral kernel-inline pass the CUDA emitter // does, so cross-kernel calls resolve in this codegen too. let mut inlined = kernel.clone(); - let k: &Kernel = match crate::passes::run_passes( - &mut inlined, - &[Box::new(crate::passes::kernel_inline::KernelInlinePass)], - ) { + let k: &Kernel = match crate::passes::run_passes(&mut inlined, &[Box::new( + crate::passes::kernel_inline::KernelInlinePass, + )]) { Ok(()) => &inlined, Err(_) => kernel, }; @@ -2142,8 +2230,16 @@ mod tests { constexpr::ConstExpr, dtype::DType, ir::{ - BinOpKind, ConstExprDecl, IndexExpr, Kernel, KernelMode, Op, Param, ParamKind, - ReduceKind, ValueId, + BinOpKind, + ConstExprDecl, + IndexExpr, + Kernel, + KernelMode, + Op, + Param, + ParamKind, + ReduceKind, + ValueId, }, shape::Shape, }; @@ -2164,12 +2260,22 @@ mod tests { k.body.push_op(Op::ProgramId { axis: 0 }, ValueId::new(0)); k.body.name_value(ValueId::new(0), "idx"); k.body.push_op( - Op::Load { src: "a".into(), indices: vec![IndexExpr::Value(ValueId::new(0))], mask: None, other: None }, + Op::Load { + src: "a".into(), + indices: vec![IndexExpr::Value(ValueId::new(0))], + mask: None, + other: None, + }, ValueId::new(1), ); k.body.name_value(ValueId::new(1), "x"); k.body.push_op( - Op::Load { src: "b".into(), indices: vec![IndexExpr::Value(ValueId::new(0))], mask: None, other: None }, + Op::Load { + src: "b".into(), + indices: vec![IndexExpr::Value(ValueId::new(0))], + mask: None, + other: None, + }, ValueId::new(2), ); k.body.name_value(ValueId::new(2), "y"); @@ -2191,12 +2297,18 @@ mod tests { let mut k = Kernel::new("row_reduce_sum"); k.mode = KernelMode::Reduction; k.params.push(Param { - name: "inp".into(), dtype: DType::F32, shape: Shape::scalar(), - is_output: false, kind: ParamKind::Tensor, + name: "inp".into(), + dtype: DType::F32, + shape: Shape::scalar(), + is_output: false, + kind: ParamKind::Tensor, }); k.params.push(Param { - name: "out".into(), dtype: DType::F32, shape: Shape::scalar(), - is_output: true, kind: ParamKind::Tensor, + name: "out".into(), + dtype: DType::F32, + shape: Shape::scalar(), + is_output: true, + kind: ParamKind::Tensor, }); k.constexprs.push(ConstExprDecl { name: ConstExpr::new("n"), @@ -2204,8 +2316,12 @@ mod tests { value: None, }); let (row, nv, rs, re, acc, res) = ( - ValueId::new(0), ValueId::new(1), ValueId::new(2), - ValueId::new(3), ValueId::new(4), ValueId::new(5), + ValueId::new(0), + ValueId::new(1), + ValueId::new(2), + ValueId::new(3), + ValueId::new(4), + ValueId::new(5), ); k.body.push_op(Op::ProgramId { axis: 0 }, row); k.body.name_value(row, "row"); @@ -2216,9 +2332,15 @@ mod tests { k.body.name_value(re, "re"); k.body.push_op( Op::StrideReduce { - src: "inp".into(), offset: rs, stride: nv, end: re, - op: ReduceKind::Sum, dtype: DType::F32, - transform: None, secondary_src: None, secondary_base: None, + src: "inp".into(), + offset: rs, + stride: nv, + end: re, + op: ReduceKind::Sum, + dtype: DType::F32, + transform: None, + secondary_src: None, + secondary_base: None, }, acc, ); @@ -2226,7 +2348,10 @@ mod tests { k.body.push_op(Op::Reduce { value: acc, axis: 0, op: ReduceKind::Sum }, res); k.body.name_value(res, "result"); k.body.push_op_no_result(Op::Store { - dst: "out".into(), indices: vec![IndexExpr::Value(row)], value: res, mask: None, + dst: "out".into(), + indices: vec![IndexExpr::Value(row)], + value: res, + mask: None, }); k } diff --git a/crates/metaltile-runtime/build.rs b/crates/metaltile-runtime/build.rs index bc442ed4..ad47af9f 100644 --- a/crates/metaltile-runtime/build.rs +++ b/crates/metaltile-runtime/build.rs @@ -34,11 +34,7 @@ fn cuda() { for sub in ["lib64", "lib64/stubs", "lib", "lib/stubs", "lib/x64"] { println!("cargo:rustc-link-search=native={cuda_root}/{sub}"); } - for p in [ - "/usr/lib/aarch64-linux-gnu", - "/usr/lib/x86_64-linux-gnu", - "/usr/lib64", - ] { + for p in ["/usr/lib/aarch64-linux-gnu", "/usr/lib/x86_64-linux-gnu", "/usr/lib64"] { println!("cargo:rustc-link-search=native={p}"); } @@ -62,14 +58,28 @@ fn cuda() { println!("cargo:rerun-if-changed={src}"); let obj = format!("{out_dir}/cutlass_moe.o"); let status = std::process::Command::new(&nvcc) - .args(["-O3", "-std=c++17", &format!("-arch={arch}"), "--expt-relaxed-constexpr", "-Xcompiler", "-fPIC"]) - .args(["-I", &format!("{cutlass_dir}/include"), "-I", &format!("{cutlass_dir}/tools/util/include")]) + .args([ + "-O3", + "-std=c++17", + &format!("-arch={arch}"), + "--expt-relaxed-constexpr", + "-Xcompiler", + "-fPIC", + ]) + .args([ + "-I", + &format!("{cutlass_dir}/include"), + "-I", + &format!("{cutlass_dir}/tools/util/include"), + ]) .args(["-c", &src, "-o", &obj]) .status() .expect("nvcc invocation for cutlass_moe.cu failed to start"); assert!(status.success(), "nvcc failed to compile cutlass_moe.cu"); let lib = format!("{out_dir}/libcutlass_moe.a"); - let ar = std::process::Command::new("ar").args(["crus", &lib, &obj]).status() + let ar = std::process::Command::new("ar") + .args(["crus", &lib, &obj]) + .status() .expect("ar failed to start"); assert!(ar.success(), "ar failed to archive cutlass_moe.o"); println!("cargo:rustc-link-search=native={out_dir}"); @@ -85,9 +95,8 @@ fn cuda() { /// `/opt/rocm/lib`. `HIP_PATH` is the canonical env var; `ROCM_PATH` /// covers Linux installs that prefer that name. fn hip() { - let hip_root = std::env::var("HIP_PATH") - .or_else(|_| std::env::var("ROCM_PATH")) - .unwrap_or_else(|_| { + let hip_root = + std::env::var("HIP_PATH").or_else(|_| std::env::var("ROCM_PATH")).unwrap_or_else(|_| { if cfg!(windows) { r"C:\Program Files\AMD\ROCm\7.1".to_string() } else { @@ -109,11 +118,7 @@ fn hip() { /// under `/Lib` on Windows, `/lib` on Linux. fn vulkan() { let vk_sdk = std::env::var("VULKAN_SDK").unwrap_or_else(|_| { - if cfg!(windows) { - r"C:\VulkanSDK\1.4.350.0".to_string() - } else { - "/usr".to_string() - } + if cfg!(windows) { r"C:\VulkanSDK\1.4.350.0".to_string() } else { "/usr".to_string() } }); for sub in ["Lib", "lib", "lib64"] { diff --git a/crates/metaltile-runtime/src/device/cuda/ffi.rs b/crates/metaltile-runtime/src/device/cuda/ffi.rs index 9b0e025e..e163f1fe 100644 --- a/crates/metaltile-runtime/src/device/cuda/ffi.rs +++ b/crates/metaltile-runtime/src/device/cuda/ffi.rs @@ -63,7 +63,12 @@ unsafe extern "C" { // cuMemcpyHtoD is always synchronous; pinned + Async is not). pub fn cuMemAllocHost_v2(pp: *mut *mut c_void, bytesize: usize) -> CUresult; pub fn cuMemFreeHost(p: *mut c_void) -> CUresult; - pub fn cuMemcpyHtoDAsync_v2(dst: CUdeviceptr, src: *const c_void, byte_count: usize, stream: CUstream) -> CUresult; + pub fn cuMemcpyHtoDAsync_v2( + dst: CUdeviceptr, + src: *const c_void, + byte_count: usize, + stream: CUstream, + ) -> CUresult; #[allow(clippy::too_many_arguments)] pub fn cuLaunchKernel( f: CUfunction, @@ -100,11 +105,7 @@ unsafe extern "C" { pub fn cuEventCreate(phEvent: *mut CUevent, flags: c_uint) -> CUresult; pub fn cuEventRecord(hEvent: CUevent, hStream: CUstream) -> CUresult; pub fn cuEventSynchronize(hEvent: CUevent) -> CUresult; - pub fn cuEventElapsedTime( - pMilliseconds: *mut f32, - hStart: CUevent, - hEnd: CUevent, - ) -> CUresult; + pub fn cuEventElapsedTime(pMilliseconds: *mut f32, hStart: CUevent, hEnd: CUevent) -> CUresult; pub fn cuEventDestroy_v2(hEvent: CUevent) -> CUresult; // ── Stream + CUDA-graph capture (replay a whole decode token as ONE graph // launch, eliminating the ~390 per-kernel launch/host-orchestration costs). ── @@ -182,7 +183,11 @@ unsafe extern "C" { pub fn cublasSetStream_v2(handle: cublasHandle_t, stream: CUstream) -> cublasStatus_t; /// Persistent workspace so cuBLAS GEMMs are CUDA-graph capture/replay-safe (the /// default workspace is a transient per-call alloc → replay reads freed memory). - pub fn cublasSetWorkspace_v2(handle: cublasHandle_t, workspace: *mut c_void, workspace_bytes: usize) -> cublasStatus_t; + pub fn cublasSetWorkspace_v2( + handle: cublasHandle_t, + workspace: *mut c_void, + workspace_bytes: usize, + ) -> cublasStatus_t; /// Allow/forbid atomic-accumulation kernels (split-K). Forbidding → deterministic. pub fn cublasSetAtomicsMode(handle: cublasHandle_t, mode: c_int) -> cublasStatus_t; pub fn cublasSetMathMode(handle: cublasHandle_t, mode: c_int) -> cublasStatus_t; @@ -325,7 +330,7 @@ pub const CUBLASLT_REDUCTION_SCHEME_NONE: u32 = 0; #[repr(C)] #[derive(Clone, Copy)] pub struct cublasLtMatmulHeuristicResult_t { - pub algo: [u64; 8], // cublasLtMatmulAlgo_t + pub algo: [u64; 8], // cublasLtMatmulAlgo_t pub workspace_size: usize, pub state: cublasStatus_t, pub waves_count: f32, @@ -334,7 +339,11 @@ pub struct cublasLtMatmulHeuristicResult_t { impl Default for cublasLtMatmulHeuristicResult_t { fn default() -> Self { cublasLtMatmulHeuristicResult_t { - algo: [0; 8], workspace_size: 0, state: 0, waves_count: 0.0, reserved: [0; 4], + algo: [0; 8], + workspace_size: 0, + state: 0, + waves_count: 0.0, + reserved: [0; 4], } } } @@ -343,15 +352,40 @@ impl Default for cublasLtMatmulHeuristicResult_t { unsafe extern "C" { pub fn cublasLtCreate(handle: *mut cublasLtHandle_t) -> cublasStatus_t; pub fn cublasLtDestroy(handle: cublasLtHandle_t) -> cublasStatus_t; - pub fn cublasLtMatmulDescCreate(desc: *mut cublasLtMatmulDesc_t, compute_type: c_int, scale_type: c_int) -> cublasStatus_t; + pub fn cublasLtMatmulDescCreate( + desc: *mut cublasLtMatmulDesc_t, + compute_type: c_int, + scale_type: c_int, + ) -> cublasStatus_t; pub fn cublasLtMatmulDescDestroy(desc: cublasLtMatmulDesc_t) -> cublasStatus_t; - pub fn cublasLtMatmulDescSetAttribute(desc: cublasLtMatmulDesc_t, attr: c_int, buf: *const c_void, size: usize) -> cublasStatus_t; - pub fn cublasLtMatrixLayoutCreate(layout: *mut cublasLtMatrixLayout_t, dtype: c_int, rows: u64, cols: u64, ld: i64) -> cublasStatus_t; - pub fn cublasLtMatrixLayoutSetAttribute(layout: cublasLtMatrixLayout_t, attr: c_int, buf: *const c_void, size: usize) -> cublasStatus_t; + pub fn cublasLtMatmulDescSetAttribute( + desc: cublasLtMatmulDesc_t, + attr: c_int, + buf: *const c_void, + size: usize, + ) -> cublasStatus_t; + pub fn cublasLtMatrixLayoutCreate( + layout: *mut cublasLtMatrixLayout_t, + dtype: c_int, + rows: u64, + cols: u64, + ld: i64, + ) -> cublasStatus_t; + pub fn cublasLtMatrixLayoutSetAttribute( + layout: cublasLtMatrixLayout_t, + attr: c_int, + buf: *const c_void, + size: usize, + ) -> cublasStatus_t; pub fn cublasLtMatrixLayoutDestroy(layout: cublasLtMatrixLayout_t) -> cublasStatus_t; pub fn cublasLtMatmulPreferenceCreate(pref: *mut cublasLtMatmulPreference_t) -> cublasStatus_t; pub fn cublasLtMatmulPreferenceDestroy(pref: cublasLtMatmulPreference_t) -> cublasStatus_t; - pub fn cublasLtMatmulPreferenceSetAttribute(pref: cublasLtMatmulPreference_t, attr: c_int, buf: *const c_void, size: usize) -> cublasStatus_t; + pub fn cublasLtMatmulPreferenceSetAttribute( + pref: cublasLtMatmulPreference_t, + attr: c_int, + buf: *const c_void, + size: usize, + ) -> cublasStatus_t; #[allow(clippy::too_many_arguments)] pub fn cublasLtMatmulAlgoGetHeuristic( handle: cublasLtHandle_t, @@ -370,11 +404,15 @@ unsafe extern "C" { handle: cublasLtHandle_t, compute_desc: cublasLtMatmulDesc_t, alpha: *const c_void, - a: CUdeviceptr, a_desc: cublasLtMatrixLayout_t, - b: CUdeviceptr, b_desc: cublasLtMatrixLayout_t, + a: CUdeviceptr, + a_desc: cublasLtMatrixLayout_t, + b: CUdeviceptr, + b_desc: cublasLtMatrixLayout_t, beta: *const c_void, - c: CUdeviceptr, c_desc: cublasLtMatrixLayout_t, - d: CUdeviceptr, d_desc: cublasLtMatrixLayout_t, + c: CUdeviceptr, + c_desc: cublasLtMatrixLayout_t, + d: CUdeviceptr, + d_desc: cublasLtMatrixLayout_t, algo: *const u64, workspace: CUdeviceptr, workspace_size: usize, diff --git a/crates/metaltile-runtime/src/device/cuda/mod.rs b/crates/metaltile-runtime/src/device/cuda/mod.rs index b3ce46b3..c2fb87c6 100644 --- a/crates/metaltile-runtime/src/device/cuda/mod.rs +++ b/crates/metaltile-runtime/src/device/cuda/mod.rs @@ -16,19 +16,20 @@ mod ffi; -use std::collections::{BTreeMap, HashMap}; -use std::ffi::{CStr, CString}; -use std::os::raw::{c_char, c_int, c_void}; -use std::ptr; -use std::sync::Mutex; +use std::{ + collections::{BTreeMap, HashMap}, + ffi::{CStr, CString}, + os::raw::{c_char, c_int, c_void}, + ptr, + sync::Mutex, +}; +use ffi::*; use metaltile_codegen::{CodegenBackend, CudaGenerator}; use metaltile_core::ir::Kernel; use crate::error::MetalTileError; -use ffi::*; - /// Synthesize a Strided param's `_shape` or `_strides` (row-major) buffer /// from the static shape, as little-endian u32s. Unknown dims default to 1. fn synth_strided_meta(shape: &metaltile_core::shape::Shape, strides: bool) -> Vec { @@ -182,7 +183,8 @@ impl Prepared<'_> { /// Build the `cuLaunchKernel` param array: param device-ptrs first, then /// scalar values, each a raw pointer into our stable `dev_ptrs`/`scalars`. fn args(&self) -> Vec<*mut c_void> { - let mut args: Vec<*mut c_void> = Vec::with_capacity(self.dev_ptrs.len() + self.scalars.len()); + let mut args: Vec<*mut c_void> = + Vec::with_capacity(self.dev_ptrs.len() + self.scalars.len()); for p in &self.dev_ptrs { args.push(p as *const CUdeviceptr as *mut c_void); } @@ -325,7 +327,20 @@ impl CudaDevice { // showed raw cuMemAlloc/cuMemFree at 62% of CUDA API time. METALTILE_POOL_ALLOC_OFF=1 // restores direct driver alloc/free (clean A/B). 4GB pool cap (POOL_CAP_BYTES). let pool_enabled = std::env::var("METALTILE_POOL_ALLOC_OFF").is_err(); - Ok(Some(CudaDevice { ctx, cc_major: major, cc_minor: minor, pool: Mutex::new(HashMap::new()), pooled_bytes: Mutex::new(0), pool_enabled, pinned_free: Mutex::new(HashMap::new()), pinned_inflight: Mutex::new(Vec::new()), stream, capturing: std::sync::atomic::AtomicBool::new(false), cublas: Mutex::new(0), cublaslt: Mutex::new((0, 0)) })) + Ok(Some(CudaDevice { + ctx, + cc_major: major, + cc_minor: minor, + pool: Mutex::new(HashMap::new()), + pooled_bytes: Mutex::new(0), + pool_enabled, + pinned_free: Mutex::new(HashMap::new()), + pinned_inflight: Mutex::new(Vec::new()), + stream, + capturing: std::sync::atomic::AtomicBool::new(false), + cublas: Mutex::new(0), + cublaslt: Mutex::new((0, 0)), + })) } } @@ -353,7 +368,10 @@ impl CudaDevice { // METALTILE_GEMM_ATOMICS=1 opts back into the nondeterministic // heuristic (for A/B perf measurement only). let atomics = std::env::var("METALTILE_GEMM_ATOMICS").ok().as_deref() == Some("1"); - cublasSetAtomicsMode(h, if atomics { CUBLAS_ATOMICS_ALLOWED } else { CUBLAS_ATOMICS_NOT_ALLOWED }); + cublasSetAtomicsMode( + h, + if atomics { CUBLAS_ATOMICS_ALLOWED } else { CUBLAS_ATOMICS_NOT_ALLOWED }, + ); } // Persistent 32MB cuBLAS workspace → GEMMs are CUDA-graph capture/replay // safe (default workspace is a transient per-call alloc that the graph @@ -361,7 +379,9 @@ impl CudaDevice { // / null read on replay). Allocated once, never freed (device-lifetime). let ws_bytes = 32usize * 1024 * 1024; if let Ok(ws) = self.alloc_raw(ws_bytes) { - unsafe { cublasSetWorkspace_v2(h, ws as *mut std::ffi::c_void, ws_bytes); } + unsafe { + cublasSetWorkspace_v2(h, ws as *mut std::ffi::c_void, ws_bytes); + } } *guard = h as usize; } @@ -381,7 +401,10 @@ impl CudaDevice { use std::sync::OnceLock; static ALGO: OnceLock = OnceLock::new(); *ALGO.get_or_init(|| { - match std::env::var("METALTILE_GEMM_ALGO").ok().and_then(|s| s.trim().parse::().ok()) { + match std::env::var("METALTILE_GEMM_ALGO") + .ok() + .and_then(|s| s.trim().parse::().ok()) + { Some(n) if (0..=15).contains(&n) => CUBLAS_GEMM_ALGO0_TENSOR_OP + n as c_int, _ => CUBLAS_GEMM_DEFAULT_TENSOR_OP, } @@ -422,7 +445,10 @@ impl CudaDevice { let cdt = match dtype { metaltile_core::DType::F16 => CUDA_R_16F, metaltile_core::DType::BF16 => CUDA_R_16BF, - other => return Err(MetalTileError::Dispatch(format!("gemm_cublas: unsupported dtype {other:?} (need f16/bf16)"))), + other => + return Err(MetalTileError::Dispatch(format!( + "gemm_cublas: unsupported dtype {other:?} (need f16/bf16)" + ))), }; let alpha: f32 = 1.0; let beta: f32 = 0.0; @@ -433,22 +459,30 @@ impl CudaDevice { let st = unsafe { cublasGemmEx( h, - CUBLAS_OP_T, // op(A) = Aᵀ - CUBLAS_OP_N, // op(B) = B - n as c_int, // rows of op(A) and C (col-major) - m as c_int, // cols of op(B) and C - k as c_int, // shared dim + CUBLAS_OP_T, // op(A) = Aᵀ + CUBLAS_OP_N, // op(B) = B + n as c_int, // rows of op(A) and C (col-major) + m as c_int, // cols of op(B) and C + k as c_int, // shared dim &alpha as *const f32 as *const c_void, - w, cdt, k as c_int, // A = W, lda = k - x, cdt, k as c_int, // B = X, ldb = k + w, + cdt, + k as c_int, // A = W, lda = k + x, + cdt, + k as c_int, // B = X, ldb = k &beta as *const f32 as *const c_void, - out, cdt, n as c_int, // C = out, ldc = n + out, + cdt, + n as c_int, // C = out, ldc = n CUBLAS_COMPUTE_32F, - Self::gemm_algo(), // deterministic explicit algo (see gemm_algo) + Self::gemm_algo(), // deterministic explicit algo (see gemm_algo) ) }; if st != CUBLAS_STATUS_SUCCESS { - return Err(MetalTileError::Dispatch(format!("cublasGemmEx failed: status {st} (m={m} n={n} k={k})"))); + return Err(MetalTileError::Dispatch(format!( + "cublasGemmEx failed: status {st} (m={m} n={n} k={k})" + ))); } Ok(()) } @@ -476,38 +510,61 @@ impl CudaDevice { /// `out[t,n] = Σ_k a[t,k]·w[eid][n,k]`. Beats cuBLAS per-expert on the skinny /// MoE shape. Errors if the runtime was built without CUTLASS. pub fn moe_grouped_cutlass( - &self, a: CUdeviceptr, w: CUdeviceptr, c: CUdeviceptr, - group_rows: &[i32], expert_ids: &[i32], n: usize, k: usize, + &self, + a: CUdeviceptr, + w: CUdeviceptr, + c: CUdeviceptr, + group_rows: &[i32], + expert_ids: &[i32], + n: usize, + k: usize, ) -> Result<(), MetalTileError> { #[cfg(have_cutlass)] { unsafe extern "C" { fn moe_grouped_gemm_cutlass( - a: *const c_void, w: *const c_void, c: *mut c_void, - group_rows: *const c_int, expert_ids: *const c_int, - n_groups: c_int, n: c_int, k: c_int, stream: *mut c_void, + a: *const c_void, + w: *const c_void, + c: *mut c_void, + group_rows: *const c_int, + expert_ids: *const c_int, + n_groups: c_int, + n: c_int, + k: c_int, + stream: *mut c_void, ) -> c_int; } if group_rows.len() != expert_ids.len() { - return Err(MetalTileError::Dispatch("moe_grouped_cutlass: group_rows/expert_ids len mismatch".into())); + return Err(MetalTileError::Dispatch( + "moe_grouped_cutlass: group_rows/expert_ids len mismatch".into(), + )); } let r = unsafe { moe_grouped_gemm_cutlass( - a as *const c_void, w as *const c_void, c as *mut c_void, - group_rows.as_ptr(), expert_ids.as_ptr(), - group_rows.len() as c_int, n as c_int, k as c_int, + a as *const c_void, + w as *const c_void, + c as *mut c_void, + group_rows.as_ptr(), + expert_ids.as_ptr(), + group_rows.len() as c_int, + n as c_int, + k as c_int, self.stream as *mut c_void, ) }; if r != 0 { - return Err(MetalTileError::Dispatch(format!("moe_grouped_gemm_cutlass failed: code {r}"))); + return Err(MetalTileError::Dispatch(format!( + "moe_grouped_gemm_cutlass failed: code {r}" + ))); } Ok(()) } #[cfg(not(have_cutlass))] { let _ = (a, w, c, group_rows, expert_ids, n, k); - Err(MetalTileError::Dispatch("moe_grouped_cutlass: runtime built without CUTLASS (set CUTLASS_DIR)".into())) + Err(MetalTileError::Dispatch( + "moe_grouped_cutlass: runtime built without CUTLASS (set CUTLASS_DIR)".into(), + )) } } @@ -557,13 +614,19 @@ impl CudaDevice { let cdt = match dtype { metaltile_core::DType::F16 => CUDA_R_16F, metaltile_core::DType::BF16 => CUDA_R_16BF, - other => return Err(MetalTileError::Dispatch(format!("gemm_cublaslt: unsupported dtype {other:?} (need f16/bf16)"))), + other => + return Err(MetalTileError::Dispatch(format!( + "gemm_cublaslt: unsupported dtype {other:?} (need f16/bf16)" + ))), }; let ddt = match out_dtype { metaltile_core::DType::F16 => CUDA_R_16F, metaltile_core::DType::BF16 => CUDA_R_16BF, metaltile_core::DType::F32 => CUDA_R_32F, - other => return Err(MetalTileError::Dispatch(format!("gemm_cublaslt: unsupported out_dtype {other:?}"))), + other => + return Err(MetalTileError::Dispatch(format!( + "gemm_cublaslt: unsupported out_dtype {other:?}" + ))), }; // Col-major identity (same as gemm_cublas): out_cm[n,m] = (W^T)·X. // op(A)=T, op(B)=N; A=W (rows=k, cols=n, ld=k), B=X (rows=k, cols=m, ld=k), @@ -573,11 +636,23 @@ impl CudaDevice { unsafe { let mut desc: cublasLtMatmulDesc_t = ptr::null_mut(); let s = cublasLtMatmulDescCreate(&mut desc, CUBLAS_COMPUTE_32F, CUDA_R_32F); - if s != CUBLAS_STATUS_SUCCESS { return Err(MetalTileError::Dispatch(format!("cublasLtMatmulDescCreate: {s}"))); } + if s != CUBLAS_STATUS_SUCCESS { + return Err(MetalTileError::Dispatch(format!("cublasLtMatmulDescCreate: {s}"))); + } let opt = CUBLAS_OP_T; let opn = CUBLAS_OP_N; - cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSA, &opt as *const c_int as *const c_void, std::mem::size_of::()); - cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSB, &opn as *const c_int as *const c_void, std::mem::size_of::()); + cublasLtMatmulDescSetAttribute( + desc, + CUBLASLT_MATMUL_DESC_TRANSA, + &opt as *const c_int as *const c_void, + std::mem::size_of::(), + ); + cublasLtMatmulDescSetAttribute( + desc, + CUBLASLT_MATMUL_DESC_TRANSB, + &opn as *const c_int as *const c_void, + std::mem::size_of::(), + ); // Layouts describe the PHYSICAL (untransposed) storage, col-major. // A=W is [n,k] row-major == [k,n] col-major → rows=k, cols=n, ld=k. @@ -593,38 +668,73 @@ impl CudaDevice { let mut pref: cublasLtMatmulPreference_t = ptr::null_mut(); cublasLtMatmulPreferenceCreate(&mut pref); let ws_bytes: usize = CUBLASLT_WORKSPACE_BYTES; - cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &ws_bytes as *const usize as *const c_void, std::mem::size_of::()); + cublasLtMatmulPreferenceSetAttribute( + pref, + CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &ws_bytes as *const usize as *const c_void, + std::mem::size_of::(), + ); // THE determinism lever: restrict allowed reduction schemes to NONE, // so the heuristic never returns a split-K (atomic-accumulate) algo. let red_mask: u32 = CUBLASLT_REDUCTION_SCHEME_NONE; - cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_REDUCTION_SCHEME_MASK, &red_mask as *const u32 as *const c_void, std::mem::size_of::()); + cublasLtMatmulPreferenceSetAttribute( + pref, + CUBLASLT_MATMUL_PREF_REDUCTION_SCHEME_MASK, + &red_mask as *const u32 as *const c_void, + std::mem::size_of::(), + ); let mut result = cublasLtMatmulHeuristicResult_t::default(); let mut returned: c_int = 0; - let hs = cublasLtMatmulAlgoGetHeuristic(lt, desc, a_l, b_l, d_l, d_l, pref, 1, &mut result, &mut returned); + let hs = cublasLtMatmulAlgoGetHeuristic( + lt, + desc, + a_l, + b_l, + d_l, + d_l, + pref, + 1, + &mut result, + &mut returned, + ); if hs != CUBLAS_STATUS_SUCCESS || returned < 1 { cublasLtMatmulPreferenceDestroy(pref); - cublasLtMatrixLayoutDestroy(a_l); cublasLtMatrixLayoutDestroy(b_l); cublasLtMatrixLayoutDestroy(d_l); + cublasLtMatrixLayoutDestroy(a_l); + cublasLtMatrixLayoutDestroy(b_l); + cublasLtMatrixLayoutDestroy(d_l); cublasLtMatmulDescDestroy(desc); - return Err(MetalTileError::Dispatch(format!("cublasLt: no deterministic algo (m={m} n={n} k={k} status={hs} returned={returned})"))); + return Err(MetalTileError::Dispatch(format!( + "cublasLt: no deterministic algo (m={m} n={n} k={k} status={hs} returned={returned})" + ))); } let mm = cublasLtMatmul( - lt, desc, + lt, + desc, &alpha as *const f32 as *const c_void, - w, a_l, - x, b_l, + w, + a_l, + x, + b_l, &beta as *const f32 as *const c_void, - out, d_l, - out, d_l, + out, + d_l, + out, + d_l, result.algo.as_ptr(), - workspace, ws_bytes, + workspace, + ws_bytes, self.stream, ); cublasLtMatmulPreferenceDestroy(pref); - cublasLtMatrixLayoutDestroy(a_l); cublasLtMatrixLayoutDestroy(b_l); cublasLtMatrixLayoutDestroy(d_l); + cublasLtMatrixLayoutDestroy(a_l); + cublasLtMatrixLayoutDestroy(b_l); + cublasLtMatrixLayoutDestroy(d_l); cublasLtMatmulDescDestroy(desc); if mm != CUBLAS_STATUS_SUCCESS { - return Err(MetalTileError::Dispatch(format!("cublasLtMatmul failed: status {mm} (m={m} n={n} k={k})"))); + return Err(MetalTileError::Dispatch(format!( + "cublasLtMatmul failed: status {mm} (m={m} n={n} k={k})" + ))); } } Ok(()) @@ -639,23 +749,43 @@ impl CudaDevice { #[allow(clippy::too_many_arguments)] pub fn gemm_cublas_strided_batched( &self, - x_base: CUdeviceptr, stride_x: i64, - w_base: CUdeviceptr, stride_w: i64, - out_base: CUdeviceptr, stride_out: i64, - m: usize, n: usize, k: usize, + x_base: CUdeviceptr, + stride_x: i64, + w_base: CUdeviceptr, + stride_w: i64, + out_base: CUdeviceptr, + stride_out: i64, + m: usize, + n: usize, + k: usize, batch_count: usize, dtype: metaltile_core::DType, ) -> Result<(), MetalTileError> { // DETERMINISTIC by default via cublasLt (split-K reductions forbidden); // see gemm_cublas. METALTILE_GEMM_NONDET=1 → legacy nondeterministic path. if std::env::var("METALTILE_GEMM_NONDET").ok().as_deref() != Some("1") { - return self.gemm_cublaslt_strided_batched(x_base, stride_x, w_base, stride_w, out_base, stride_out, m, n, k, batch_count, dtype); + return self.gemm_cublaslt_strided_batched( + x_base, + stride_x, + w_base, + stride_w, + out_base, + stride_out, + m, + n, + k, + batch_count, + dtype, + ); } let h = self.cublas_handle()?; let cdt = match dtype { metaltile_core::DType::F16 => CUDA_R_16F, metaltile_core::DType::BF16 => CUDA_R_16BF, - other => return Err(MetalTileError::Dispatch(format!("gemm_cublas_strided_batched: unsupported dtype {other:?}"))), + other => + return Err(MetalTileError::Dispatch(format!( + "gemm_cublas_strided_batched: unsupported dtype {other:?}" + ))), }; let alpha: f32 = 1.0; let beta: f32 = 0.0; @@ -667,19 +797,32 @@ impl CudaDevice { h, CUBLAS_OP_T, CUBLAS_OP_N, - n as c_int, m as c_int, k as c_int, + n as c_int, + m as c_int, + k as c_int, &alpha as *const f32 as *const c_void, - w_base, cdt, k as c_int, stride_w / el, - x_base, cdt, k as c_int, stride_x / el, + w_base, + cdt, + k as c_int, + stride_w / el, + x_base, + cdt, + k as c_int, + stride_x / el, &beta as *const f32 as *const c_void, - out_base, cdt, n as c_int, stride_out / el, + out_base, + cdt, + n as c_int, + stride_out / el, batch_count as c_int, CUBLAS_COMPUTE_32F, - Self::gemm_algo(), // deterministic explicit algo (see gemm_algo) + Self::gemm_algo(), // deterministic explicit algo (see gemm_algo) ) }; if st != CUBLAS_STATUS_SUCCESS { - return Err(MetalTileError::Dispatch(format!("cublasGemmStridedBatchedEx failed: {st} (m={m} n={n} k={k} batch={batch_count})"))); + return Err(MetalTileError::Dispatch(format!( + "cublasGemmStridedBatchedEx failed: {st} (m={m} n={n} k={k} batch={batch_count})" + ))); } Ok(()) } @@ -689,10 +832,15 @@ impl CudaDevice { #[allow(clippy::too_many_arguments)] pub fn gemm_cublaslt_strided_batched( &self, - x_base: CUdeviceptr, stride_x: i64, - w_base: CUdeviceptr, stride_w: i64, - out_base: CUdeviceptr, stride_out: i64, - m: usize, n: usize, k: usize, + x_base: CUdeviceptr, + stride_x: i64, + w_base: CUdeviceptr, + stride_w: i64, + out_base: CUdeviceptr, + stride_out: i64, + m: usize, + n: usize, + k: usize, batch_count: usize, dtype: metaltile_core::DType, ) -> Result<(), MetalTileError> { @@ -700,7 +848,10 @@ impl CudaDevice { let cdt = match dtype { metaltile_core::DType::F16 => CUDA_R_16F, metaltile_core::DType::BF16 => CUDA_R_16BF, - other => return Err(MetalTileError::Dispatch(format!("gemm_cublaslt_strided_batched: unsupported dtype {other:?}"))), + other => + return Err(MetalTileError::Dispatch(format!( + "gemm_cublaslt_strided_batched: unsupported dtype {other:?}" + ))), }; let el = 2i64; // bytes per f16/bf16 element let bc = batch_count as i32; @@ -713,10 +864,25 @@ impl CudaDevice { unsafe { let mut desc: cublasLtMatmulDesc_t = ptr::null_mut(); let s = cublasLtMatmulDescCreate(&mut desc, CUBLAS_COMPUTE_32F, CUDA_R_32F); - if s != CUBLAS_STATUS_SUCCESS { return Err(MetalTileError::Dispatch(format!("cublasLtMatmulDescCreate(strided): {s}"))); } - let opt = CUBLAS_OP_T; let opn = CUBLAS_OP_N; - cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSA, &opt as *const c_int as *const c_void, std::mem::size_of::()); - cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSB, &opn as *const c_int as *const c_void, std::mem::size_of::()); + if s != CUBLAS_STATUS_SUCCESS { + return Err(MetalTileError::Dispatch(format!( + "cublasLtMatmulDescCreate(strided): {s}" + ))); + } + let opt = CUBLAS_OP_T; + let opn = CUBLAS_OP_N; + cublasLtMatmulDescSetAttribute( + desc, + CUBLASLT_MATMUL_DESC_TRANSA, + &opt as *const c_int as *const c_void, + std::mem::size_of::(), + ); + cublasLtMatmulDescSetAttribute( + desc, + CUBLASLT_MATMUL_DESC_TRANSB, + &opn as *const c_int as *const c_void, + std::mem::size_of::(), + ); // Same per-matrix layouts as gemm_cublaslt, plus batch count + stride. let mut a_l: cublasLtMatrixLayout_t = ptr::null_mut(); @@ -726,8 +892,18 @@ impl CudaDevice { cublasLtMatrixLayoutCreate(&mut b_l, cdt, k as u64, m as u64, k as i64); cublasLtMatrixLayoutCreate(&mut d_l, cdt, n as u64, m as u64, n as i64); let set_batch = |layout: cublasLtMatrixLayout_t, stride: i64| { - cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &bc as *const i32 as *const c_void, std::mem::size_of::()); - cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &stride as *const i64 as *const c_void, std::mem::size_of::()); + cublasLtMatrixLayoutSetAttribute( + layout, + CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, + &bc as *const i32 as *const c_void, + std::mem::size_of::(), + ); + cublasLtMatrixLayoutSetAttribute( + layout, + CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, + &stride as *const i64 as *const c_void, + std::mem::size_of::(), + ); }; set_batch(a_l, off_a); set_batch(b_l, off_b); @@ -736,36 +912,71 @@ impl CudaDevice { let mut pref: cublasLtMatmulPreference_t = ptr::null_mut(); cublasLtMatmulPreferenceCreate(&mut pref); let ws_bytes: usize = CUBLASLT_WORKSPACE_BYTES; - cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &ws_bytes as *const usize as *const c_void, std::mem::size_of::()); + cublasLtMatmulPreferenceSetAttribute( + pref, + CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &ws_bytes as *const usize as *const c_void, + std::mem::size_of::(), + ); let red_mask: u32 = CUBLASLT_REDUCTION_SCHEME_NONE; - cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_REDUCTION_SCHEME_MASK, &red_mask as *const u32 as *const c_void, std::mem::size_of::()); + cublasLtMatmulPreferenceSetAttribute( + pref, + CUBLASLT_MATMUL_PREF_REDUCTION_SCHEME_MASK, + &red_mask as *const u32 as *const c_void, + std::mem::size_of::(), + ); let mut result = cublasLtMatmulHeuristicResult_t::default(); let mut returned: c_int = 0; - let hs = cublasLtMatmulAlgoGetHeuristic(lt, desc, a_l, b_l, d_l, d_l, pref, 1, &mut result, &mut returned); + let hs = cublasLtMatmulAlgoGetHeuristic( + lt, + desc, + a_l, + b_l, + d_l, + d_l, + pref, + 1, + &mut result, + &mut returned, + ); if hs != CUBLAS_STATUS_SUCCESS || returned < 1 { cublasLtMatmulPreferenceDestroy(pref); - cublasLtMatrixLayoutDestroy(a_l); cublasLtMatrixLayoutDestroy(b_l); cublasLtMatrixLayoutDestroy(d_l); + cublasLtMatrixLayoutDestroy(a_l); + cublasLtMatrixLayoutDestroy(b_l); + cublasLtMatrixLayoutDestroy(d_l); cublasLtMatmulDescDestroy(desc); - return Err(MetalTileError::Dispatch(format!("cublasLt(strided): no deterministic algo (m={m} n={n} k={k} batch={batch_count} status={hs} returned={returned})"))); + return Err(MetalTileError::Dispatch(format!( + "cublasLt(strided): no deterministic algo (m={m} n={n} k={k} batch={batch_count} status={hs} returned={returned})" + ))); } let mm = cublasLtMatmul( - lt, desc, + lt, + desc, &alpha as *const f32 as *const c_void, - w_base, a_l, - x_base, b_l, + w_base, + a_l, + x_base, + b_l, &beta as *const f32 as *const c_void, - out_base, d_l, - out_base, d_l, + out_base, + d_l, + out_base, + d_l, result.algo.as_ptr(), - workspace, ws_bytes, + workspace, + ws_bytes, self.stream, ); cublasLtMatmulPreferenceDestroy(pref); - cublasLtMatrixLayoutDestroy(a_l); cublasLtMatrixLayoutDestroy(b_l); cublasLtMatrixLayoutDestroy(d_l); + cublasLtMatrixLayoutDestroy(a_l); + cublasLtMatrixLayoutDestroy(b_l); + cublasLtMatrixLayoutDestroy(d_l); cublasLtMatmulDescDestroy(desc); if mm != CUBLAS_STATUS_SUCCESS { - return Err(MetalTileError::Dispatch(format!("cublasLtMatmul(strided) failed: status {mm} (m={m} n={n} k={k} batch={batch_count})"))); + return Err(MetalTileError::Dispatch(format!( + "cublasLtMatmul(strided) failed: status {mm} (m={m} n={n} k={k} batch={batch_count})" + ))); } } Ok(()) @@ -783,25 +994,30 @@ impl CudaDevice { #[allow(clippy::too_many_arguments)] pub fn gemm_cublas_grouped( &self, - x_ptrs: &[u64], // [group_count] device pointers for X - w_ptrs: &[u64], // [group_count] device pointers for W - out_ptrs: &[u64], // [group_count] device pointers for Out - m_per_group: &[i32], // [group_count] rows per group (token count) - n: usize, // shared output dim - k: usize, // shared input dim + x_ptrs: &[u64], // [group_count] device pointers for X + w_ptrs: &[u64], // [group_count] device pointers for W + out_ptrs: &[u64], // [group_count] device pointers for Out + m_per_group: &[i32], // [group_count] rows per group (token count) + n: usize, // shared output dim + k: usize, // shared input dim dtype: metaltile_core::DType, ) -> Result<(), MetalTileError> { let group_count = x_ptrs.len(); assert_eq!(w_ptrs.len(), group_count); assert_eq!(out_ptrs.len(), group_count); assert_eq!(m_per_group.len(), group_count); - if group_count == 0 { return Ok(()); } + if group_count == 0 { + return Ok(()); + } let h = self.cublas_handle()?; let cdt = match dtype { metaltile_core::DType::F16 => CUDA_R_16F, metaltile_core::DType::BF16 => CUDA_R_16BF, - other => return Err(MetalTileError::Dispatch(format!("gemm_cublas_grouped: unsupported dtype {other:?}"))), + other => + return Err(MetalTileError::Dispatch(format!( + "gemm_cublas_grouped: unsupported dtype {other:?}" + ))), }; // cublasGemmGroupedBatchedEx: col-major same as GemmEx. // Row-major C[m_i,n] = X[m_i,k] · W[n,k]^T @@ -811,16 +1027,16 @@ impl CudaDevice { let k_i = k as i32; // alpha/beta: single scalar shared across all groups (same as GemmEx) let alpha_scalar: f32 = 1.0; - let beta_scalar: f32 = 0.0; + let beta_scalar: f32 = 0.0; let transas: Vec = vec![CUBLAS_OP_T; group_count]; let transbs: Vec = vec![CUBLAS_OP_N; group_count]; // m_array[i] = n (cm rows = output cols), n_array[i] = m_per_group[i] (cm cols = token rows) let m_arr: Vec = vec![n_i; group_count]; let n_arr: Vec = m_per_group.to_vec(); let k_arr: Vec = vec![k_i; group_count]; - let lda: Vec = vec![k_i; group_count]; // A=W, lda=k - let ldb: Vec = vec![k_i; group_count]; // B=X, ldb=k - let ldc: Vec = vec![n_i; group_count]; // C=out, ldc=n + let lda: Vec = vec![k_i; group_count]; // A=W, lda=k + let ldb: Vec = vec![k_i; group_count]; // B=X, ldb=k + let ldc: Vec = vec![n_i; group_count]; // C=out, ldc=n let group_sizes: Vec = vec![1i32; group_count]; // Build void* pointer arrays (W, X, Out) for cuBLAS. @@ -837,15 +1053,16 @@ impl CudaDevice { let triple_bytes = ptr_bytes * 3; // Layout: [A(W) | B(X) | C(Out)] contiguous. let mut staging: Vec = Vec::with_capacity(group_count * 3); - staging.extend_from_slice(w_ptrs); // A = W device pointers - staging.extend_from_slice(x_ptrs); // B = X device pointers + staging.extend_from_slice(w_ptrs); // A = W device pointers + staging.extend_from_slice(x_ptrs); // B = X device pointers staging.extend_from_slice(out_ptrs); // C = Out device pointers let base = self.alloc_raw(triple_bytes)?; let a_dev: CUdeviceptr = base; let b_dev: CUdeviceptr = base + ptr_bytes as CUdeviceptr; let c_dev: CUdeviceptr = base + (2 * ptr_bytes) as CUdeviceptr; - let all_bytes = unsafe { std::slice::from_raw_parts(staging.as_ptr() as *const u8, triple_bytes) }; + let all_bytes = + unsafe { std::slice::from_raw_parts(staging.as_ptr() as *const u8, triple_bytes) }; // Single synchronous H2D: `staging` is pageable host memory, so a sync copy // guarantees it's fully consumed before this Vec drops (correctness over the // marginal async win for a ≤few-KB pointer array). cuBLAS reads a_dev/b_dev/ @@ -854,7 +1071,17 @@ impl CudaDevice { // after any prior GEMM's stream read of a recycled copy (see the same // fix + rationale in gemm_cublas_batched). Pageable src ⇒ host-sync copy, // so `staging` stays valid through the call. - cu_check(unsafe { cuMemcpyHtoDAsync_v2(a_dev, all_bytes.as_ptr() as *const c_void, triple_bytes, self.stream) }, "cuMemcpyHtoDAsync(grouped_ptrs)")?; + cu_check( + unsafe { + cuMemcpyHtoDAsync_v2( + a_dev, + all_bytes.as_ptr() as *const c_void, + triple_bytes, + self.stream, + ) + }, + "cuMemcpyHtoDAsync(grouped_ptrs)", + )?; let st = unsafe { cublasGemmGroupedBatchedEx( @@ -908,18 +1135,25 @@ impl CudaDevice { x_ptrs: &[u64], // [batch] device ptr per X (rows operand) w_ptrs: &[u64], // [batch] device ptr per W (weight operand) out_ptrs: &[u64], // [batch] device ptr per Out - m: usize, n: usize, k: usize, + m: usize, + n: usize, + k: usize, dtype: metaltile_core::DType, ) -> Result<(), MetalTileError> { let batch_count = x_ptrs.len(); assert_eq!(w_ptrs.len(), batch_count); assert_eq!(out_ptrs.len(), batch_count); - if batch_count == 0 { return Ok(()); } + if batch_count == 0 { + return Ok(()); + } let h = self.cublas_handle()?; let cdt = match dtype { metaltile_core::DType::F16 => CUDA_R_16F, metaltile_core::DType::BF16 => CUDA_R_16BF, - other => return Err(MetalTileError::Dispatch(format!("gemm_cublas_batched: unsupported dtype {other:?}"))), + other => + return Err(MetalTileError::Dispatch(format!( + "gemm_cublas_batched: unsupported dtype {other:?}" + ))), }; let alpha: f32 = 1.0; let beta: f32 = 0.0; @@ -931,14 +1165,15 @@ impl CudaDevice { let ptr_bytes = batch_count * 8; let triple_bytes = ptr_bytes * 3; let mut staging: Vec = Vec::with_capacity(batch_count * 3); - staging.extend_from_slice(w_ptrs); // A = W - staging.extend_from_slice(x_ptrs); // B = X + staging.extend_from_slice(w_ptrs); // A = W + staging.extend_from_slice(x_ptrs); // B = X staging.extend_from_slice(out_ptrs); // C = Out let base = self.alloc_raw(triple_bytes)?; let a_dev: CUdeviceptr = base; let b_dev: CUdeviceptr = base + ptr_bytes as CUdeviceptr; let c_dev: CUdeviceptr = base + (2 * ptr_bytes) as CUdeviceptr; - let all_bytes = unsafe { std::slice::from_raw_parts(staging.as_ptr() as *const u8, triple_bytes) }; + let all_bytes = + unsafe { std::slice::from_raw_parts(staging.as_ptr() as *const u8, triple_bytes) }; // H2D on self.stream (NOT the null stream): the pooled ptr-array buffer is // recycled across calls, and the PRIOR batched/grouped GEMM reading its // copy of the array runs on self.stream. A synchronous null-stream copy @@ -947,18 +1182,36 @@ impl CudaDevice { // chunk of a 4-chunk SSD scan corrupted). Issuing the copy on self.stream // makes stream order serialize prior-read → this-write. (Pageable src ⇒ // the async copy is host-synchronous, so `staging` stays valid.) - cu_check(unsafe { cuMemcpyHtoDAsync_v2(a_dev, all_bytes.as_ptr() as *const c_void, triple_bytes, self.stream) }, "cuMemcpyHtoDAsync(batched_ptrs)")?; + cu_check( + unsafe { + cuMemcpyHtoDAsync_v2( + a_dev, + all_bytes.as_ptr() as *const c_void, + triple_bytes, + self.stream, + ) + }, + "cuMemcpyHtoDAsync(batched_ptrs)", + )?; let st = unsafe { cublasGemmBatchedEx( h, CUBLAS_OP_T, CUBLAS_OP_N, - n as c_int, m as c_int, k as c_int, + n as c_int, + m as c_int, + k as c_int, &alpha as *const f32 as *const c_void, - a_dev as *const *const c_void, cdt, k as c_int, - b_dev as *const *const c_void, cdt, k as c_int, + a_dev as *const *const c_void, + cdt, + k as c_int, + b_dev as *const *const c_void, + cdt, + k as c_int, &beta as *const f32 as *const c_void, - c_dev as *const *mut c_void, cdt, n as c_int, + c_dev as *const *mut c_void, + cdt, + n as c_int, batch_count as c_int, CUBLAS_COMPUTE_32F, Self::gemm_algo(), @@ -997,8 +1250,9 @@ impl CudaDevice { )?; // Compile to the device's virtual architecture. - let arch = CString::new(format!("--gpu-architecture=compute_{}{}", self.cc_major, self.cc_minor)) - .unwrap(); + let arch = + CString::new(format!("--gpu-architecture=compute_{}{}", self.cc_major, self.cc_minor)) + .unwrap(); // NVRTC does not auto-include the toolkit headers (cuda_fp16.h, // cuda_bf16.h) — point it at /include. let cuda_root = std::env::var("CUDA_PATH") @@ -1015,17 +1269,24 @@ impl CudaDevice { let fixed1 = format!("{cuda_root}/targets/sbsa-linux/include/cccl"); let fixed2 = format!("{cuda_root}/targets/x86_64-linux/include/cccl"); // Also try the versioned cuda-13.x path (some distros install both). - let by_ver = std::fs::read_dir(format!("{cuda_root}-{}.{}/targets", - self.cc_major, self.cc_minor)) - .ok() - .and_then(|mut rd| rd.next()) - .and_then(|e| e.ok()) - .map(|e| format!("{}/include/cccl", e.path().display())); - if std::path::Path::new(&fixed1).exists() { Some(fixed1) } - else if std::path::Path::new(&fixed2).exists() { Some(fixed2) } - else { by_ver.filter(|p| std::path::Path::new(p).exists()) } + let by_ver = std::fs::read_dir(format!( + "{cuda_root}-{}.{}/targets", + self.cc_major, self.cc_minor + )) + .ok() + .and_then(|mut rd| rd.next()) + .and_then(|e| e.ok()) + .map(|e| format!("{}/include/cccl", e.path().display())); + if std::path::Path::new(&fixed1).exists() { + Some(fixed1) + } else if std::path::Path::new(&fixed2).exists() { + Some(fixed2) + } else { + by_ver.filter(|p| std::path::Path::new(p).exists()) + } }; - let cccl_inc = cccl_path.as_ref().map(|p| CString::new(format!("--include-path={p}")).unwrap()); + let cccl_inc = + cccl_path.as_ref().map(|p| CString::new(format!("--include-path={p}")).unwrap()); // Disable contraction of `a*b+c` into FMA by default: the CPU oracle // uses non-fused IEEE arithmetic, so default FMA fusion drifts in // accumulation-heavy kernels (conv, attention, recurrence). Matching @@ -1041,26 +1302,29 @@ impl CudaDevice { // MT_FAST_MATH=1: enable --use_fast_math (implies --fmad=true + fast // intrinsics: __expf, __sinf, etc.). Trades ~1-2 ULP precision for // ~2-4x faster transcendentals. Safe for inference softmax. - let fast_math_on = std::env::var("MT_FAST_MATH").map(|v| v == "1" || v == "true").unwrap_or(false); + let fast_math_on = + std::env::var("MT_FAST_MATH").map(|v| v == "1" || v == "true").unwrap_or(false); let compile_res = match (fast_math_on, cccl_inc.as_ref()) { (true, Some(cccl)) => { let fast = CString::new("--use_fast_math").unwrap(); - let opts: [*const c_char; 4] = [arch.as_ptr(), inc.as_ptr(), cccl.as_ptr(), fast.as_ptr()]; + let opts: [*const c_char; 4] = + [arch.as_ptr(), inc.as_ptr(), cccl.as_ptr(), fast.as_ptr()]; unsafe { nvrtcCompileProgram(prog, opts.len() as _, opts.as_ptr()) } - } + }, (false, Some(cccl)) => { - let opts: [*const c_char; 4] = [arch.as_ptr(), inc.as_ptr(), cccl.as_ptr(), fmad.as_ptr()]; + let opts: [*const c_char; 4] = + [arch.as_ptr(), inc.as_ptr(), cccl.as_ptr(), fmad.as_ptr()]; unsafe { nvrtcCompileProgram(prog, opts.len() as _, opts.as_ptr()) } - } + }, (true, None) => { let fast = CString::new("--use_fast_math").unwrap(); let opts: [*const c_char; 3] = [arch.as_ptr(), inc.as_ptr(), fast.as_ptr()]; unsafe { nvrtcCompileProgram(prog, opts.len() as _, opts.as_ptr()) } - } + }, (false, None) => { let opts: [*const c_char; 3] = [arch.as_ptr(), inc.as_ptr(), fmad.as_ptr()]; unsafe { nvrtcCompileProgram(prog, opts.len() as _, opts.as_ptr()) } - } + }, }; // Always fetch the log — it carries the actual compiler diagnostics. @@ -1140,10 +1404,7 @@ impl CudaDevice { }, "cuMemcpyHtoDAsync(upload)", )?; - cu_check( - unsafe { cuStreamSynchronize(self.stream) }, - "cuStreamSynchronize(upload)", - )?; + cu_check(unsafe { cuStreamSynchronize(self.stream) }, "cuStreamSynchronize(upload)")?; } Ok(buf) } @@ -1265,21 +1526,20 @@ impl CudaDevice { // interleaved with the on-stream expert GEMMs). if len > 262_144 { return cu_check( - unsafe { cuMemcpyHtoDAsync_v2(ptr, bytes.as_ptr() as *const c_void, len, self.stream) }, + unsafe { + cuMemcpyHtoDAsync_v2(ptr, bytes.as_ptr() as *const c_void, len, self.stream) + }, "cuMemcpyHtoDAsync(large)", ); } - let pinned = self - .pinned_free - .lock() - .unwrap() - .get_mut(&len) - .and_then(|v| v.pop()) - .unwrap_or_else(|| { - let mut p: *mut c_void = ptr::null_mut(); - unsafe { cuMemAllocHost_v2(&mut p, len) }; - p as usize - }); + let pinned = + self.pinned_free.lock().unwrap().get_mut(&len).and_then(|v| v.pop()).unwrap_or_else( + || { + let mut p: *mut c_void = ptr::null_mut(); + unsafe { cuMemAllocHost_v2(&mut p, len) }; + p as usize + }, + ); unsafe { ptr::copy_nonoverlapping(bytes.as_ptr(), pinned as *mut u8, len) }; cu_check( unsafe { cuMemcpyHtoDAsync_v2(ptr, pinned as *const c_void, len, self.stream) }, @@ -1349,8 +1609,12 @@ impl CudaDevice { unsafe { cuLaunchKernel( func.func, - grid[0], grid[1], grid[2], - block[0], block[1], block[2], + grid[0], + grid[1], + grid[2], + block[0], + block[1], + block[2], shared_bytes, self.stream, args.as_mut_ptr(), @@ -1380,8 +1644,12 @@ impl CudaDevice { unsafe { cuLaunchKernel( func.func, - grid[0], grid[1], grid[2], - block[0], block[1], block[2], + grid[0], + grid[1], + grid[2], + block[0], + block[1], + block[2], shared_bytes, self.stream, args.as_mut_ptr(), @@ -1410,8 +1678,12 @@ impl CudaDevice { unsafe { cuLaunchCooperativeKernel( func.func, - grid[0], grid[1], grid[2], - block[0], block[1], block[2], + grid[0], + grid[1], + grid[2], + block[0], + block[1], + block[2], shared_bytes, self.stream, args.as_mut_ptr(), @@ -1424,16 +1696,17 @@ impl CudaDevice { /// Returns `true` if a CUDA-graph capture is currently in progress on this /// device's stream. Callers that use non-capturable launches (e.g. cooperative /// kernel) must fall back to a capturable alternative during capture. - pub fn is_capturing(&self) -> bool { - self.capturing.load(std::sync::atomic::Ordering::SeqCst) - } + pub fn is_capturing(&self) -> bool { self.capturing.load(std::sync::atomic::Ordering::SeqCst) } /// Begin recording all subsequent stream work into a CUDA graph (phase-1 /// megakernel). Caller must run a NO-host-sync (all-device) sequence, then /// `end_capture`. THREAD_LOCAL mode scopes capture to this thread's stream. pub fn begin_capture(&self) -> Result<(), MetalTileError> { self.capturing.store(true, std::sync::atomic::Ordering::SeqCst); - cu_check(unsafe { cuStreamBeginCapture_v2(self.stream, CU_STREAM_CAPTURE_MODE_THREAD_LOCAL) }, "cuStreamBeginCapture") + cu_check( + unsafe { cuStreamBeginCapture_v2(self.stream, CU_STREAM_CAPTURE_MODE_THREAD_LOCAL) }, + "cuStreamBeginCapture", + ) } /// Finish capture → instantiate an executable graph. Replay with `graph_launch`. @@ -1442,7 +1715,10 @@ impl CudaDevice { cu_check(unsafe { cuStreamEndCapture(self.stream, &mut graph) }, "cuStreamEndCapture")?; self.capturing.store(false, std::sync::atomic::Ordering::SeqCst); let mut exec: CUgraphExec = ptr::null_mut(); - cu_check(unsafe { cuGraphInstantiateWithFlags(&mut exec, graph, 0) }, "cuGraphInstantiate")?; + cu_check( + unsafe { cuGraphInstantiateWithFlags(&mut exec, graph, 0) }, + "cuGraphInstantiate", + )?; unsafe { cuGraphDestroy(graph) }; Ok(exec) } @@ -1538,9 +1814,9 @@ impl CudaDevice { let mut scalars: Vec> = Vec::new(); for ce in &kernel.constexprs { let name = ce.name.name(); - let bytes = buffers.get(name).ok_or_else(|| { - MetalTileError::Dispatch(format!("missing constexpr '{name}'")) - })?; + let bytes = buffers + .get(name) + .ok_or_else(|| MetalTileError::Dispatch(format!("missing constexpr '{name}'")))?; scalars.push(bytes.clone()); } if kernel.mode == metaltile_core::ir::KernelMode::Elementwise { @@ -1557,15 +1833,7 @@ impl CudaDevice { scalars.push(n_elems.to_le_bytes().to_vec()); } - Ok(Prepared { - _module: module, - func, - dev_bufs, - dev_ptrs, - scalars, - out_meta, - shared_bytes, - }) + Ok(Prepared { _module: module, func, dev_bufs, dev_ptrs, scalars, out_meta, shared_bytes }) } /// End-to-end generic dispatch: generate CUDA for `kernel`, compile, @@ -1644,8 +1912,12 @@ impl CudaDevice { unsafe { cuLaunchKernel( prep.func.func, - grid[0], grid[1], grid[2], - block[0], block[1], block[2], + grid[0], + grid[1], + grid[2], + block[0], + block[1], + block[2], prep.shared_bytes, ptr::null_mut(), args.as_mut_ptr(), diff --git a/crates/metaltile-runtime/src/device/hip/ffi.rs b/crates/metaltile-runtime/src/device/hip/ffi.rs index 785e9439..3639c8ae 100644 --- a/crates/metaltile-runtime/src/device/hip/ffi.rs +++ b/crates/metaltile-runtime/src/device/hip/ffi.rs @@ -55,11 +55,7 @@ pub const HIP_EVENT_DEFAULT: c_uint = 0; unsafe extern "C" { pub fn hipInit(flags: c_uint) -> hipError_t; pub fn hipDeviceGet(device: *mut hipDevice_t, ordinal: c_int) -> hipError_t; - pub fn hipDeviceGetAttribute( - pi: *mut c_int, - attrib: c_int, - dev: hipDevice_t, - ) -> hipError_t; + pub fn hipDeviceGetAttribute(pi: *mut c_int, attrib: c_int, dev: hipDevice_t) -> hipError_t; pub fn hipCtxCreate(pctx: *mut hipCtx_t, flags: c_uint, dev: hipDevice_t) -> hipError_t; pub fn hipCtxDestroy(ctx: hipCtx_t) -> hipError_t; pub fn hipCtxSynchronize() -> hipError_t; @@ -71,23 +67,11 @@ unsafe extern "C" { module: hipModule_t, name: *const c_char, ) -> hipError_t; - pub fn hipFuncSetAttribute( - func: hipFunction_t, - attrib: c_int, - value: c_int, - ) -> hipError_t; + pub fn hipFuncSetAttribute(func: hipFunction_t, attrib: c_int, value: c_int) -> hipError_t; pub fn hipMalloc(dptr: *mut hipDeviceptr_t, bytesize: usize) -> hipError_t; pub fn hipFree(dptr: hipDeviceptr_t) -> hipError_t; - pub fn hipMemcpyHtoD( - dst: hipDeviceptr_t, - src: *const c_void, - byte_count: usize, - ) -> hipError_t; - pub fn hipMemcpyDtoH( - dst: *mut c_void, - src: hipDeviceptr_t, - byte_count: usize, - ) -> hipError_t; + pub fn hipMemcpyHtoD(dst: hipDeviceptr_t, src: *const c_void, byte_count: usize) -> hipError_t; + pub fn hipMemcpyDtoH(dst: *mut c_void, src: hipDeviceptr_t, byte_count: usize) -> hipError_t; #[allow(clippy::too_many_arguments)] pub fn hipModuleLaunchKernel( f: hipFunction_t, diff --git a/crates/metaltile-runtime/src/device/hip/mod.rs b/crates/metaltile-runtime/src/device/hip/mod.rs index 821483a6..d9e3575e 100644 --- a/crates/metaltile-runtime/src/device/hip/mod.rs +++ b/crates/metaltile-runtime/src/device/hip/mod.rs @@ -19,18 +19,19 @@ mod ffi; -use std::collections::BTreeMap; -use std::ffi::{CStr, CString}; -use std::os::raw::{c_char, c_int, c_void}; -use std::ptr; +use std::{ + collections::BTreeMap, + ffi::{CStr, CString}, + os::raw::{c_char, c_int, c_void}, + ptr, +}; +use ffi::*; use metaltile_codegen::{CodegenBackend, HipGenerator}; use metaltile_core::ir::Kernel; use crate::error::MetalTileError; -use ffi::*; - fn synth_strided_meta(shape: &metaltile_core::shape::Shape, strides: bool) -> Vec { use metaltile_core::shape::Dim; let dims: Vec = (0..shape.rank()) @@ -204,8 +205,7 @@ impl HipDevice { // `METALTILE_HIP_GFX` if set, else default to gfx1201 (RDNA 4 / // RX 9070 XT — the user's primary target). Override for anything // else: gfx1100 RDNA 3, gfx942 MI300, gfx950 MI350. - let gfx = std::env::var("METALTILE_HIP_GFX") - .unwrap_or_else(|_| "gfx1201".to_string()); + let gfx = std::env::var("METALTILE_HIP_GFX").unwrap_or_else(|_| "gfx1201".to_string()); // Derive wave size from gfx family. The CUDA-style attribute // query (HIP_DEVICE_ATTRIBUTE_WARP_SIZE) is unreliable on Windows // ROCm 7.x — the enum index drifts between releases — so we use @@ -248,15 +248,10 @@ impl HipDevice { /// Compile HIP C++ source → AMDGPU code-object → loaded module via /// hipRTC + `hipModuleLoadData`. - pub fn compile( - &self, - src: &str, - prog_name: &str, - ) -> Result { - let csrc = - CString::new(src).map_err(|e| MetalTileError::Compilation(e.to_string()))?; - let cname = CString::new(prog_name) - .map_err(|e| MetalTileError::Compilation(e.to_string()))?; + pub fn compile(&self, src: &str, prog_name: &str) -> Result { + let csrc = CString::new(src).map_err(|e| MetalTileError::Compilation(e.to_string()))?; + let cname = + CString::new(prog_name).map_err(|e| MetalTileError::Compilation(e.to_string()))?; let mut prog: hiprtcProgram = ptr::null_mut(); hiprtc_check( @@ -299,10 +294,9 @@ impl HipDevice { } }); let inc = CString::new(format!("-I{hip_root}/include")).unwrap(); - let opts: [*const c_char; 4] = [arch.as_ptr(), no_fma.as_ptr(), prec_div.as_ptr(), inc.as_ptr()]; - let compile_res = unsafe { - hiprtcCompileProgram(prog, opts.len() as _, opts.as_ptr()) - }; + let opts: [*const c_char; 4] = + [arch.as_ptr(), no_fma.as_ptr(), prec_div.as_ptr(), inc.as_ptr()]; + let compile_res = unsafe { hiprtcCompileProgram(prog, opts.len() as _, opts.as_ptr()) }; let log = unsafe { let mut log_size: usize = 0; @@ -310,8 +304,7 @@ impl HipDevice { if log_size > 1 { let mut buf = vec![0u8; log_size]; hiprtcGetProgramLog(prog, buf.as_mut_ptr() as *mut c_char); - String::from_utf8_lossy(&buf[..log_size.saturating_sub(1)]) - .into_owned() + String::from_utf8_lossy(&buf[..log_size.saturating_sub(1)]).into_owned() } else { String::new() } @@ -320,9 +313,7 @@ impl HipDevice { if compile_res != HIPRTC_SUCCESS { unsafe { hiprtcDestroyProgram(&mut prog) }; let msg = unsafe { - CStr::from_ptr(hiprtcGetErrorString(compile_res)) - .to_string_lossy() - .into_owned() + CStr::from_ptr(hiprtcGetErrorString(compile_res)).to_string_lossy().into_owned() }; return Err(MetalTileError::Compilation(format!( "hiprtcCompileProgram failed: {msg}\n--- log ---\n{log}" @@ -334,10 +325,7 @@ impl HipDevice { let mut sz: usize = 0; hiprtc_check(hiprtcGetCodeSize(prog, &mut sz), "hiprtcGetCodeSize")?; let mut buf = vec![0u8; sz]; - hiprtc_check( - hiprtcGetCode(prog, buf.as_mut_ptr() as *mut c_char), - "hiprtcGetCode", - )?; + hiprtc_check(hiprtcGetCode(prog, buf.as_mut_ptr() as *mut c_char), "hiprtcGetCode")?; buf }; unsafe { hiprtcDestroyProgram(&mut prog) }; @@ -363,20 +351,14 @@ impl HipDevice { let buf = self.alloc(data.len())?; if !data.is_empty() { hip_check( - unsafe { - hipMemcpyHtoD(buf.ptr, data.as_ptr() as *const c_void, data.len()) - }, + unsafe { hipMemcpyHtoD(buf.ptr, data.as_ptr() as *const c_void, data.len()) }, "hipMemcpyHtoD", )?; } Ok(buf) } - pub fn download( - &self, - buf: &HipBuffer, - out: &mut [u8], - ) -> Result<(), MetalTileError> { + pub fn download(&self, buf: &HipBuffer, out: &mut [u8]) -> Result<(), MetalTileError> { let n = out.len().min(buf.len); if n == 0 { return Ok(()); @@ -431,9 +413,7 @@ impl HipDevice { // still over an implementation-specific cap. Report it // explicitly so the harness can bucket it. let msg = unsafe { - CStr::from_ptr(hipGetErrorString(attr_res)) - .to_string_lossy() - .into_owned() + CStr::from_ptr(hipGetErrorString(attr_res)).to_string_lossy().into_owned() }; return Err(MetalTileError::Dispatch(format!( "hipFuncSetAttribute(MaxDynamicSharedMemorySize={shared_bytes}): {msg}" @@ -444,8 +424,12 @@ impl HipDevice { unsafe { hipModuleLaunchKernel( func.func, - grid[0], grid[1], grid[2], - block[0], block[1], block[2], + grid[0], + grid[1], + grid[2], + block[0], + block[1], + block[2], shared_bytes, ptr::null_mut(), args.as_mut_ptr(), @@ -474,10 +458,7 @@ impl HipDevice { let mut out_meta: Vec> = Vec::new(); for p in &kernel.params { let bytes = buffers.get(&p.name).ok_or_else(|| { - MetalTileError::Dispatch(format!( - "missing buffer for param '{}'", - p.name - )) + MetalTileError::Dispatch(format!("missing buffer for param '{}'", p.name)) })?; let buf = self.upload(bytes)?; dev_ptrs.push(buf.device_ptr()); @@ -502,9 +483,9 @@ impl HipDevice { let mut scalars: Vec> = Vec::new(); for ce in &kernel.constexprs { let name = ce.name.name(); - let bytes = buffers.get(name).ok_or_else(|| { - MetalTileError::Dispatch(format!("missing constexpr '{name}'")) - })?; + let bytes = buffers + .get(name) + .ok_or_else(|| MetalTileError::Dispatch(format!("missing constexpr '{name}'")))?; scalars.push(bytes.clone()); } if kernel.mode == metaltile_core::ir::KernelMode::Elementwise { @@ -514,23 +495,13 @@ impl HipDevice { .position(|p| p.is_output) .and_then(|i| { let p = &kernel.params[i]; - buffers - .get(&p.name) - .map(|b| (b.len() / p.dtype.size_bytes().max(1)) as u32) + buffers.get(&p.name).map(|b| (b.len() / p.dtype.size_bytes().max(1)) as u32) }) .unwrap_or(0); scalars.push(n_elems.to_le_bytes().to_vec()); } - Ok(Prepared { - _module: module, - func, - dev_bufs, - dev_ptrs, - scalars, - out_meta, - shared_bytes, - }) + Ok(Prepared { _module: module, func, dev_bufs, dev_ptrs, scalars, out_meta, shared_bytes }) } /// End-to-end generic dispatch (CUDA `run_kernel` analog). diff --git a/crates/metaltile-runtime/src/device/vulkan/ffi.rs b/crates/metaltile-runtime/src/device/vulkan/ffi.rs index bda33cb2..b072bb31 100644 --- a/crates/metaltile-runtime/src/device/vulkan/ffi.rs +++ b/crates/metaltile-runtime/src/device/vulkan/ffi.rs @@ -670,11 +670,7 @@ unsafe extern "C" { pAllocator: *const c_void, pPipelines: *mut VkPipeline_, ) -> VkResult; - pub fn vkDestroyPipeline( - device: VkDevice, - pipeline: VkPipeline_, - pAllocator: *const c_void, - ); + pub fn vkDestroyPipeline(device: VkDevice, pipeline: VkPipeline_, pAllocator: *const c_void); pub fn vkCreateDescriptorPool( device: VkDevice, @@ -925,11 +921,7 @@ unsafe extern "C" { pAllocator: *const c_void, pPipelines: *mut VkPipeline_, ) -> VkResult; - pub fn vkDestroyPipeline( - device: VkDevice, - pipeline: VkPipeline_, - pAllocator: *const c_void, - ); + pub fn vkDestroyPipeline(device: VkDevice, pipeline: VkPipeline_, pAllocator: *const c_void); pub fn vkCreateDescriptorPool( device: VkDevice, @@ -1110,9 +1102,7 @@ unsafe extern "C" { ) -> shaderc_compilation_status; pub fn shaderc_result_get_length(result: shaderc_compilation_result_t) -> usize; pub fn shaderc_result_get_bytes(result: shaderc_compilation_result_t) -> *const u8; - pub fn shaderc_result_get_error_message( - result: shaderc_compilation_result_t, - ) -> *const c_char; + pub fn shaderc_result_get_error_message(result: shaderc_compilation_result_t) -> *const c_char; } #[cfg(not(windows))] @@ -1146,9 +1136,7 @@ unsafe extern "C" { ) -> shaderc_compilation_status; pub fn shaderc_result_get_length(result: shaderc_compilation_result_t) -> usize; pub fn shaderc_result_get_bytes(result: shaderc_compilation_result_t) -> *const u8; - pub fn shaderc_result_get_error_message( - result: shaderc_compilation_result_t, - ) -> *const c_char; + pub fn shaderc_result_get_error_message(result: shaderc_compilation_result_t) -> *const c_char; } // Silence unused param warnings on c_uint imports if any. diff --git a/crates/metaltile-runtime/src/device/vulkan/mod.rs b/crates/metaltile-runtime/src/device/vulkan/mod.rs index 6f9992fe..54e7194d 100644 --- a/crates/metaltile-runtime/src/device/vulkan/mod.rs +++ b/crates/metaltile-runtime/src/device/vulkan/mod.rs @@ -27,18 +27,19 @@ mod ffi; -use std::collections::BTreeMap; -use std::ffi::{CStr, CString}; -use std::os::raw::{c_char, c_void}; -use std::ptr; +use std::{ + collections::BTreeMap, + ffi::{CStr, CString}, + os::raw::{c_char, c_void}, + ptr, +}; +use ffi::*; use metaltile_codegen::{CodegenBackend, GlslGenerator, spirv::GlslBindingPlan}; use metaltile_core::{dtype::DType, ir::Kernel}; use crate::error::MetalTileError; -use ffi::*; - const ENTRY_POINT: &[u8] = b"main\0"; /// Synthesize a Strided param's `_shape` or `_strides` companion buffer @@ -220,11 +221,7 @@ impl VulkanDevice { // Pick a queue family supporting compute. let mut qcount: u32 = 0; - vkGetPhysicalDeviceQueueFamilyProperties( - physical_device, - &mut qcount, - ptr::null_mut(), - ); + vkGetPhysicalDeviceQueueFamilyProperties(physical_device, &mut qcount, ptr::null_mut()); let mut qprops: Vec = (0..qcount as usize) .map(|_| VkQueueFamilyProperties { queueFlags: 0, @@ -307,13 +304,10 @@ impl VulkanDevice { // feature + the device extension. The subgroup size stays // pinned at 32 (the kernels' staging math assumes 32-lane // subgroups; on wave32 the 16x16 fragment spans 32 lanes). - let mut coop_feat: VkPhysicalDeviceCooperativeMatrixFeaturesKHR = - std::mem::zeroed(); - let coop_ext_ptr = - VK_KHR_COOPERATIVE_MATRIX_EXTENSION_NAME.as_ptr() as *const i8; + let mut coop_feat: VkPhysicalDeviceCooperativeMatrixFeaturesKHR = std::mem::zeroed(); + let coop_ext_ptr = VK_KHR_COOPERATIVE_MATRIX_EXTENSION_NAME.as_ptr() as *const i8; if coopmat { - coop_feat.sType = - VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR; + coop_feat.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR; coop_feat.cooperativeMatrix = VK_TRUE; coop_feat.pNext = feat11.pNext; feat11.pNext = &mut coop_feat as *mut _ as *mut c_void; @@ -327,11 +321,7 @@ impl VulkanDevice { enabledLayerCount: 0, ppEnabledLayerNames: ptr::null(), enabledExtensionCount: if coopmat { 1 } else { 0 }, - ppEnabledExtensionNames: if coopmat { - &coop_ext_ptr - } else { - ptr::null() - }, + ppEnabledExtensionNames: if coopmat { &coop_ext_ptr } else { ptr::null() }, pEnabledFeatures: ptr::null(), }; let mut device: VkDevice = ptr::null_mut(); @@ -340,10 +330,7 @@ impl VulkanDevice { // Retry without the f16/bf16/i8 chain — old drivers, or // devices that don't support these features, get the // Phase-1 f32-only path back. - let plain_ci = VkDeviceCreateInfo { - pNext: ptr::null(), - ..dev_ci - }; + let plain_ci = VkDeviceCreateInfo { pNext: ptr::null(), ..dev_ci }; vk_check( vkCreateDevice(physical_device, &plain_ci, ptr::null(), &mut device), "vkCreateDevice(plain)", @@ -409,15 +396,10 @@ impl VulkanDevice { /// Find a memory type matching `mem_type_bits` (from /// `vkGetBufferMemoryRequirements`) with all the requested property flags. - fn find_memory_type( - &self, - mem_type_bits: u32, - flags: u32, - ) -> Result { + fn find_memory_type(&self, mem_type_bits: u32, flags: u32) -> Result { for i in 0..self.memory_properties.memoryTypeCount { if (mem_type_bits & (1u32 << i)) != 0 - && (self.memory_properties.memoryTypes[i as usize].propertyFlags & flags) - == flags + && (self.memory_properties.memoryTypes[i as usize].propertyFlags & flags) == flags { return Ok(i); } @@ -448,11 +430,7 @@ impl VulkanDevice { vkCreateBuffer(self.device, &bci, ptr::null(), &mut buffer), "vkCreateBuffer", )?; - let mut req = VkMemoryRequirements { - size: 0, - alignment: 0, - memoryTypeBits: 0, - }; + let mut req = VkMemoryRequirements { size: 0, alignment: 0, memoryTypeBits: 0 }; vkGetBufferMemoryRequirements(self.device, buffer, &mut req); let mem_type_index = self.find_memory_type( req.memoryTypeBits, @@ -469,10 +447,7 @@ impl VulkanDevice { vkAllocateMemory(self.device, &ai, ptr::null(), &mut memory), "vkAllocateMemory", )?; - vk_check( - vkBindBufferMemory(self.device, buffer, memory, 0), - "vkBindBufferMemory", - )?; + vk_check(vkBindBufferMemory(self.device, buffer, memory, 0), "vkBindBufferMemory")?; Ok(VulkanBuffer { buffer, memory, size, dev: self }) } } @@ -569,12 +544,7 @@ impl VulkanDevice { }; let mut set_layout: VkDescriptorSetLayout = VK_NULL_HANDLE; vk_check( - vkCreateDescriptorSetLayout( - self.device, - &dsl_ci, - ptr::null(), - &mut set_layout, - ), + vkCreateDescriptorSetLayout(self.device, &dsl_ci, ptr::null(), &mut set_layout), "vkCreateDescriptorSetLayout", )?; @@ -609,8 +579,7 @@ impl VulkanDevice { pNext: ptr::null_mut(), requiredSubgroupSize: 32, }; - let entry = - CStr::from_bytes_with_nul(ENTRY_POINT).unwrap().as_ptr(); + let entry = CStr::from_bytes_with_nul(ENTRY_POINT).unwrap().as_ptr(); let stage = VkPipelineShaderStageCreateInfo { sType: VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO, pNext: &req_subgroup as *const _ as *const c_void, @@ -709,9 +678,9 @@ impl VulkanDevice { let mut push: Vec = Vec::with_capacity(plan.push_constant_bytes as usize); for ce in &kernel.constexprs { let name = ce.name.name(); - let bytes = buffers.get(name).ok_or_else(|| { - MetalTileError::Dispatch(format!("missing constexpr '{name}'")) - })?; + let bytes = buffers + .get(name) + .ok_or_else(|| MetalTileError::Dispatch(format!("missing constexpr '{name}'")))?; push.extend_from_slice(bytes); } if plan.has_n_elems { @@ -747,11 +716,7 @@ impl VulkanDevice { // references our slice). Hold it for the whole unsafe block. let buf_infos: Vec = dev_bufs .iter() - .map(|b| VkDescriptorBufferInfo { - buffer: b.buffer, - offset: 0, - range: b.size, - }) + .map(|b| VkDescriptorBufferInfo { buffer: b.buffer, offset: 0, range: b.size }) .collect(); let writes: Vec = plan .bindings @@ -861,10 +826,7 @@ impl VulkanDevice { signalSemaphoreCount: 0, pSignalSemaphores: ptr::null(), }; - vk_check( - vkQueueSubmit(self.queue, 1, &submit, VK_NULL_HANDLE), - "vkQueueSubmit", - )?; + vk_check(vkQueueSubmit(self.queue, 1, &submit, VK_NULL_HANDLE), "vkQueueSubmit")?; vk_check(vkQueueWaitIdle(self.queue), "vkQueueWaitIdle")?; // No vkFreeCommandBuffers wired up — they live with the pool; // we destroy the pool at device-drop. @@ -1024,49 +986,47 @@ impl VulkanDevice { /// a resident weight read in a decode GEMV runs at device bandwidth instead /// of host bandwidth. Upload is one-time (staged); reads are device-local. /// The returned handle is freed with [`free_raw`] exactly like `alloc_raw`. - pub fn alloc_raw_device_local( - &self, - data: &[u8], - ) -> Result { + pub fn alloc_raw_device_local(&self, data: &[u8]) -> Result { let size = (data.len().max(4)) as u64; - let make_buffer = |usage: u32, props: u32| -> Result<(VkBuffer_, VkDeviceMemory), MetalTileError> { - let bci = VkBufferCreateInfo { - sType: VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO, - pNext: ptr::null(), - flags: 0, - size, - usage, - sharingMode: VK_SHARING_MODE_EXCLUSIVE, - queueFamilyIndexCount: 0, - pQueueFamilyIndices: ptr::null(), - }; - unsafe { - let mut buffer: VkBuffer_ = VK_NULL_HANDLE; - vk_check( - vkCreateBuffer(self.device, &bci, ptr::null(), &mut buffer), - "vkCreateBuffer(devlocal)", - )?; - let mut req = VkMemoryRequirements { size: 0, alignment: 0, memoryTypeBits: 0 }; - vkGetBufferMemoryRequirements(self.device, buffer, &mut req); - let mti = self.find_memory_type(req.memoryTypeBits, props)?; - let ai = VkMemoryAllocateInfo { - sType: VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO, + let make_buffer = + |usage: u32, props: u32| -> Result<(VkBuffer_, VkDeviceMemory), MetalTileError> { + let bci = VkBufferCreateInfo { + sType: VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO, pNext: ptr::null(), - allocationSize: req.size, - memoryTypeIndex: mti, + flags: 0, + size, + usage, + sharingMode: VK_SHARING_MODE_EXCLUSIVE, + queueFamilyIndexCount: 0, + pQueueFamilyIndices: ptr::null(), }; - let mut memory: VkDeviceMemory = VK_NULL_HANDLE; - vk_check( - vkAllocateMemory(self.device, &ai, ptr::null(), &mut memory), - "vkAllocateMemory(devlocal)", - )?; - vk_check( - vkBindBufferMemory(self.device, buffer, memory, 0), - "vkBindBufferMemory(devlocal)", - )?; - Ok((buffer, memory)) - } - }; + unsafe { + let mut buffer: VkBuffer_ = VK_NULL_HANDLE; + vk_check( + vkCreateBuffer(self.device, &bci, ptr::null(), &mut buffer), + "vkCreateBuffer(devlocal)", + )?; + let mut req = VkMemoryRequirements { size: 0, alignment: 0, memoryTypeBits: 0 }; + vkGetBufferMemoryRequirements(self.device, buffer, &mut req); + let mti = self.find_memory_type(req.memoryTypeBits, props)?; + let ai = VkMemoryAllocateInfo { + sType: VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO, + pNext: ptr::null(), + allocationSize: req.size, + memoryTypeIndex: mti, + }; + let mut memory: VkDeviceMemory = VK_NULL_HANDLE; + vk_check( + vkAllocateMemory(self.device, &ai, ptr::null(), &mut memory), + "vkAllocateMemory(devlocal)", + )?; + vk_check( + vkBindBufferMemory(self.device, buffer, memory, 0), + "vkBindBufferMemory(devlocal)", + )?; + Ok((buffer, memory)) + } + }; // Device-local destination (shader reads it fast). let (dst_buf, dst_mem) = make_buffer( @@ -1224,7 +1184,13 @@ impl VulkanDevice { }) .collect(); unsafe { - vkUpdateDescriptorSets(self.device, writes.len() as u32, writes.as_ptr(), 0, ptr::null()); + vkUpdateDescriptorSets( + self.device, + writes.len() as u32, + writes.as_ptr(), + 0, + ptr::null(), + ); } // Record + submit the command buffer. @@ -1504,9 +1470,7 @@ impl VulkanDevice { pub fn name(&self) -> &str { "vulkan-device" } /// Physical handle (for future direct queries via `vkGetPhysicalDeviceProperties`). - pub fn physical_device_handle(&self) -> VkPhysicalDevice { - self.physical_device - } + pub fn physical_device_handle(&self) -> VkPhysicalDevice { self.physical_device } /// Queue family index in use. pub fn queue_family(&self) -> u32 { self.queue_family_index } @@ -1534,21 +1498,14 @@ impl Drop for VulkanDevice { /// GLSL → SPIR-V via shaderc. The result is a byte-vector whose length is /// a multiple of 4 (SPIR-V is a stream of u32 words). -pub fn compile_glsl_to_spv( - glsl_src: &str, - file_name: &str, -) -> Result, MetalTileError> { - let csrc = - CString::new(glsl_src).map_err(|e| MetalTileError::Compilation(e.to_string()))?; - let cfile = - CString::new(file_name).map_err(|e| MetalTileError::Compilation(e.to_string()))?; +pub fn compile_glsl_to_spv(glsl_src: &str, file_name: &str) -> Result, MetalTileError> { + let csrc = CString::new(glsl_src).map_err(|e| MetalTileError::Compilation(e.to_string()))?; + let cfile = CString::new(file_name).map_err(|e| MetalTileError::Compilation(e.to_string()))?; let centry = CString::new("main").unwrap(); unsafe { let compiler = shaderc_compiler_initialize(); if compiler.is_null() { - return Err(MetalTileError::Compilation( - "shaderc_compiler_initialize failed".into(), - )); + return Err(MetalTileError::Compilation("shaderc_compiler_initialize failed".into())); } let opts = shaderc_compile_options_initialize(); shaderc_compile_options_set_target_env( diff --git a/crates/metaltile-runtime/src/lib.rs b/crates/metaltile-runtime/src/lib.rs index 33b1f1bf..9944e9ed 100644 --- a/crates/metaltile-runtime/src/lib.rs +++ b/crates/metaltile-runtime/src/lib.rs @@ -16,16 +16,18 @@ mod dispatch; pub mod error; pub use context::{Context, DispatchResult, DispatchSpec, ResidentBuffer}; -pub use device::gpu_family::GpuFamily; -pub use error::MetalTileError; - #[cfg(feature = "cuda")] pub use device::cuda::{CudaDevice, CudaFunction, CudaModule, DeviceBuffer}; - +pub use device::gpu_family::GpuFamily; #[cfg(feature = "hip")] pub use device::hip::{HipBuffer, HipDevice, HipKernel, HipModuleHandle}; - #[cfg(feature = "vulkan")] pub use device::vulkan::{ - BatchDispatch, VulkanBuffer, VulkanDevice, VulkanPipeline, VulkanRawBuffer, compile_glsl_to_spv, + BatchDispatch, + VulkanBuffer, + VulkanDevice, + VulkanPipeline, + VulkanRawBuffer, + compile_glsl_to_spv, }; +pub use error::MetalTileError; diff --git a/crates/metaltile-runtime/tests/cuda_smoke.rs b/crates/metaltile-runtime/tests/cuda_smoke.rs index 828a9c4b..9d9c7368 100644 --- a/crates/metaltile-runtime/tests/cuda_smoke.rs +++ b/crates/metaltile-runtime/tests/cuda_smoke.rs @@ -17,7 +17,15 @@ use metaltile_core::{ constexpr::ConstExpr, dtype::DType, ir::{ - BinOpKind, ConstExprDecl, IndexExpr, Kernel, Op, Param, ParamKind, ReduceKind, UnaryOpKind, + BinOpKind, + ConstExprDecl, + IndexExpr, + Kernel, + Op, + Param, + ParamKind, + ReduceKind, + UnaryOpKind, ValueId, }, shape::Shape, @@ -40,12 +48,22 @@ fn vector_add_ir() -> Kernel { k.body.push_op(Op::ProgramId { axis: 0 }, ValueId::new(0)); k.body.name_value(ValueId::new(0), "idx"); k.body.push_op( - Op::Load { src: "a".into(), indices: vec![IndexExpr::Value(ValueId::new(0))], mask: None, other: None }, + Op::Load { + src: "a".into(), + indices: vec![IndexExpr::Value(ValueId::new(0))], + mask: None, + other: None, + }, ValueId::new(1), ); k.body.name_value(ValueId::new(1), "x"); k.body.push_op( - Op::Load { src: "b".into(), indices: vec![IndexExpr::Value(ValueId::new(0))], mask: None, other: None }, + Op::Load { + src: "b".into(), + indices: vec![IndexExpr::Value(ValueId::new(0))], + mask: None, + other: None, + }, ValueId::new(2), ); k.body.name_value(ValueId::new(2), "y"); @@ -133,11 +151,7 @@ fn row_reduce_sum_ir() -> Kernel { is_output: true, kind: ParamKind::Tensor, }); - k.constexprs.push(ConstExprDecl { - name: ConstExpr::new("n"), - dtype: DType::U32, - value: None, - }); + k.constexprs.push(ConstExprDecl { name: ConstExpr::new("n"), dtype: DType::U32, value: None }); let row = ValueId::new(0); let nval = ValueId::new(1); let rs = ValueId::new(2); @@ -196,12 +210,37 @@ fn rms_norm_ir() -> Kernel { }); } k.constexprs.push(ConstExprDecl { name: ConstExpr::new("n"), dtype: DType::U32, value: None }); - k.constexprs.push(ConstExprDecl { name: ConstExpr::new("eps"), dtype: DType::F32, value: None }); + k.constexprs.push(ConstExprDecl { + name: ConstExpr::new("eps"), + dtype: DType::F32, + value: None, + }); let v = |i| ValueId::new(i); - let (row, tidv, nval, rs, col, x, sq, ssq, nf, msq, epsv, t, inv, w, xn, outv) = - (v(0), v(1), v(2), v(3), v(4), v(5), v(6), v(7), v(8), v(9), v(10), v(11), v(12), v(13), v(14), v(15)); - let ld = |src: &str, idx: Vec| Op::Load { src: src.into(), indices: idx, mask: None, other: None }; + let (row, tidv, nval, rs, col, x, sq, ssq, nf, msq, epsv, t, inv, w, xn, outv) = ( + v(0), + v(1), + v(2), + v(3), + v(4), + v(5), + v(6), + v(7), + v(8), + v(9), + v(10), + v(11), + v(12), + v(13), + v(14), + v(15), + ); + let ld = |src: &str, idx: Vec| Op::Load { + src: src.into(), + indices: idx, + mask: None, + other: None, + }; k.body.push_op(Op::ProgramId { axis: 0 }, row); k.body.name_value(row, "row"); @@ -234,9 +273,7 @@ fn rms_norm_ir() -> Kernel { k } -fn f32s_to_bytes(v: &[f32]) -> Vec { - v.iter().flat_map(|x| x.to_ne_bytes()).collect() -} +fn f32s_to_bytes(v: &[f32]) -> Vec { v.iter().flat_map(|x| x.to_ne_bytes()).collect() } fn bytes_to_f32s(b: &[u8]) -> Vec { b.chunks_exact(4).map(|c| f32::from_ne_bytes([c[0], c[1], c[2], c[3]])).collect() } @@ -362,9 +399,8 @@ fn row_reduce_sum_cuda_end_to_end() { const ROWS: usize = 128; const N: usize = 256; let inp: Vec = (0..ROWS * N).map(|i| ((i % 17) as f32 - 8.0) * 0.01).collect(); - let expected: Vec = (0..ROWS) - .map(|r| inp[r * N..(r + 1) * N].iter().sum::()) - .collect(); + let expected: Vec = + (0..ROWS).map(|r| inp[r * N..(r + 1) * N].iter().sum::()).collect(); let dinp = dev.upload(&f32s_to_bytes(&inp)).expect("upload inp"); let dout = dev.alloc(ROWS * 4).expect("alloc out"); diff --git a/crates/metaltile-runtime/tests/hip_smoke.rs b/crates/metaltile-runtime/tests/hip_smoke.rs index dfe06b3d..947803f4 100644 --- a/crates/metaltile-runtime/tests/hip_smoke.rs +++ b/crates/metaltile-runtime/tests/hip_smoke.rs @@ -16,8 +16,17 @@ use metaltile_core::{ constexpr::ConstExpr, dtype::DType, ir::{ - BinOpKind, ConstExprDecl, IndexExpr, Kernel, KernelMode, Op, Param, ParamKind, - ReduceKind, UnaryOpKind, ValueId, + BinOpKind, + ConstExprDecl, + IndexExpr, + Kernel, + KernelMode, + Op, + Param, + ParamKind, + ReduceKind, + UnaryOpKind, + ValueId, }, shape::Shape, }; @@ -39,12 +48,22 @@ fn vector_add_ir() -> Kernel { k.body.push_op(Op::ProgramId { axis: 0 }, ValueId::new(0)); k.body.name_value(ValueId::new(0), "idx"); k.body.push_op( - Op::Load { src: "a".into(), indices: vec![IndexExpr::Value(ValueId::new(0))], mask: None, other: None }, + Op::Load { + src: "a".into(), + indices: vec![IndexExpr::Value(ValueId::new(0))], + mask: None, + other: None, + }, ValueId::new(1), ); k.body.name_value(ValueId::new(1), "x"); k.body.push_op( - Op::Load { src: "b".into(), indices: vec![IndexExpr::Value(ValueId::new(0))], mask: None, other: None }, + Op::Load { + src: "b".into(), + indices: vec![IndexExpr::Value(ValueId::new(0))], + mask: None, + other: None, + }, ValueId::new(2), ); k.body.name_value(ValueId::new(2), "y"); @@ -80,11 +99,11 @@ fn hip_vector_add_f32_bit_exact() { Ok(None) => { eprintln!("hip_smoke: no HIP device — skipping"); return; - } + }, Err(e) => { eprintln!("hip_smoke: HIP init failed ({e:?}) — skipping"); return; - } + }, }; eprintln!( "hip_smoke: device='{}' gfx={} warp_size={}", @@ -99,9 +118,7 @@ fn hip_vector_add_f32_bit_exact() { let oracle: Vec = a.iter().zip(&b).map(|(x, y)| x + y).collect(); let mut bufs = BTreeMap::new(); - let to_bytes = |v: &[f32]| -> Vec { - v.iter().flat_map(|x| x.to_le_bytes()).collect() - }; + let to_bytes = |v: &[f32]| -> Vec { v.iter().flat_map(|x| x.to_le_bytes()).collect() }; bufs.insert("a".to_string(), to_bytes(&a)); bufs.insert("b".to_string(), to_bytes(&b)); bufs.insert("c".to_string(), vec![0u8; N * 4]); @@ -109,9 +126,7 @@ fn hip_vector_add_f32_bit_exact() { let block = 256u32; let grid = (N as u32).div_ceil(block); let k = vector_add_ir(); - let out = dev - .run_kernel(&k, &bufs, [grid, 1, 1], [block, 1, 1]) - .expect("hip run_kernel"); + let out = dev.run_kernel(&k, &bufs, [grid, 1, 1], [block, 1, 1]).expect("hip run_kernel"); let c_bytes = out.get("c").expect("output `c` present"); let c: Vec = c_bytes @@ -148,8 +163,13 @@ fn scale_add_exp_ir() -> Kernel { value: None, }); let (idx, x, sc, mul, y, sum, e) = ( - ValueId::new(0), ValueId::new(1), ValueId::new(2), - ValueId::new(3), ValueId::new(4), ValueId::new(5), ValueId::new(6), + ValueId::new(0), + ValueId::new(1), + ValueId::new(2), + ValueId::new(3), + ValueId::new(4), + ValueId::new(5), + ValueId::new(6), ); k.body.push_op(Op::ProgramId { axis: 0 }, idx); k.body.name_value(idx, "idx"); @@ -189,7 +209,7 @@ fn hip_scale_add_exp_f32_tight_tol() { _ => { eprintln!("hip_smoke: no HIP device — skipping"); return; - } + }, }; const N: usize = 4096; // Small magnitudes so exp doesn't saturate; tests `expf` precision. @@ -199,9 +219,7 @@ fn hip_scale_add_exp_f32_tight_tol() { let oracle: Vec = a.iter().zip(&b).map(|(x, y)| (x * scale + y).exp()).collect(); let mut bufs = BTreeMap::new(); - let to_bytes = |v: &[f32]| -> Vec { - v.iter().flat_map(|x| x.to_le_bytes()).collect() - }; + let to_bytes = |v: &[f32]| -> Vec { v.iter().flat_map(|x| x.to_le_bytes()).collect() }; bufs.insert("a".into(), to_bytes(&a)); bufs.insert("b".into(), to_bytes(&b)); bufs.insert("c".into(), vec![0u8; N * 4]); @@ -210,9 +228,7 @@ fn hip_scale_add_exp_f32_tight_tol() { let block = 256u32; let grid = (N as u32).div_ceil(block); let k = scale_add_exp_ir(); - let out = dev - .run_kernel(&k, &bufs, [grid, 1, 1], [block, 1, 1]) - .expect("hip run_kernel"); + let out = dev.run_kernel(&k, &bufs, [grid, 1, 1], [block, 1, 1]).expect("hip run_kernel"); let c: Vec = out["c"] .chunks_exact(4) @@ -235,21 +251,27 @@ fn row_reduce_sum_ir() -> Kernel { let mut k = Kernel::new("row_reduce_sum"); k.mode = KernelMode::Reduction; k.params.push(Param { - name: "inp".into(), dtype: DType::F32, shape: Shape::scalar(), - is_output: false, kind: ParamKind::Tensor, + name: "inp".into(), + dtype: DType::F32, + shape: Shape::scalar(), + is_output: false, + kind: ParamKind::Tensor, }); k.params.push(Param { - name: "out".into(), dtype: DType::F32, shape: Shape::scalar(), - is_output: true, kind: ParamKind::Tensor, - }); - k.constexprs.push(ConstExprDecl { - name: ConstExpr::new("n"), - dtype: DType::U32, - value: None, + name: "out".into(), + dtype: DType::F32, + shape: Shape::scalar(), + is_output: true, + kind: ParamKind::Tensor, }); + k.constexprs.push(ConstExprDecl { name: ConstExpr::new("n"), dtype: DType::U32, value: None }); let (row, nv, rs, re, acc, res) = ( - ValueId::new(0), ValueId::new(1), ValueId::new(2), - ValueId::new(3), ValueId::new(4), ValueId::new(5), + ValueId::new(0), + ValueId::new(1), + ValueId::new(2), + ValueId::new(3), + ValueId::new(4), + ValueId::new(5), ); k.body.push_op(Op::ProgramId { axis: 0 }, row); k.body.name_value(row, "row"); @@ -260,9 +282,15 @@ fn row_reduce_sum_ir() -> Kernel { k.body.name_value(re, "re"); k.body.push_op( Op::StrideReduce { - src: "inp".into(), offset: rs, stride: nv, end: re, - op: ReduceKind::Sum, dtype: DType::F32, - transform: None, secondary_src: None, secondary_base: None, + src: "inp".into(), + offset: rs, + stride: nv, + end: re, + op: ReduceKind::Sum, + dtype: DType::F32, + transform: None, + secondary_src: None, + secondary_base: None, }, acc, ); @@ -270,7 +298,10 @@ fn row_reduce_sum_ir() -> Kernel { k.body.push_op(Op::Reduce { value: acc, axis: 0, op: ReduceKind::Sum }, res); k.body.name_value(res, "result"); k.body.push_op_no_result(Op::Store { - dst: "out".into(), indices: vec![IndexExpr::Value(row)], value: res, mask: None, + dst: "out".into(), + indices: vec![IndexExpr::Value(row)], + value: res, + mask: None, }); k } @@ -282,21 +313,16 @@ fn hip_row_reduce_sum_f32() { _ => { eprintln!("hip_smoke: no HIP device — skipping"); return; - } + }, }; const ROWS: usize = 32; const COLS: usize = 4096; - let inp: Vec = (0..ROWS * COLS) - .map(|i| ((i as i32 % 257) as f32) * 0.001 - 0.1) - .collect(); - let oracle: Vec = (0..ROWS) - .map(|r| inp[r * COLS..(r + 1) * COLS].iter().sum::()) - .collect(); + let inp: Vec = (0..ROWS * COLS).map(|i| ((i as i32 % 257) as f32) * 0.001 - 0.1).collect(); + let oracle: Vec = + (0..ROWS).map(|r| inp[r * COLS..(r + 1) * COLS].iter().sum::()).collect(); let mut bufs = BTreeMap::new(); - let to_bytes = |v: &[f32]| -> Vec { - v.iter().flat_map(|x| x.to_le_bytes()).collect() - }; + let to_bytes = |v: &[f32]| -> Vec { v.iter().flat_map(|x| x.to_le_bytes()).collect() }; bufs.insert("inp".into(), to_bytes(&inp)); bufs.insert("out".into(), vec![0u8; ROWS * 4]); bufs.insert("n".into(), (COLS as u32).to_le_bytes().to_vec()); @@ -304,9 +330,7 @@ fn hip_row_reduce_sum_f32() { let block = 256u32; let grid = ROWS as u32; let k = row_reduce_sum_ir(); - let out = dev - .run_kernel(&k, &bufs, [grid, 1, 1], [block, 1, 1]) - .expect("hip run_kernel"); + let out = dev.run_kernel(&k, &bufs, [grid, 1, 1], [block, 1, 1]).expect("hip run_kernel"); let got: Vec = out["out"] .chunks_exact(4) diff --git a/crates/metaltile-runtime/tests/vulkan_smoke.rs b/crates/metaltile-runtime/tests/vulkan_smoke.rs index fef9c2e3..1d5185f6 100644 --- a/crates/metaltile-runtime/tests/vulkan_smoke.rs +++ b/crates/metaltile-runtime/tests/vulkan_smoke.rs @@ -16,8 +16,17 @@ use metaltile_core::{ constexpr::ConstExpr, dtype::DType, ir::{ - BinOpKind, ConstExprDecl, IndexExpr, Kernel, KernelMode, Op, Param, ParamKind, - ReduceKind, UnaryOpKind, ValueId, + BinOpKind, + ConstExprDecl, + IndexExpr, + Kernel, + KernelMode, + Op, + Param, + ParamKind, + ReduceKind, + UnaryOpKind, + ValueId, }, shape::Shape, }; @@ -37,12 +46,22 @@ fn vector_add_ir() -> Kernel { k.body.push_op(Op::ProgramId { axis: 0 }, ValueId::new(0)); k.body.name_value(ValueId::new(0), "idx"); k.body.push_op( - Op::Load { src: "a".into(), indices: vec![IndexExpr::Value(ValueId::new(0))], mask: None, other: None }, + Op::Load { + src: "a".into(), + indices: vec![IndexExpr::Value(ValueId::new(0))], + mask: None, + other: None, + }, ValueId::new(1), ); k.body.name_value(ValueId::new(1), "x"); k.body.push_op( - Op::Load { src: "b".into(), indices: vec![IndexExpr::Value(ValueId::new(0))], mask: None, other: None }, + Op::Load { + src: "b".into(), + indices: vec![IndexExpr::Value(ValueId::new(0))], + mask: None, + other: None, + }, ValueId::new(2), ); k.body.name_value(ValueId::new(2), "y"); @@ -74,8 +93,8 @@ fn shaderc_glsl_to_spv_for_vector_add() { // a Vulkan device — it exercises the GLSL→SPIR-V path in isolation. let g = GlslGenerator::new(); let glsl = g.generate(&vector_add_ir()).unwrap(); - let spv = metaltile_runtime::compile_glsl_to_spv(&glsl, "vector_add.comp") - .expect("shaderc compile"); + let spv = + metaltile_runtime::compile_glsl_to_spv(&glsl, "vector_add.comp").expect("shaderc compile"); // SPIR-V magic number `0x07230203` (little-endian). assert!(spv.len() >= 4); assert_eq!(&spv[..4], &[0x03, 0x02, 0x23, 0x07]); @@ -90,17 +109,13 @@ fn vulkan_vector_add_f32_bit_exact() { Ok(None) => { eprintln!("vulkan_smoke: no Vulkan device — skipping"); return; - } + }, Err(e) => { eprintln!("vulkan_smoke: Vulkan init failed ({e:?}) — skipping"); return; - } + }, }; - eprintln!( - "vulkan_smoke: device='{}' qfam={}", - dev.name(), - dev.queue_family() - ); + eprintln!("vulkan_smoke: device='{}' qfam={}", dev.name(), dev.queue_family()); const N: usize = 16 * 1024; let a: Vec = (0..N).map(|i| (i as f32) * 0.5).collect(); @@ -108,9 +123,7 @@ fn vulkan_vector_add_f32_bit_exact() { let oracle: Vec = a.iter().zip(&b).map(|(x, y)| x + y).collect(); let mut bufs = BTreeMap::new(); - let to_bytes = |v: &[f32]| -> Vec { - v.iter().flat_map(|x| x.to_le_bytes()).collect() - }; + let to_bytes = |v: &[f32]| -> Vec { v.iter().flat_map(|x| x.to_le_bytes()).collect() }; bufs.insert("a".to_string(), to_bytes(&a)); bufs.insert("b".to_string(), to_bytes(&b)); bufs.insert("c".to_string(), vec![0u8; N * 4]); @@ -118,9 +131,7 @@ fn vulkan_vector_add_f32_bit_exact() { let block = 256u32; let grid = (N as u32).div_ceil(block); let k = vector_add_ir(); - let out = dev - .run_kernel(&k, &bufs, [grid, 1, 1], [block, 1, 1]) - .expect("vulkan run_kernel"); + let out = dev.run_kernel(&k, &bufs, [grid, 1, 1], [block, 1, 1]).expect("vulkan run_kernel"); let c_bytes = out.get("c").expect("output `c` present"); let c: Vec = c_bytes @@ -157,8 +168,13 @@ fn scale_add_exp_ir() -> Kernel { value: None, }); let (idx, x, sc, mul, y, sum, e) = ( - ValueId::new(0), ValueId::new(1), ValueId::new(2), - ValueId::new(3), ValueId::new(4), ValueId::new(5), ValueId::new(6), + ValueId::new(0), + ValueId::new(1), + ValueId::new(2), + ValueId::new(3), + ValueId::new(4), + ValueId::new(5), + ValueId::new(6), ); k.body.push_op(Op::ProgramId { axis: 0 }, idx); k.body.name_value(idx, "idx"); @@ -196,7 +212,7 @@ fn vulkan_scale_add_exp_f32_tight_tol() { _ => { eprintln!("vulkan_smoke: no Vulkan device — skipping"); return; - } + }, }; const N: usize = 4096; @@ -206,9 +222,7 @@ fn vulkan_scale_add_exp_f32_tight_tol() { let oracle: Vec = a.iter().zip(&b).map(|(x, y)| (x * scale + y).exp()).collect(); let mut bufs = BTreeMap::new(); - let to_bytes = |v: &[f32]| -> Vec { - v.iter().flat_map(|x| x.to_le_bytes()).collect() - }; + let to_bytes = |v: &[f32]| -> Vec { v.iter().flat_map(|x| x.to_le_bytes()).collect() }; bufs.insert("a".into(), to_bytes(&a)); bufs.insert("b".into(), to_bytes(&b)); bufs.insert("c".into(), vec![0u8; N * 4]); @@ -217,9 +231,7 @@ fn vulkan_scale_add_exp_f32_tight_tol() { let block = 256u32; let grid = (N as u32).div_ceil(block); let k = scale_add_exp_ir(); - let out = dev - .run_kernel(&k, &bufs, [grid, 1, 1], [block, 1, 1]) - .expect("vulkan run_kernel"); + let out = dev.run_kernel(&k, &bufs, [grid, 1, 1], [block, 1, 1]).expect("vulkan run_kernel"); let c: Vec = out["c"] .chunks_exact(4) @@ -245,21 +257,27 @@ fn row_reduce_sum_ir() -> Kernel { let mut k = Kernel::new("row_reduce_sum"); k.mode = KernelMode::Reduction; k.params.push(Param { - name: "inp".into(), dtype: DType::F32, shape: Shape::scalar(), - is_output: false, kind: ParamKind::Tensor, + name: "inp".into(), + dtype: DType::F32, + shape: Shape::scalar(), + is_output: false, + kind: ParamKind::Tensor, }); k.params.push(Param { - name: "out".into(), dtype: DType::F32, shape: Shape::scalar(), - is_output: true, kind: ParamKind::Tensor, - }); - k.constexprs.push(ConstExprDecl { - name: ConstExpr::new("n"), - dtype: DType::U32, - value: None, + name: "out".into(), + dtype: DType::F32, + shape: Shape::scalar(), + is_output: true, + kind: ParamKind::Tensor, }); + k.constexprs.push(ConstExprDecl { name: ConstExpr::new("n"), dtype: DType::U32, value: None }); let (row, nv, rs, re, acc, res) = ( - ValueId::new(0), ValueId::new(1), ValueId::new(2), - ValueId::new(3), ValueId::new(4), ValueId::new(5), + ValueId::new(0), + ValueId::new(1), + ValueId::new(2), + ValueId::new(3), + ValueId::new(4), + ValueId::new(5), ); k.body.push_op(Op::ProgramId { axis: 0 }, row); k.body.name_value(row, "row"); @@ -270,9 +288,15 @@ fn row_reduce_sum_ir() -> Kernel { k.body.name_value(re, "re"); k.body.push_op( Op::StrideReduce { - src: "inp".into(), offset: rs, stride: nv, end: re, - op: ReduceKind::Sum, dtype: DType::F32, - transform: None, secondary_src: None, secondary_base: None, + src: "inp".into(), + offset: rs, + stride: nv, + end: re, + op: ReduceKind::Sum, + dtype: DType::F32, + transform: None, + secondary_src: None, + secondary_base: None, }, acc, ); @@ -280,7 +304,10 @@ fn row_reduce_sum_ir() -> Kernel { k.body.push_op(Op::Reduce { value: acc, axis: 0, op: ReduceKind::Sum }, res); k.body.name_value(res, "result"); k.body.push_op_no_result(Op::Store { - dst: "out".into(), indices: vec![IndexExpr::Value(row)], value: res, mask: None, + dst: "out".into(), + indices: vec![IndexExpr::Value(row)], + value: res, + mask: None, }); k } @@ -292,22 +319,17 @@ fn vulkan_row_reduce_sum_f32() { _ => { eprintln!("vulkan_smoke: no Vulkan device — skipping"); return; - } + }, }; const ROWS: usize = 32; const COLS: usize = 4096; - let inp: Vec = (0..ROWS * COLS) - .map(|i| ((i as i32 % 257) as f32) * 0.001 - 0.1) - .collect(); - let oracle: Vec = (0..ROWS) - .map(|r| inp[r * COLS..(r + 1) * COLS].iter().sum::()) - .collect(); + let inp: Vec = (0..ROWS * COLS).map(|i| ((i as i32 % 257) as f32) * 0.001 - 0.1).collect(); + let oracle: Vec = + (0..ROWS).map(|r| inp[r * COLS..(r + 1) * COLS].iter().sum::()).collect(); let mut bufs = BTreeMap::new(); - let to_bytes = |v: &[f32]| -> Vec { - v.iter().flat_map(|x| x.to_le_bytes()).collect() - }; + let to_bytes = |v: &[f32]| -> Vec { v.iter().flat_map(|x| x.to_le_bytes()).collect() }; bufs.insert("inp".into(), to_bytes(&inp)); bufs.insert("out".into(), vec![0u8; ROWS * 4]); bufs.insert("n".into(), (COLS as u32).to_le_bytes().to_vec()); @@ -315,9 +337,7 @@ fn vulkan_row_reduce_sum_f32() { let block = 256u32; let grid = ROWS as u32; let k = row_reduce_sum_ir(); - let out = dev - .run_kernel(&k, &bufs, [grid, 1, 1], [block, 1, 1]) - .expect("vulkan run_kernel"); + let out = dev.run_kernel(&k, &bufs, [grid, 1, 1], [block, 1, 1]).expect("vulkan run_kernel"); let got: Vec = out["out"] .chunks_exact(4) diff --git a/crates/metaltile-std/Cargo.toml b/crates/metaltile-std/Cargo.toml index a45bdc2a..179c6071 100644 --- a/crates/metaltile-std/Cargo.toml +++ b/crates/metaltile-std/Cargo.toml @@ -37,3 +37,6 @@ insta.workspace = true bytemuck.workspace = true # GPU dispatch in probe / GPU-correctness integration tests. metaltile-runtime.workspace = true +# In-module smoke tests in src/ffai/ssm.rs drive codegen + core IR directly. +metaltile-core.workspace = true +metaltile-codegen.workspace = true From c9d13eb52c26bf74518c3db2ecd70b2e080166e4 Mon Sep 17 00:00:00 2001 From: TheTom Date: Thu, 23 Jul 2026 06:51:54 -0500 Subject: [PATCH 4/5] chore: all-features clippy clean + typos config for domain identifiers The CI clippy gate runs with --all-features, which pulls the cuda and vulkan device modules into the lint set: allow attrs on the FFI-handle wrappers (not_unsafe_ptr_arg_deref on graph_launch/_batch, dead_code on the CUDA FFI constant surface, too_many_arguments on the grouped-CUTLASS entry), a dropped never-read init, and doc-comment list formatting. Typos config gains dout/ERRO (d-out signal name, log-level tag). No behavior change. --- _typos.toml | 3 +++ crates/metaltile-runtime/src/device/cuda/ffi.rs | 4 +++- crates/metaltile-runtime/src/device/cuda/mod.rs | 14 ++++++++++---- .../metaltile-runtime/src/device/vulkan/ffi.rs | 6 +++--- .../metaltile-runtime/src/device/vulkan/mod.rs | 4 ++-- crates/metaltile-runtime/tests/hip_smoke.rs | 4 ++-- crates/metaltile-std/tests/cuda_kernel_corpus.rs | 13 +++++++------ crates/metaltile-std/tests/hip_kernel_corpus.rs | 13 +++++++------ .../metaltile-std/tests/vulkan_kernel_corpus.rs | 16 ++++++++-------- 9 files changed, 45 insertions(+), 32 deletions(-) diff --git a/_typos.toml b/_typos.toml index 1589097e..5816de79 100644 --- a/_typos.toml +++ b/_typos.toml @@ -8,3 +8,6 @@ bb = "bb" nb = "nb" optin = "optin" OPTIN = "OPTIN" +dout = "dout" +ERRO = "ERRO" +erro = "erro" diff --git a/crates/metaltile-runtime/src/device/cuda/ffi.rs b/crates/metaltile-runtime/src/device/cuda/ffi.rs index e163f1fe..648ee51a 100644 --- a/crates/metaltile-runtime/src/device/cuda/ffi.rs +++ b/crates/metaltile-runtime/src/device/cuda/ffi.rs @@ -7,7 +7,9 @@ //! is under the NVIDIA Software License (not Apache) and the stack is alpha //! / Linux-only, so Phase 1 hand-rolls this small, stable surface for //! control. Re-evaluate `cuda-core` adoption when it leaves alpha. - +// FFI surface: bindings and constants are declared ahead of use per CUDA +// header groupings; unused ones stay for completeness. +#![allow(dead_code)] #![allow(non_camel_case_types)] use std::os::raw::{c_char, c_int, c_uint, c_void}; diff --git a/crates/metaltile-runtime/src/device/cuda/mod.rs b/crates/metaltile-runtime/src/device/cuda/mod.rs index c2fb87c6..1e51829e 100644 --- a/crates/metaltile-runtime/src/device/cuda/mod.rs +++ b/crates/metaltile-runtime/src/device/cuda/mod.rs @@ -226,7 +226,7 @@ fn size_bucket(len: usize) -> usize { // worst-case slack to ~12.5% while keeping the bucket count small. let pow2 = len.next_power_of_two(); let step = (pow2 / 8).max(MIN_BLOCK); - (len + step - 1) / step * step + len.div_ceil(step) * step } } @@ -509,6 +509,7 @@ impl CudaDevice { /// (HOST slices) describe each contiguous token group + its weight slab. /// `out[t,n] = Σ_k a[t,k]·w[eid][n,k]`. Beats cuBLAS per-expert on the skinny /// MoE shape. Errors if the runtime was built without CUTLASS. + #[allow(clippy::too_many_arguments)] pub fn moe_grouped_cutlass( &self, a: CUdeviceptr, @@ -1369,6 +1370,7 @@ impl CudaDevice { /// Allocate `len` bytes of device memory. pub fn alloc(&self, len: usize) -> Result, MetalTileError> { + #[allow(unused_assignments)] let mut ptr: CUdeviceptr = 0; if len == 0 { return Ok(DeviceBuffer { ptr: 0, len: 0, _dev: self }); @@ -1446,6 +1448,7 @@ impl CudaDevice { *self.pooled_bytes.lock().unwrap() -= bucket; return Ok(ptr); } + #[allow(unused_assignments)] let mut ptr: CUdeviceptr = 0; cu_check(unsafe { cuMemAlloc_v2(&mut ptr, bucket) }, "cuMemAlloc")?; return Ok(ptr); @@ -1454,6 +1457,7 @@ impl CudaDevice { if let Some(ptr) = self.pool.lock().unwrap().get_mut(&len).and_then(|v| v.pop()) { return Ok(ptr); } + #[allow(unused_assignments)] let mut ptr: CUdeviceptr = 0; cu_check(unsafe { cuMemAlloc_v2(&mut ptr, len) }, "cuMemAlloc")?; Ok(ptr) @@ -1725,6 +1729,7 @@ impl CudaDevice { /// Replay a captured decode token: ONE host launch replaces ~390 — no /// per-kernel enqueue, no inter-kernel host bubbles. Syncs the stream after. + #[allow(clippy::not_unsafe_ptr_arg_deref)] // FFI-handle wrapper, same idiom as the launch fns pub fn graph_launch(&self, exec: CUgraphExec) -> Result<(), MetalTileError> { cu_check(unsafe { cuGraphLaunch(exec, self.stream) }, "cuGraphLaunch")?; self.synchronize() @@ -1737,6 +1742,7 @@ impl CudaDevice { /// `graph_launch` (sync-per-token) incurs, giving the maximum throughput /// for the captured graph. Use only for throughput benchmarking — state /// (KV cache, SSM state) is overwritten sequentially and not meaningful. + #[allow(clippy::not_unsafe_ptr_arg_deref)] // FFI-handle wrapper, same idiom as the launch fns pub fn graph_launch_batch(&self, exec: CUgraphExec, n: usize) -> Result<(), MetalTileError> { for _ in 0..n { cu_check(unsafe { cuGraphLaunch(exec, self.stream) }, "cuGraphLaunch(batch)")?; @@ -1761,7 +1767,7 @@ impl CudaDevice { ) -> Result, MetalTileError> { // 1. IR → CUDA C++ → module. let cg = CudaGenerator::new(); - let src = cg.generate(kernel).map_err(|e| MetalTileError::Codegen(e))?; + let src = cg.generate(kernel).map_err(MetalTileError::Codegen)?; // MT_DUMP_CUDA_SRC=: write generated CUDA C++ for every kernel. if let Ok(dir) = std::env::var("MT_DUMP_CUDA_SRC") { let path = format!("{}/{}.cu", dir, kernel.name); @@ -1840,8 +1846,8 @@ impl CudaDevice { /// allocate + upload every param (by name from `buffers`), pack kernel /// args in signature order, launch over (`grid`×`block`), and read back /// the output params. The CUDA analog of `Context::dispatch_with_grid` - /// + `SingleDispatch`, used to run the registered kernel-test corpus on - /// CUDA. `buffers` must contain every param's bytes (inputs AND + /// plus `SingleDispatch`, used to run the registered kernel-test corpus + /// on CUDA. `buffers` must contain every param's bytes (inputs AND /// pre-sized outputs) plus each constexpr's name→LE-bytes. pub fn run_kernel( &self, diff --git a/crates/metaltile-runtime/src/device/vulkan/ffi.rs b/crates/metaltile-runtime/src/device/vulkan/ffi.rs index b072bb31..ec427fd8 100644 --- a/crates/metaltile-runtime/src/device/vulkan/ffi.rs +++ b/crates/metaltile-runtime/src/device/vulkan/ffi.rs @@ -12,9 +12,9 @@ //! ## Subset //! //! We implement only the path needed for: create instance → enumerate -//! physical device → create logical device + compute queue → create buffers -//! + memory + descriptor set → create compute pipeline from SPIR-V → -//! `vkCmdDispatch` → readback. Everything else (windowing, graphics, +//! physical device → create logical device and compute queue → create +//! buffers, memory, and descriptor set → create compute pipeline from +//! SPIR-V → `vkCmdDispatch` → readback. Everything else (windowing, graphics, //! sparse memory, ray-tracing) is intentionally out of scope. #![allow(non_camel_case_types, non_upper_case_globals, non_snake_case, dead_code)] diff --git a/crates/metaltile-runtime/src/device/vulkan/mod.rs b/crates/metaltile-runtime/src/device/vulkan/mod.rs index 54e7194d..82ab11c8 100644 --- a/crates/metaltile-runtime/src/device/vulkan/mod.rs +++ b/crates/metaltile-runtime/src/device/vulkan/mod.rs @@ -857,8 +857,8 @@ impl VulkanDevice { Ok(out) } - /// Query name (best-effort): we don't link the extra "PhysicalDeviceProperties" - /// FFI in Phase 1, so this is a placeholder. + // Query name (best-effort): we don't link the extra "PhysicalDeviceProperties" + // FFI in Phase 1, so this is a placeholder. // ────────────────────────────────────────────────────────────────── // Resident-buffer + cached-pipeline seam (mirrors CudaDevice::alloc_raw diff --git a/crates/metaltile-runtime/tests/hip_smoke.rs b/crates/metaltile-runtime/tests/hip_smoke.rs index 947803f4..3b0f236e 100644 --- a/crates/metaltile-runtime/tests/hip_smoke.rs +++ b/crates/metaltile-runtime/tests/hip_smoke.rs @@ -237,7 +237,7 @@ fn hip_scale_add_exp_f32_tight_tol() { let max_rel: f32 = c .iter() .zip(&oracle) - .map(|(g, w)| ((g - w).abs() / w.abs().max(1e-30))) + .map(|(g, w)| (g - w).abs() / w.abs().max(1e-30)) .fold(0.0f32, f32::max); eprintln!("hip_smoke: scale_add_exp max_rel = {max_rel:e}"); assert!(max_rel < 5e-7, "scale_add_exp tol broken: max_rel = {max_rel:e}"); @@ -341,7 +341,7 @@ fn hip_row_reduce_sum_f32() { let max_rel: f32 = got .iter() .zip(&oracle) - .map(|(g, w)| ((g - w).abs() / w.abs().max(1e-30))) + .map(|(g, w)| (g - w).abs() / w.abs().max(1e-30)) .fold(0.0f32, f32::max); eprintln!("hip_smoke: row_reduce_sum max_rel = {max_rel:e}"); assert!(max_rel < 1e-5, "row_reduce_sum tol broken: max_rel = {max_rel:e}"); diff --git a/crates/metaltile-std/tests/cuda_kernel_corpus.rs b/crates/metaltile-std/tests/cuda_kernel_corpus.rs index 3d8e27e0..74a32262 100644 --- a/crates/metaltile-std/tests/cuda_kernel_corpus.rs +++ b/crates/metaltile-std/tests/cuda_kernel_corpus.rs @@ -147,12 +147,13 @@ fn run_corpus_on_cuda() { let label = format!("{} [{dt}]", t.name()); // Debug: DUMP= prints its generated CUDA. - if let Ok(want) = std::env::var("DUMP") { - if t.name() == want && dt == DType::F32 { - use metaltile_codegen::{CodegenBackend, CudaGenerator}; - if let Ok(src) = CudaGenerator::new().generate(kernel) { - eprintln!("==== {} ====\n{src}\n==== end ====", t.name()); - } + if let Ok(want) = std::env::var("DUMP") + && t.name() == want + && dt == DType::F32 + { + use metaltile_codegen::{CodegenBackend, CudaGenerator}; + if let Ok(src) = CudaGenerator::new().generate(kernel) { + eprintln!("==== {} ====\n{src}\n==== end ====", t.name()); } } diff --git a/crates/metaltile-std/tests/hip_kernel_corpus.rs b/crates/metaltile-std/tests/hip_kernel_corpus.rs index b11de1de..de529857 100644 --- a/crates/metaltile-std/tests/hip_kernel_corpus.rs +++ b/crates/metaltile-std/tests/hip_kernel_corpus.rs @@ -146,12 +146,13 @@ fn run_corpus_on_hip() { let label = format!("{} [{dt}]", t.name()); // Debug: DUMP_HIP= prints the generated HIP source. - if let Ok(want) = std::env::var("DUMP_HIP") { - if t.name() == want && dt == DType::F32 { - use metaltile_codegen::{CodegenBackend, HipGenerator}; - if let Ok(src) = HipGenerator::new().generate(kernel) { - eprintln!("==== {} (HIP) ====\n{src}\n==== end ====", t.name()); - } + if let Ok(want) = std::env::var("DUMP_HIP") + && t.name() == want + && dt == DType::F32 + { + use metaltile_codegen::{CodegenBackend, HipGenerator}; + if let Ok(src) = HipGenerator::new().generate(kernel) { + eprintln!("==== {} (HIP) ====\n{src}\n==== end ====", t.name()); } } diff --git a/crates/metaltile-std/tests/vulkan_kernel_corpus.rs b/crates/metaltile-std/tests/vulkan_kernel_corpus.rs index 10a7eec4..2dd4afbd 100644 --- a/crates/metaltile-std/tests/vulkan_kernel_corpus.rs +++ b/crates/metaltile-std/tests/vulkan_kernel_corpus.rs @@ -145,14 +145,14 @@ fn run_corpus_on_vulkan() { let grid = setup.grid(); let label = format!("{} [{dt}]", t.name()); - if let Ok(want) = std::env::var("DUMP_VK") { - if t.name() == want && dt == DType::F32 { - use metaltile_codegen::{CodegenBackend, GlslGenerator}; - if let Ok(src) = - GlslGenerator::new().with_local_size_3d(grid.tpg).generate(kernel) - { - eprintln!("==== {} (Vulkan/GLSL) ====\n{src}\n==== end ====", t.name()); - } + if let Ok(want) = std::env::var("DUMP_VK") + && t.name() == want + && dt == DType::F32 + { + use metaltile_codegen::{CodegenBackend, GlslGenerator}; + if let Ok(src) = GlslGenerator::new().with_local_size_3d(grid.tpg).generate(kernel) + { + eprintln!("==== {} (Vulkan/GLSL) ====\n{src}\n==== end ====", t.name()); } } From 141d400d08518811dcf5cc02beb7d70656420df3 Mon Sep 17 00:00:00 2001 From: TheTom Date: Thu, 23 Jul 2026 06:59:03 -0500 Subject: [PATCH 5/5] chore: extend the CI typos config for kernel identifier vocabulary The typos gate reads .github/configs/typos-cli.toml explicitly; add the short GEMV accumulator/nibble identifiers (ba/bb/na/nb) and the dout/erro/optin word stems there and drop the stray root config the checker never consults. --- .github/configs/typos-cli.toml | 9 +++++++++ _typos.toml | 13 ------------- 2 files changed, 9 insertions(+), 13 deletions(-) delete mode 100644 _typos.toml diff --git a/.github/configs/typos-cli.toml b/.github/configs/typos-cli.toml index 83aa0357..db16202e 100644 --- a/.github/configs/typos-cli.toml +++ b/.github/configs/typos-cli.toml @@ -33,11 +33,20 @@ extend-ignore-re = [ ] [default.extend-identifiers] +# Short accumulator/nibble names in the hand-rolled GEMV kernels (a/b row pair). +ba = "ba" +bb = "bb" +na = "na" +nb = "nb" # Weird syntax, but this how you ignore corrections for certain words. [default.extend-words] # DSL / IR op names arange = "arange" +# d-out (derivative/output stream) identifier stem, ERRO log-level tag. +dout = "dout" +erro = "erro" +optin = "optin" Arange = "Arange" ARANGE = "ARANGE" # MLX kernel suffix for non-Apple-Silicon variants diff --git a/_typos.toml b/_typos.toml deleted file mode 100644 index 5816de79..00000000 --- a/_typos.toml +++ /dev/null @@ -1,13 +0,0 @@ -# Short accumulator and nibble identifiers in the hand-rolled GEMV kernels -# (ba/bb = per-block accumulators, na/nb = nibbles) read as typos to the -# checker; they are intentional. -[default.extend-identifiers] -ba = "ba" -na = "na" -bb = "bb" -nb = "nb" -optin = "optin" -OPTIN = "OPTIN" -dout = "dout" -ERRO = "ERRO" -erro = "erro"