From eed3b444330c371010414b423484d3ac5e46d2b8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 6 Sep 2026 19:30:00 +0200 Subject: [PATCH 01/33] fix(webgpu): declare SPIR-V entry points for passthrough shader modules --- CHANGELOG.md | 3 ++ crates/khal/src/backend/webgpu.rs | 66 +++++++++++++++++++++++++++++++ 2 files changed, 69 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4d22aa5..5e9f7be 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,9 @@ _Disclaimer: this changelog is updated using generative AI, but is still verifie ### Added - `atomic_sub_u32` in `khal-std`: atomic wrapping subtraction on a storage-buffer `u32`, returning the old value (`OpAtomicISub` on SPIR-V, `fetch_sub` elsewhere). +### Fixed +- WebGPU SPIR-V passthrough loading (`WebGpu::load_module_spirv_passthrough`) now declares the module's entry points (parsed from `OpEntryPoint`/`OpExecutionMode LocalSize`) in the wgpu passthrough descriptor. wgpu 30 validates compute pipelines against that list, so with it left empty every passthrough kernel failed to load on Vulkan with `Unable to find entry point`. + ## v0.3.0 ### Changed diff --git a/crates/khal/src/backend/webgpu.rs b/crates/khal/src/backend/webgpu.rs index 1fae606..f1facc2 100644 --- a/crates/khal/src/backend/webgpu.rs +++ b/crates/khal/src/backend/webgpu.rs @@ -282,6 +282,10 @@ impl WebGpu { } let spirv = wgpu::util::make_spirv_raw(spirv_bytes); + // wgpu doesn't reflect passthrough modules, so it only knows the entry points we + // declare here. `create_compute_pipeline` then rejects any entry point missing + // from that list, so declare every entry point the module actually contains. + let entry_points = spirv_entry_points(&spirv); // SAFETY: the caller has marked this shader as requiring passthrough loading, // meaning it uses SPIR-V features not supported by naga. The SPIR-V bytecode @@ -290,6 +294,7 @@ impl WebGpu { self.device .create_shader_module_passthrough(wgpu::ShaderModuleDescriptorPassthrough { spirv: Some(spirv), + entry_points: Cow::Owned(entry_points), ..Default::default() }) }; @@ -305,6 +310,67 @@ impl WebGpu { } } +/// Lists the entry points declared by a SPIR-V module, with the workgroup size of each. +/// +/// This reads the `OpEntryPoint` and `OpExecutionMode LocalSize` instructions directly +/// from the SPIR-V words, without going through naga (passthrough modules are exactly the +/// ones naga can't handle). Entry points without a `LocalSize` execution mode get a +/// workgroup size of `(0, 0, 0)`; wgpu only uses that field on Metal anyway. +fn spirv_entry_points(words: &[u32]) -> Vec> { + const HEADER_LEN: usize = 5; + const OP_ENTRY_POINT: u32 = 15; + const OP_EXECUTION_MODE: u32 = 16; + const EXECUTION_MODE_LOCAL_SIZE: u32 = 17; + + // (function id, name) + let mut entry_points: Vec<(u32, String)> = Vec::new(); + // (function id, workgroup size) + let mut local_sizes: Vec<(u32, (u32, u32, u32))> = Vec::new(); + + let mut i = HEADER_LEN; + while i < words.len() { + let word_count = (words[i] >> 16) as usize; + let opcode = words[i] & 0xffff; + if word_count == 0 { + break; // Malformed module; stop rather than loop forever. + } + let end = (i + word_count).min(words.len()); + + match opcode { + // OpEntryPoint: execution model, function id, literal name, interface ids... + OP_ENTRY_POINT if word_count >= 4 => { + let function_id = words[i + 2]; + let bytes: Vec = words[i + 3..end] + .iter() + .flat_map(|w| w.to_le_bytes()) + .collect(); + let len = bytes.iter().position(|&b| b == 0).unwrap_or(bytes.len()); + let name = String::from_utf8_lossy(&bytes[..len]).into_owned(); + entry_points.push((function_id, name)); + } + // OpExecutionMode: function id, mode, literals... (LocalSize has x, y, z) + OP_EXECUTION_MODE if word_count >= 6 && words[i + 2] == EXECUTION_MODE_LOCAL_SIZE => { + local_sizes.push((words[i + 1], (words[i + 3], words[i + 4], words[i + 5]))); + } + _ => {} + } + + i = end; + } + + entry_points + .into_iter() + .map(|(function_id, name)| wgpu::PassthroughShaderEntryPoint { + name: Cow::Owned(name), + workgroup_size: local_sizes + .iter() + .find(|(id, _)| *id == function_id) + .map(|(_, size)| *size) + .unwrap_or((0, 0, 0)), + }) + .collect() +} + /// Errors specific to the WebGPU backend. #[derive(thiserror::Error, Debug)] pub enum WebGpuBackendError { From 8a2cb1f34a343e9a383e03453fe9b64df798d010 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Sat, 6 Jun 2026 12:22:06 +0200 Subject: [PATCH 02/33] khal-derive + khal-std: cuda-oxide (native-CUDA) device backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a cuda-oxide code path so khal shader crates compile to PTX/cubin via the cuda-oxide Rust→PTX backend (alternative to Rust-CUDA's rustc_codegen_nvvm, which is wall-walled on LLVM 7), enabling `vortx --features cuda` and the nexus3d physics shaders to run on native CUDA. - khal-derive/spirv_bindgen: generate a `_cuda_entry_` cuda-oxide kernel (typed params, inlined body, builtins → khal_std::arch::cuda, `#[spirv(workgroup)]` → SharedArray/SmemBuf). CRITICAL ABI fix: the host pushes a (ptr, byte_len) pair for every storage binding, so sized-array (`&[T; N]`) and scalar (`&T`) storage params are received as slices and the original reference rebuilt in a prelude — otherwise they'd consume one fewer kernel arg than the host pushes and shift every later binding (illegal-address at runtime). Uniforms stay single-pointer. - khal-std: `cuda-oxide` feature — own `llvm.nvvm.*` sreg/barrier intrinsics, a libdevice-backed `Float`, `SmemBuf` shared-memory glue, generic `MaybeIndexUnchecked`; `cuda_std` made optional (it doesn't build under cuda-oxide's `-Zbuild-std`). - khal/backend/cuda: load a pre-linked CUBIN (ELF magic) in addition to PTX text — libdevice kernels (exp/sqrt/…) must arrive as a cubin since their PTX carries unresolved `__nv_*` externs. Optional `KHAL_CUDA_TRACE` launch trace. --- CUDA_OXIDE_AUTOPORT.md | 229 +++++++++++++++++++ crates/khal-derive/src/spirv_bindgen/cuda.rs | 220 +++++++++++++++++- crates/khal-derive/src/spirv_bindgen/mod.rs | 10 + crates/khal-std/Cargo.toml | 26 ++- crates/khal-std/src/arch/cuda.rs | 47 +++- crates/khal-std/src/lib.rs | 36 ++- crates/khal-std/src/num_traits.rs | 67 +++++- crates/khal/src/backend/cuda.rs | 19 +- crates/khal/src/backend/webgpu.rs | 67 +++++- 9 files changed, 698 insertions(+), 23 deletions(-) create mode 100644 CUDA_OXIDE_AUTOPORT.md diff --git a/CUDA_OXIDE_AUTOPORT.md b/CUDA_OXIDE_AUTOPORT.md new file mode 100644 index 0000000..2f4efea --- /dev/null +++ b/CUDA_OXIDE_AUTOPORT.md @@ -0,0 +1,229 @@ +# Running rust-gpu `#[spirv]` shaders on native CUDA via cuda-oxide + +> Status: working prototype (2026-06-05). `gpu_add` and `gemm_tiled` run **bit-exact** +> on the RTX 5090 (sm_120) through khal's CUDA backend, compiled by cuda-oxide. +> A `#[spirv_kernel]` proc-macro auto-lowers the *verbatim* shader source. +> +> **Update — the upstream `--target nvptx64` fix removes the per-shader shims.** +> The original prototype below compiled the device crate on the *host* target, so +> `cfg(target_arch="nvptx64")` was false — which forced a string of workarounds +> (glam `scalar-math`, `khal_std` barrier → `cuda_device::sync_threads`, +> `Vec4::new` for `Vec4::ZERO`). cuda-oxide now supports building a **device-only +> crate for the real `nvptx64-nvidia-cuda` target** (`cargo-oxide build --device`), +> so the cfg is correct and the *verbatim* `khal_std` `link_llvm_intrinsics` +> barrier lowers to a real `bar.sync`. See "Upstream fix" at the bottom. The +> shim story below is kept for the unified-crate (host-target) path. + +## The problem + +The whole stack (vortx tensors, nexus3d physics) is written once in **rust-gpu +`#[spirv]`** and compiled to SPIR-V for the **WebGPU/Vulkan** backend. khal also +has a **CUDA backend** (`khal/src/backend/cuda.rs`, cudarc-based) — but its +shader→PTX compiler, `cargo-cuda` → `rustc_codegen_nvvm`, needs **LLVM 7**, which +is effectively unbuildable on modern systems. So `--features cuda` could never +get PTX, and the CUDA backend sat unused. + +## The fix: cuda-oxide as the PTX compiler + +[cuda-oxide](https://github.com/NVlabs/cuda-oxide) is a modern **Rust → PTX** +backend (LLVM 21). It replaces the dead step: compile the kernels with cuda-oxide, +load the PTX/cubin through khal's existing `cudarc` backend. Same vortx ops, now +on native CUDA. + +Three layers had to line up: + +### 1. ABI — already compatible + +khal dispatches a kernel by **entry-point name** with positional `u64` params, +sorted by `(binding.space, binding.index)`: + +| `#[spirv(...)]` param | khal passes | cuda-oxide `#[kernel]` lowering | +|---|---|---| +| `storage_buffer` `&[f32]` / `&mut [f32]` | `(device_ptr, byte_len)` | `&[f32]` / `DisjointSlice` → `(ptr, len)` | +| `uniform` `&Shape` | `(device_ptr)` | `&Shape` → `(ptr)` | +| `push_constant` | raw bytes | scalar params | + +cuda-oxide's `#[kernel] fn(a: &[f32])` already lowers to `(ptr, i64)` — the same +shape khal sends. The entry-point name the host loads is +`Self::CUDA_ENTRY_POINT` = `format!("{fn}_cuda_entry_{hash}")` (from +`spirv_bindgen`); read it from a probe crate depending on `vortx-shaders`. + +### 2. libdevice — one patch + a cubin step + +Arithmetic kernels (GEMM) load as **PTX text** directly. But `exp`/ELU lower to +`__nv_expf` (libdevice), whose PTX has an unresolved extern that the driver JIT +rejects. Fix: + +- **Patch khal's loader** (`backend/cuda.rs::load_module_bytes`): detect a CUBIN + (ELF magic `\x7fELF`) → `cudarc::nvrtc::Ptx::from_binary(bytes)`, else PTX text + via `from_src`. `cuModuleLoadData` auto-detects the format. +- **Produce a self-contained cubin** from cuda-oxide's `.ll`: + `cuda_host::ltoir::build_cubin_from_ll(ll, "sm_120")` (libNVVM + libdevice + + nvJitLink). Needs CUDA 12.8+ `libnvvm`/`libnvJitLink` (pip wheels + `nvidia-cuda-nvcc-cu12` / `nvidia-nvjitlink-cu12`, pointed at via + `LIBNVVM_PATH` / `LIBNVJITLINK_PATH`) and the cuda-oxide nightly. + +A combined module (GEMM PTX + ELU libdevice) is built as **one cubin** — load it +as `shaders.ptx`; khal's patched loader takes it as a cubin. + +> The same wall hits TMA + FMA: TMA's nvvm intrinsic lowers only on the direct-PTX +> path, FMA (`mul_add`) only on the libdevice path — they're mutually exclusive in +> cuda-oxide today. + +### 3. Auto-port — compile the existing `#[spirv]` source unchanged + +The big realisation: **don't re-author kernels**. rust-gpu shaders are +`spirv-std`-flavored Rust; cuda-oxide speaks `cuda-device`. But: + +- `khal-std`, `glamx`, `spirv-std` **already compile for the nvptx target** under + cuda-oxide (`khal-std` is `#![cfg_attr(target_arch="nvptx64", …)]` by design). +- The shader **body** compiles unchanged (grid-stride loops, `khal_std` + `at_mut`/`read`, `Shape` math, `glamx` vectors). +- Only the **entry signature** needs a mechanical transform → the `#[spirv_kernel]` + proc-macro (`crates/spirv2cuda_macros`, or `~/spirv2cuda_macros` on the box). + +#### The `#[spirv_kernel]` macro + +Apply it where `#[spirv(compute(...))]` would go; the `#[spirv(...)]` *param* +attributes stay in place. It emits a `#[cuda_device::kernel]` fn: + +| input | output | +|---|---| +| `#[spirv(global/local_invocation_id\|workgroup_id)] id: UVec3` | dropped param; `let id = UVec3::new(thread::index_1d().get() as u32, 0, 0)` (etc.) in a prelude | +| `#[spirv(uniform)] x: &T` | kept (`&T` → ptr) | +| `#[spirv(storage_buffer)] x: &[f32]` | kept (`(ptr, len)`) | +| `#[spirv(storage_buffer)] x: &mut [f32]` | `DisjointSlice` param + `let mut x = MutBuf(x)` | +| body | emitted **unchanged** | + +The mutable mapping is the crux: cuda-oxide **cannot write through `&mut [T]` +params** (both `a[i]=v` — "Deref→Index assignment not yet implemented" — and +`get_unchecked_mut` stores silently no-op). Its working mutable type is +`DisjointSlice`. So the macro maps mutable storage buffers to `DisjointSlice` and +provides a tiny glue: + +```rust +pub struct MutBuf<'a>(pub DisjointSlice<'a, f32>); +impl<'a> MaybeIndexUnchecked for MutBuf<'a> { + fn at_mut(&mut self, i: usize) -> &mut f32 { unsafe { self.0.get_unchecked_mut(i) } } + fn write(&mut self, i: usize, v: f32) { unsafe { *self.0.get_unchecked_mut(i) = v; } } + // gpu_add never reads its output buffer; DisjointSlice has no &self read. + fn read(&self, _i: usize) -> f32 { 0.0 } + fn at(&self, _i: usize) -> &f32 { static Z: f32 = 0.0; &Z } // (&T->&mut T cast is now a hard error) +} +``` + +So the body's verbatim `*a.at_mut(ia) += b.read(ib)` lands on a persisting store. + +#### Proof + +`examples/autoport2` is the **byte-for-byte vortx `gpu_add`** — signature, body, +and `#[spirv(...)]` attributes all in place — under `#[spirv_kernel]`. It compiles +to PTX, loads through khal's CUDA backend, and returns `a + b` with **0.0 error** +over 120 elements. No wrapper written by hand. + +## Results (RTX 5090) + +- **Ops** (bit-exact vs WebGPU/CPU): `gpu_add` (0.0), `gemm_tiled` (0.0, boundary- + tested), `elu` (5.96e-8, libdevice cubin), and vortx's own `OpAssign::launch` + host op on cuda-oxide PTX (0.0). +- **GEMM ladder** (1024³): naive 7.4 → register-blocked 10.5 → +FMA 16.3 → + +double-buffer **22.8 TFLOPS**. +- **Policy forward** (actor+critic), WebGPU vs cuda-oxide: **5.5–13×** faster + (0.34 vs 4.5 ms/step at N=8192; 14.4 vs ~1.1 TFLOPS) — WebGPU is dispatch-bound. +- **Hybrid full iteration** (physics stays WebGPU, policy + PPO update on + cuda-oxide): **1.44–1.52×** end-to-end — the update (3.2×) is the real lever; + the forward is ~2% of the physics-bound rollout. + +## What's left + +The mechanism is proven; the rest is **bounded extension**, not research: + +1. **Macro coverage** for cases `gpu_add` doesn't hit: + `#[spirv(workgroup)] &mut [T; N]` → `SharedArray`; push-constant cfg branches; + more thread builtins; non-`f32` element types; **atomics** (the physics solver). +2. **A real cuda-oxide fix** (optional): support write-back through `&mut [T]` + params (the mature `rustc_codegen_nvvm` does) — would remove the `MutBuf` + mapping. Also unsupported: `core::slice::from_raw_parts_mut`. +3. **Zero-copy hybrid transfer**: Vulkan↔CUDA external-memory interop + (`VK_KHR_external_memory` ↔ `cudaImportExternalMemory`) via wgpu-hal raw + handles, to erase the per-step obs copy. +4. **Full port**: apply the macro across all vortx + nexus shaders → a fully + native-CUDA pipeline (no cross-API), at which point the hybrid's transfer wart + disappears. + +## Upstream fix — compile device crates for the real `nvptx64` target + +The prototype above worked around a single root cause: `cargo-oxide` ran plain +`cargo build` (host target), so `cfg(target_arch="nvptx64")` was **false** and +every nvptx-gated code path (glam scalar math, `khal_std`'s `bar.sync` barrier, +indexing) silently took its host fallback. The proper fix is to compile the +**device crate** for `nvptx64-nvidia-cuda`. Four small cuda-oxide changes make +this a first-class mode — and then the *verbatim* shader source compiles with +**no shims at all**: + +1. **cargo-oxide** — new flag `cargo-oxide build --device [--arch sm_120] ` + adds `--target nvptx64-nvidia-cuda -Z build-std=core` (the nvptx target ships + no precompiled `core`) and points `CUDA_OXIDE_PTX_DIR` at the crate dir. +2. **cuda-macros** — gate the host-side `cuda_host::CudaKernel` marker behind + `#[cfg(not(target_arch="nvptx64"))]`, so the same `#[kernel]` source compiles + both as a unified host+device crate *and* as a pure device crate (where + `cuda_host` isn't in scope). +3. **mir-importer** — honor `#[link_name="llvm.*"]` externs (the + `link_llvm_intrinsics` pattern rust-gpu / `khal_std` use). rustc resolves such + a foreign item to a symbol whose mangled name *is* the link name, so the + importer surfaces `llvm.nvvm.barrier0` as the dispatch name and routes it to + the same `Barrier0Op` as `cuda_device::sync_threads`. General fix: any crate's + `llvm.nvvm.*` intrinsic externs now lower instead of failing PTX verification + with "Symbol … not found". +4. **rustc-codegen-cuda** — skip the host-object *embed* step for an nvptx + `llvm_target` (a device-only crate has no host object to embed into); the + standalone `.ptx` is the deliverable. + +**Proof** (`examples/nvptx_probe`): a `#![no_std]` device-only crate whose kernel +uses the *verbatim* `khal_std` barrier and bakes `cfg!(target_arch="nvptx64")` +into its output. `cargo-oxide build --device --arch sm_120 nvptx_probe` → +PTX with `bar.sync` present and the cfg marker = nvptx-true → runs on the 5090 +through khal/cuda with **0.0 error**. The per-shader workarounds are gone. + +**Bridge to the full vortx stack — closed.** `khal_std`'s nvptx path used to +depend on `cuda_std` (Rust-CUDA), which transitively pulls `half`, and `half` +fails under `-Z build-std` for nvptx. `cuda_std` was used only for thread-index +sregs and `GpuFloat` math. khal_std now has a **`cuda-oxide` feature** (alongside +the default `rust-cuda`) that: + +- declares the 12 `%tid`/`%ctaid`/`%ntid`/`%nctaid` reads as khal_std's own + `llvm.nvvm.read.ptx.sreg.*` externs (lowered by the importer's link-name + honoring — the same mechanism as the barrier; the importer also gained + `llvm.nvvm.read.ptx.sreg.*` aliases on its existing sreg dispatch arms); +- routes `Float` to `num_traits` (libm) instead of `cuda_std::float::GpuFloat`; +- drops the `cuda_std` dependency entirely, so the device crate builds under + `-Z build-std`. + +(It also needs `#![feature(asm_experimental_arch)]` on nvptx for the warp-shuffle +`asm!` in `sync::subgroup`.) Build shader crates with +`khal-std = { default-features = false, features = ["cuda-oxide", …] }`. + +**Capstone proof** (`examples/gemm_device`): the *verbatim* vortx `gemm_tiled` +as a `#![no_std]` device-only crate — verbatim `khal_std` barrier (no +`sync_threads` shim), glamx `nostd-libm` (no `scalar-math` override), +khal_std `cuda-oxide` feature. `cargo-oxide build --device --arch sm_120 +gemm_device` → `bar.sync:2`, two `.shared` arrays, `.entry gemm_tiled` → runs on +the 5090 (100×50 @ 50×70, boundary-tested) with **0.0 error**. Every per-shader +workaround is gone. Remaining: `elu` (libdevice / `__nv_expf` cubin path) and the +rest of the vortx kernels, then nexus physics (atomics). + +## Where things live (on the baguette box) + +- cuda-oxide toolchain: `~/cuda-oxide-src`, LLVM 21 at `~/llvm21`, nightly + `2026-04-03`, `cargo-oxide` on PATH. Env per build: `CUDA_OXIDE_LLC`, + `LIBCLANG_PATH`, `LD_LIBRARY_PATH`, `CUDA_HOME`, and for libdevice + `LIBNVVM_PATH` / `LIBNVJITLINK_PATH` (12.9 wheels in `~/nvvm-wheel`, + `~/nvjit-wheel`). +- proc-macro: `~/spirv2cuda_macros`. cubin tool: `~/make_cubin`. +- examples: `~/cuda-oxide-src/crates/rustc-codegen-cuda/examples/{autoport2, + khal_gemm,khal_elu,khal_fwd,khal_opassign,…}`. khal host demos: `~/khal_*_demo`, + `~/fwd_bench`, `~/upd_bench`, `~/vortx_cuda_host`. +- patches (local, baguette only): `khal/src/backend/cuda.rs` (cubin loader), + `khal-builder/src/lib.rs` (`build_ptx` injects cuda-oxide PTX, skipping + cargo-cuda). Backups in `/tmp/*.bak`. WebGPU training is unaffected (these only + fire under `--features cuda`). diff --git a/crates/khal-derive/src/spirv_bindgen/cuda.rs b/crates/khal-derive/src/spirv_bindgen/cuda.rs index 93b64e2..5063f1e 100644 --- a/crates/khal-derive/src/spirv_bindgen/cuda.rs +++ b/crates/khal-derive/src/spirv_bindgen/cuda.rs @@ -215,7 +215,7 @@ pub(super) fn generate_cuda_entry_block( // Only generate if there are bindings (otherwise it's not a real kernel) if !sorted_bindings.is_empty() || !sorted_push_constants.is_empty() { quote! { - #[cfg(target_arch = "nvptx64")] + #[cfg(all(target_arch = "nvptx64", not(feature = "cuda-oxide")))] #[khal_std::cuda_std::kernel] pub unsafe fn #cuda_entry_ident( #(#cuda_params),* @@ -228,3 +228,221 @@ pub(super) fn generate_cuda_entry_block( quote! {} } } + +/// Generate a CUDA entry point for the **cuda-oxide** PTX backend. +/// +/// cuda-oxide kernels take TYPED params (`&T` uniforms, `&[T]`/`&mut [T]` +/// storage buffers — khal passes each as `(ptr, len)`), so this inlines the +/// shader body directly: bindings become kernel params, builtins are computed +/// from `khal_std::arch::cuda`, and `#[spirv(workgroup)] &mut [f32; N]` becomes +/// a function-local `SharedArray` wrapped in `SmemBuf` (shared memory). The +/// body runs verbatim against those bindings. +pub(super) fn generate_cuda_oxide_entry_block( + func: &syn::ItemFn, + original_params: &[OriginalParam], + bindings: &[ShaderBinding], + cuda_entry_ident: &syn::Ident, +) -> proc_macro2::TokenStream { + let all_params: Vec<&OriginalParam> = original_params + .iter() + .filter(|p| { + p.cfg_attrs.is_empty() + || p.cfg_attrs.iter().any(|a| { + let s = quote::quote!(#a).to_string(); + s.contains("not") && s.contains("push_constants") + }) + }) + .collect(); + + // Push-constant kernels have no uniform fallback here -> skip. + if all_params + .iter() + .any(|p| matches!(p.kind, OriginalParamKind::PushConstant)) + { + return quote! {}; + } + + // Kernel params = bindings, in (descriptor_set, binding) order, typed. + let mut sorted_bindings: Vec<&OriginalParam> = all_params + .iter() + .filter(|p| matches!(p.kind, OriginalParamKind::Binding { .. })) + .copied() + .collect(); + sorted_bindings.sort_by_key(|p| { + bindings + .iter() + .find(|b| b.name == p.name) + .map(|b| (b.descriptor_set, b.binding)) + .unwrap_or((0, 0)) + }); + // The host (khal CudaDispatch) pushes a `(ptr, byte_len)` pair for EVERY + // storage binding. A `&[T]`/`&mut [T]` typed param already lowers to that + // 2-arg ABI, but a sized-array ref `&[T; N]`/`&mut [T; N]` lowers to a single + // thin pointer — so the kernel would consume one fewer arg than the host + // pushes, shifting every later binding's pointer by a slot (a following + // buffer pointer becomes a length value → illegal-address on first access). + // Receive array storage bindings as SLICES here (matching the 2-arg ABI) and + // reconstruct the `&[T; N]`/`&mut [T; N]` in a prelude below. + let mut array_reconstructions: Vec = Vec::new(); + let kernel_params: Vec = sorted_bindings + .iter() + .map(|p| { + let name = &p.name; + let ty = &p.ty; + // Uniform bindings are pushed by the host as a single pointer, so keep + // them typed (thin pointer ABI). Storage bindings are pushed as a + // `(ptr, byte_len)` pair: slice refs already match that 2-arg ABI, but + // sized-array refs (`&[T; N]`) and scalar refs (`&T`) lower to a single + // thin pointer and would under-consume the host's args, shifting every + // later binding. Receive those as slices and reconstruct the original + // reference in a prelude so the ABI stays aligned. + let is_uniform = matches!( + p.kind, + OriginalParamKind::Binding { is_uniform: true, .. } + ); + if is_uniform { + return quote! { #name: #ty }; + } + if let syn::Type::Reference(r) = ty { + match &*r.elem { + // slice storage already lowers to (ptr, len) + syn::Type::Slice(_) => quote! { #name: #ty }, + // sized array storage: receive as slice, rebuild `[T; N]` ref + syn::Type::Array(arr) => { + let elem_ty = &*arr.elem; + let len_expr = &arr.len; + let buf_name = syn::Ident::new(&format!("{}_buf", name), name.span()); + if r.mutability.is_some() { + array_reconstructions.push(quote! { + let #name: &mut [#elem_ty; #len_expr] = unsafe { + &mut *(#buf_name.as_mut_ptr() as *mut [#elem_ty; #len_expr]) + }; + }); + quote! { #buf_name: &mut [#elem_ty] } + } else { + array_reconstructions.push(quote! { + let #name: &[#elem_ty; #len_expr] = unsafe { + &*(#buf_name.as_ptr() as *const [#elem_ty; #len_expr]) + }; + }); + quote! { #buf_name: &[#elem_ty] } + } + } + // scalar storage (`&T` / `&mut T`): receive as slice, rebuild ref + scalar => { + let elem_ty = scalar; + let buf_name = syn::Ident::new(&format!("{}_buf", name), name.span()); + if r.mutability.is_some() { + array_reconstructions.push(quote! { + let #name: &mut #elem_ty = + unsafe { &mut *#buf_name.as_mut_ptr() }; + }); + quote! { #buf_name: &mut [#elem_ty] } + } else { + array_reconstructions.push(quote! { + let #name: &#elem_ty = unsafe { &*#buf_name.as_ptr() }; + }); + quote! { #buf_name: &[#elem_ty] } + } + } + } + } else { + quote! { #name: #ty } + } + }) + .collect(); + + // Preludes for builtins and workgroup shared memory, in original order. + let mut preludes: Vec = Vec::new(); + for p in &all_params { + let name = &p.name; + match &p.kind { + OriginalParamKind::Builtin(kind) => { + let expr = match kind { + BuiltinKind::GlobalInvocationId => { + quote! { khal_std::arch::cuda::global_invocation_id() } + } + BuiltinKind::LocalInvocationId => { + quote! { khal_std::arch::cuda::local_invocation_id() } + } + BuiltinKind::WorkgroupId => { + quote! { khal_std::arch::cuda::workgroup_id() } + } + BuiltinKind::NumWorkgroups => { + quote! { khal_std::arch::cuda::num_workgroups() } + } + BuiltinKind::LocalInvocationIndex => { + quote! {{ + let __tid = khal_std::arch::cuda::thread_idx(); + __tid.x + }} + } + _ => quote! { Default::default() }, + }; + preludes.push(quote! { let #name = #expr; }); + } + OriginalParamKind::Workgroup => { + // `#[spirv(workgroup)] x: &mut [T; N]` -> a function-local + // `SharedArray` (shared address space) viewed as a real + // `&mut [T; N]`; and `x: &mut T` (scalar broadcast slot) -> + // `SharedArray` viewed as `&mut T`. A real reference (not a + // wrapper) so the body can index it directly AND pass it to + // helpers that take `&mut [T; N]` / `&mut T`. + let static_name = + syn::Ident::new(&format!("__smem_{}", name), name.span()); + match &p.ty { + syn::Type::Reference(r) => match &*r.elem { + // `&mut [f32; N]` shared tile -> `&mut SmemBuf` (direct + // `SharedArray` indexing -> st.shared). Body indexes it via + // MaybeIndexUnchecked and can reborrow it across helpers + // generic over `&mut impl MaybeIndexUnchecked`. + syn::Type::Array(arr) => { + let elem_ty = &*arr.elem; + let len_expr = &arr.len; + let sb_name = + syn::Ident::new(&format!("__smembuf_{}", name), name.span()); + preludes.push(quote! { + static mut #static_name: + khal_std::cuda_device::SharedArray<#elem_ty, { #len_expr }> = + khal_std::cuda_device::SharedArray::UNINIT; + let mut #sb_name = khal_std::cuda_oxide_glue::SmemBuf( + unsafe { &mut *core::ptr::addr_of_mut!(#static_name) } + ); + let #name = &mut #sb_name; + }); + } + // `&mut T` scalar broadcast slot -> `SharedArray` -> + // `&mut T` (real shared reference; works through helpers). + scalar => { + let elem_ty = scalar; + let sref_name = + syn::Ident::new(&format!("__sref_{}", name), name.span()); + preludes.push(quote! { + static mut #static_name: + khal_std::cuda_device::SharedArray<#elem_ty, 1> = + khal_std::cuda_device::SharedArray::UNINIT; + let #sref_name: &mut khal_std::cuda_device::SharedArray<#elem_ty, 1> = + unsafe { &mut *core::ptr::addr_of_mut!(#static_name) }; + let #name: &mut #elem_ty = &mut #sref_name[0]; + }); + } + }, + _ => {} + } + } + _ => {} + } + } + + let body = &func.block; + + quote! { + #[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] + #[khal_std::cuda_device::kernel] + pub fn #cuda_entry_ident(#(#kernel_params),*) { + #(#preludes)* + #(#array_reconstructions)* + #body + } + } +} diff --git a/crates/khal-derive/src/spirv_bindgen/mod.rs b/crates/khal-derive/src/spirv_bindgen/mod.rs index 48a0854..0b850cf 100644 --- a/crates/khal-derive/src/spirv_bindgen/mod.rs +++ b/crates/khal-derive/src/spirv_bindgen/mod.rs @@ -660,6 +660,15 @@ pub(crate) fn spirv_bindgen(attr: TokenStream, item: TokenStream) -> TokenStream &cuda_entry_ident, ); + // cuda-oxide PTX backend variant (typed params; selected by the + // `cuda-oxide` feature on the shader crate). + let cuda_oxide_entry_block = cuda::generate_cuda_oxide_entry_block( + &func, + &original_params, + &bindings, + &cuda_entry_ident, + ); + let cuda_entry_name_str = &cuda_entry_name; // Extract doc attributes from the original function to propagate to the generated struct. @@ -820,6 +829,7 @@ pub(crate) fn spirv_bindgen(attr: TokenStream, item: TokenStream) -> TokenStream #args_struct_def #wrapper_def #cuda_entry_block + #cuda_oxide_entry_block }; output.into() diff --git a/crates/khal-std/Cargo.toml b/crates/khal-std/Cargo.toml index eae18d9..19074b1 100644 --- a/crates/khal-std/Cargo.toml +++ b/crates/khal-std/Cargo.toml @@ -8,6 +8,7 @@ edition = "2024" license = "MIT OR Apache-2.0" [features] +default = ["rust-cuda"] cpu = ["corosensei"] cpu-parallel = ["cpu", "rayon"] cuda = [] @@ -24,6 +25,21 @@ cuda = [] # for deploying to the web. unsafe_remove_boundchecks = [] +# --- nvptx64 device backend selection --- +# `rust-cuda` (default): compile device code with Rust-CUDA's rustc_codegen_nvvm, +# pulling `cuda_std` for thread/grid intrinsics and `GpuFloat` math. +# `cuda-oxide`: compile device code with the cuda-oxide PTX backend instead. +# cuda_std does not build under cuda-oxide's `-Zbuild-std`, so this path uses +# khal_std's own `llvm.nvvm.*` intrinsic declarations (lowered by cuda-oxide) +# and a libm-backed `Float`. Build shader crates with +# `default-features = false, features = ["cuda-oxide", ...]`. +rust-cuda = ["dep:cuda_std"] +# NOTE: the cuda-oxide backend (device-only nvptx build) also needs a +# `cuda-device` dependency, but that crate lives in the cuda-oxide tree and is +# only present on the build box; it is added there. This manifest keeps the +# feature host-buildable everywhere else. +cuda-oxide = ["dep:num-traits"] + [dependencies] glamx = { version = "0.3", default-features = false, features = ["nostd-libm", "bytemuck", "u32", "i32", "f64"] } rayon = { version = "1", optional = true } @@ -38,5 +54,11 @@ workspace = true spirv-std = "0.10.0-alpha.1" [target.'cfg(target_arch = "nvptx64")'.dependencies] -# Fixes the UVec3::element_product bug -#cuda_std = { git = "https://github.com/sebcrozet/Rust-CUDA.git", branch = "glam-up" } +# Fixes the UVec3::element_product bug. Optional: only the `rust-cuda` backend +# (rustc_codegen_nvvm) uses it; the `cuda-oxide` backend declares its own +# intrinsics instead (cuda_std does not build under `-Zbuild-std`). +cuda_std = { git = "https://github.com/sebcrozet/Rust-CUDA.git", branch = "glam-up", optional = true } +# `Float` for the cuda-oxide backend (libm software math, no_std). +num-traits = { version = "0.2", default-features = false, features = ["libm"], optional = true } +# (cuda-oxide device build only, added on the build box:) +# cuda-device = { path = ".../cuda-oxide/crates/cuda-device", optional = true } diff --git a/crates/khal-std/src/arch/cuda.rs b/crates/khal-std/src/arch/cuda.rs index 251b611..6bf9070 100644 --- a/crates/khal-std/src/arch/cuda.rs +++ b/crates/khal-std/src/arch/cuda.rs @@ -1,8 +1,47 @@ //! CUDA intrinsic wrappers for thread and block indexing. +//! +//! Two backends provide the underlying `%tid`/`%ctaid`/`%ntid`/`%nctaid` +//! special-register reads: +//! - `rust-cuda` (default): `cuda_std::thread` (Rust-CUDA / rustc_codegen_nvvm). +//! - `cuda-oxide`: khal_std's own `llvm.nvvm.read.ptx.sreg.*` externs, lowered +//! by the cuda-oxide PTX backend. `cuda_std` does not build under cuda-oxide. -use cuda_std::thread; use glamx::UVec3; +#[cfg(not(feature = "cuda-oxide"))] +use cuda_std::thread; + +/// khal_std's own nvptx special-register reads, used by the `cuda-oxide` +/// backend in place of `cuda_std::thread`. Each lowers to a single +/// `mov.u32 %r, %sreg` in PTX. +#[cfg(feature = "cuda-oxide")] +mod thread { + macro_rules! sreg { + ($name:ident, $intr:literal) => { + #[inline(always)] + pub fn $name() -> u32 { + unsafe extern "C" { + #[link_name = $intr] + fn f() -> u32; + } + unsafe { f() } + } + }; + } + sreg!(thread_idx_x, "llvm.nvvm.read.ptx.sreg.tid.x"); + sreg!(thread_idx_y, "llvm.nvvm.read.ptx.sreg.tid.y"); + sreg!(thread_idx_z, "llvm.nvvm.read.ptx.sreg.tid.z"); + sreg!(block_idx_x, "llvm.nvvm.read.ptx.sreg.ctaid.x"); + sreg!(block_idx_y, "llvm.nvvm.read.ptx.sreg.ctaid.y"); + sreg!(block_idx_z, "llvm.nvvm.read.ptx.sreg.ctaid.z"); + sreg!(block_dim_x, "llvm.nvvm.read.ptx.sreg.ntid.x"); + sreg!(block_dim_y, "llvm.nvvm.read.ptx.sreg.ntid.y"); + sreg!(block_dim_z, "llvm.nvvm.read.ptx.sreg.ntid.z"); + sreg!(grid_dim_x, "llvm.nvvm.read.ptx.sreg.nctaid.x"); + sreg!(grid_dim_y, "llvm.nvvm.read.ptx.sreg.nctaid.y"); + sreg!(grid_dim_z, "llvm.nvvm.read.ptx.sreg.nctaid.z"); +} + /// Returns the thread index within the current block as a `UVec3`. #[inline(always)] pub fn thread_idx() -> UVec3 { @@ -55,8 +94,8 @@ pub fn workgroup_id() -> UVec3 { #[inline(always)] pub fn num_workgroups() -> UVec3 { UVec3::new( - cuda_std::thread::grid_dim_x(), - cuda_std::thread::grid_dim_y(), - cuda_std::thread::grid_dim_z(), + thread::grid_dim_x(), + thread::grid_dim_y(), + thread::grid_dim_z(), ) } diff --git a/crates/khal-std/src/lib.rs b/crates/khal-std/src/lib.rs index dd39cbd..12c9b19 100644 --- a/crates/khal-std/src/lib.rs +++ b/crates/khal-std/src/lib.rs @@ -5,6 +5,13 @@ #![cfg_attr(any(target_arch = "spirv", target_arch = "nvptx64"), no_std)] #![cfg_attr(target_arch = "nvptx64", feature(link_llvm_intrinsics))] +// nvptx inline asm (warp shuffle in `sync::subgroup`) is unstable; the +// cuda-oxide backend compiles for the real nvptx64 target so it needs the gate. +#![cfg_attr(target_arch = "nvptx64", feature(asm_experimental_arch))] +// `Float` for the cuda-oxide backend calls `core::intrinsics::{expf32,..}` +// (lowered to libdevice `__nv_*` by cuda-oxide) instead of software libm. +#![cfg_attr(all(target_arch = "nvptx64", feature = "cuda-oxide"), feature(core_intrinsics))] +#![cfg_attr(all(target_arch = "nvptx64", feature = "cuda-oxide"), allow(internal_features))] /// Architecture-specific runtime support (CPU coroutines, CUDA intrinsics). pub mod arch; @@ -32,8 +39,35 @@ pub use build_script::*; /// Re-export of the `glamx` math library. pub use glamx; -#[cfg(target_arch = "nvptx64")] +#[cfg(all(target_arch = "nvptx64", not(feature = "cuda-oxide")))] pub use cuda_std; +/// cuda-oxide device crate (re-exported so the `spirv_bindgen`-generated CUDA +/// entry points can name `khal_std::cuda_device::kernel` without the shader +/// crate depending on cuda-device directly). +#[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] +pub use cuda_device; + +/// Glue for the cuda-oxide `spirv_bindgen` entry transform: wraps a +/// `&'static mut SharedArray` so a `#[spirv(workgroup)] &mut [T; N]` parameter +/// body (which calls `MaybeIndexUnchecked::read/write/at_mut`) lands on shared +/// memory. +#[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] +pub mod cuda_oxide_glue { + use crate::index::MaybeIndexUnchecked; + use cuda_device::SharedArray; + + pub struct SmemBuf(pub &'static mut SharedArray); + impl MaybeIndexUnchecked for SmemBuf { + #[inline(always)] + fn read(&self, i: usize) -> f32 { self.0[i] } + #[inline(always)] + fn write(&mut self, i: usize, v: f32) { self.0[i] = v; } + #[inline(always)] + fn at_mut(&mut self, i: usize) -> &mut f32 { &mut self.0[i] } + #[inline(always)] + fn at(&self, i: usize) -> &f32 { &self.0[i] } + } +} #[cfg(not(any(target_arch = "spirv", target_arch = "nvptx64")))] pub use std::println; diff --git a/crates/khal-std/src/num_traits.rs b/crates/khal-std/src/num_traits.rs index d7f52c2..6fba8c6 100644 --- a/crates/khal-std/src/num_traits.rs +++ b/crates/khal-std/src/num_traits.rs @@ -1,9 +1,64 @@ -// Re-export num_traits at crate root. -// On non-nvptx64: from spirv_std. On nvptx64: a compat module that re-exports -// cuda_std::GpuFloat as Float so `use crate::num_traits::Float` works. -/// On nvptx64, re-export `cuda_std::GpuFloat` as `Float` so that -/// `use crate::num_traits::Float` works the same as `spirv_std::num_traits::Float`. -#[cfg(target_arch = "nvptx64")] +// Re-export a `Float` trait at crate root, backend-dependent. +// - non-nvptx (spirv/cpu): `spirv_std::num_traits::Float`. +// - nvptx + rust-cuda: `cuda_std::float::GpuFloat` (hardware intrinsics). +// - nvptx + cuda-oxide: a local trait whose methods call `core::intrinsics` +// transcendentals; cuda-oxide lowers those to libdevice `__nv_*`. (libm +// software math does not resolve under cuda-oxide's PTX backend.) +#[cfg(all(target_arch = "nvptx64", not(feature = "cuda-oxide")))] pub use cuda_std::float::GpuFloat as Float; #[cfg(not(target_arch = "nvptx64"))] pub use spirv_std::num_traits::Float; + +#[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] +pub use cuda_oxide_float::Float; + +#[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] +mod cuda_oxide_float { + /// Minimal `Float` for the cuda-oxide PTX backend. Methods lower to + /// libdevice (`__nv_expf`, `__nv_sqrtf`, …) via `core::intrinsics`. + pub trait Float: Copy { + fn exp(self) -> Self; + fn ln(self) -> Self; + fn sqrt(self) -> Self; + fn powf(self, n: Self) -> Self; + fn abs(self) -> Self; + fn floor(self) -> Self; + fn ceil(self) -> Self; + fn max(self, other: Self) -> Self; + fn min(self, other: Self) -> Self; + fn atan(self) -> Self; + } + + macro_rules! float_impl { + ($ty:ty, $exp:ident, $ln:ident, $sqrt:ident, $pow:ident, $floor:ident, $ceil:ident, $nvatan:ident) => { + impl Float for $ty { + #[inline(always)] + fn exp(self) -> $ty { unsafe { core::intrinsics::$exp(self) } } + #[inline(always)] + fn ln(self) -> $ty { unsafe { core::intrinsics::$ln(self) } } + #[inline(always)] + fn sqrt(self) -> $ty { unsafe { core::intrinsics::$sqrt(self) } } + #[inline(always)] + fn powf(self, n: $ty) -> $ty { unsafe { core::intrinsics::$pow(self, n) } } + #[inline(always)] + fn floor(self) -> $ty { unsafe { core::intrinsics::$floor(self) } } + #[inline(always)] + fn ceil(self) -> $ty { unsafe { core::intrinsics::$ceil(self) } } + // abs/max/min via plain ops (no stable intrinsic names here). + #[inline(always)] + fn abs(self) -> $ty { if self < 0.0 { -self } else { self } } + #[inline(always)] + fn max(self, other: $ty) -> $ty { if self >= other { self } else { other } } + #[inline(always)] + fn min(self, other: $ty) -> $ty { if self <= other { self } else { other } } + #[inline(always)] + fn atan(self) -> $ty { + unsafe extern "C" { fn $nvatan(x: $ty) -> $ty; } + unsafe { $nvatan(self) } + } + } + }; + } + float_impl!(f32, expf32, logf32, sqrtf32, powf32, floorf32, ceilf32, __nv_atanf); + float_impl!(f64, expf64, logf64, sqrtf64, powf64, floorf64, ceilf64, __nv_atan); +} diff --git a/crates/khal/src/backend/cuda.rs b/crates/khal/src/backend/cuda.rs index 2cba404..6c23581 100644 --- a/crates/khal/src/backend/cuda.rs +++ b/crates/khal/src/backend/cuda.rs @@ -121,6 +121,7 @@ pub struct CudaModule { #[derive(Clone)] pub struct CudaFunction { pub(crate) func: driver::CudaFunction, + pub(crate) name: String, } // ── Encoder / Pass ───────────────────────────────────────────────────── @@ -319,7 +320,7 @@ impl Backend for Cuda { _push_constant_size: u32, ) -> Result { let func = module.inner.load_function(entry_point)?; - Ok(CudaFunction { func }) + Ok(CudaFunction { func, name: entry_point.to_string() }) } fn load_function_with_layouts( @@ -607,9 +608,25 @@ impl<'a> Dispatch<'a, Cuda> for CudaDispatch<'a> { shared_mem_bytes: 0, }; + let trace = std::env::var_os("KHAL_CUDA_TRACE").is_some(); + if trace { + eprintln!( + "[khal-cuda launch] {} nargs={} grid={:?} block={:?}", + self.function.name, param_values.len(), grid_dim, block_dim + ); + } unsafe { builder.launch(cfg)?; } + if trace { + match self.stream.synchronize() { + Ok(()) => eprintln!("[khal-cuda ok ] {}", self.function.name), + Err(e) => { + eprintln!("[khal-cuda FAIL ] {} -> {:?}", self.function.name, e); + return Err(e.into()); + } + } + } Ok(()) } diff --git a/crates/khal/src/backend/webgpu.rs b/crates/khal/src/backend/webgpu.rs index f1facc2..fae5ba1 100644 --- a/crates/khal/src/backend/webgpu.rs +++ b/crates/khal/src/backend/webgpu.rs @@ -8,8 +8,9 @@ use bytemuck::{AnyBitPattern, NoUninit}; use regex::Regex; use smallvec::SmallVec; use std::borrow::Cow; +use std::collections::HashMap; use std::ops::RangeBounds; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use wgpu::util::{BufferInitDescriptor, DeviceExt}; use wgpu::wgt::CommandEncoderDescriptor; use wgpu::{ @@ -110,6 +111,14 @@ pub struct WebGpu { spirv_passthrough_enabled: bool, /// Whether `TIMESTAMP_QUERY` is supported on this device. timestamp_supported: bool, + /// Per-size pool of `MAP_READ | COPY_DST` staging buffers, reused across + /// `slow_read_buffer` calls so we don't pay a fresh `create_buffer` / + /// `MTLBuffer.makeBuffer` (and its inverse on drop) every call. Keyed by + /// buffer size in bytes. `Vec` per size so concurrent callers with the + /// same payload size each get an exclusive buffer. `Arc` so cloned + /// `WebGpu` handles (which already share the same underlying device / + /// queue via wgpu's internal Arcs) also share the pool. + read_staging_pool: Arc>>>, } impl WebGpu { @@ -191,6 +200,7 @@ impl WebGpu { hacks: vec![], spirv_passthrough_enabled, timestamp_supported, + read_staging_pool: Arc::new(Mutex::new(HashMap::new())), }) } @@ -213,9 +223,43 @@ impl WebGpu { hacks: vec![], spirv_passthrough_enabled, timestamp_supported, + read_staging_pool: Arc::new(Mutex::new(HashMap::new())), } } + /// Pop a `MAP_READ | COPY_DST` staging buffer of exactly `bytes_len` bytes + /// from the readback pool, allocating one if the pool is empty for that + /// size. Pairs with [`release_read_staging`]. Used by [`slow_read_buffer`] + /// to avoid allocating a fresh GPU buffer (and freeing it) on every call. + fn acquire_read_staging(&self, bytes_len: u64) -> Buffer { + if let Some(buf) = self + .read_staging_pool + .lock() + .expect("read_staging_pool poisoned") + .get_mut(&bytes_len) + .and_then(|v| v.pop()) + { + return buf; + } + self.device.create_buffer(&BufferDescriptor { + label: Some("khal-read-staging"), + size: bytes_len, + usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }) + } + + /// Return a staging buffer to the pool so the next `slow_read_buffer` of + /// the same size can reuse it. Caller must have already `unmap`ed it. + fn release_read_staging(&self, bytes_len: u64, buf: Buffer) { + self.read_staging_pool + .lock() + .expect("read_staging_pool poisoned") + .entry(bytes_len) + .or_default() + .push(buf); + } + /// Adds a regex-based text replacement to apply to WGSL source before compilation. pub fn append_hack(&mut self, regex: Regex, replace_pattern: String) { self.hacks.push((regex, replace_pattern)); @@ -694,18 +738,25 @@ impl Backend for WebGpu { buffer: &Self::Buffer, out: &mut [T], ) -> Result<(), Self::Error> { - // Create staging buffer. - let bytes_len = buffer.size() as usize; - let staging = - self.uninit_buffer::(bytes_len, BufferUsages::MAP_READ | BufferUsages::COPY_DST)?; + // Acquire (or allocate) a pooled staging buffer of the exact size. + // For a steady-state RL loop reading the same buffer every step this + // pool warms up after the first call and we stop paying the per-step + // `create_buffer` / drop round-trip (especially expensive on Metal + // where it's a fresh MTLBuffer allocation each time). + let bytes_len = buffer.size(); + let staging = self.acquire_read_staging(bytes_len); let mut encoder = self.begin_encoding(); encoder .encoder - .copy_buffer_to_buffer(buffer, 0, &staging, 0, bytes_len as u64); + .copy_buffer_to_buffer(buffer, 0, &staging, 0, bytes_len); self.submit(encoder)?; - // Read the buffer. - self.read_buffer(&staging, out).await + // `read_buffer` does the map_async + memcpy + unmap. After this call + // the staging buffer is back in `MAP_READ | COPY_DST` ready state and + // safe to return to the pool. + let result = self.read_buffer(&staging, out).await; + self.release_read_staging(bytes_len, staging); + result } } From d94f6abf81d29772556f93c1db8f7e5e97a5624a Mon Sep 17 00:00:00 2001 From: haixuantao Date: Sun, 7 Jun 2026 19:12:25 +0200 Subject: [PATCH 03/33] khal/cuda: print entry name on load_function failure (debug aid for cubin/host hash mismatches) --- crates/khal/src/backend/cuda.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/crates/khal/src/backend/cuda.rs b/crates/khal/src/backend/cuda.rs index 6c23581..c9d7a52 100644 --- a/crates/khal/src/backend/cuda.rs +++ b/crates/khal/src/backend/cuda.rs @@ -319,7 +319,13 @@ impl Backend for Cuda { entry_point: &str, _push_constant_size: u32, ) -> Result { - let func = module.inner.load_function(entry_point)?; + let func = match module.inner.load_function(entry_point) { + Ok(f) => f, + Err(e) => { + eprintln!("[khal-cuda load_function FAIL] {} -> {:?}", entry_point, e); + return Err(e.into()); + } + }; Ok(CudaFunction { func, name: entry_point.to_string() }) } From 7ac12dbff07c3befb24546f9f984457bc3e7a403 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Sun, 7 Jun 2026 22:06:22 +0200 Subject: [PATCH 04/33] khal-std: cap glam <0.33 for the spirv-std shader build glam 0.33 dropped UVec3/UVec4 under default-features=false, which breaks spirv-std 0.10.0-alpha.1's compile (112 errors). The `cargo gpu build` sub-invocation re-resolves without the workspace lock and would otherwise pick the newest 0.33.x via spirv-std's open `glam >=0.30.8`. Pin to 0.32.1, the newest pre-0.33 version that still works (mirrors the pin in vortx-shaders). --- crates/khal-std/Cargo.toml | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/crates/khal-std/Cargo.toml b/crates/khal-std/Cargo.toml index 19074b1..989b41c 100644 --- a/crates/khal-std/Cargo.toml +++ b/crates/khal-std/Cargo.toml @@ -52,6 +52,13 @@ workspace = true [target.'cfg(not(target_arch = "nvptx64"))'.dependencies] spirv-std = "0.10.0-alpha.1" +# Cap glam < 0.33 for the shader build: spirv-std 0.10.0-alpha.1 declares glam +# ">=0.30.8" (open-ended), but glam 0.33 dropped `UVec4`/`UVec3` under +# default-features=false, which breaks spirv-std's compile (112 errors). The +# `cargo gpu build` sub-invocation re-resolves without the workspace lock and +# would otherwise pick the newest 0.33.x. 0.32.1 is the newest pre-0.33 version +# that still works. (Mirrors the same pin in vortx-shaders.) +glam = { version = "=0.32.1", default-features = false } [target.'cfg(target_arch = "nvptx64")'.dependencies] # Fixes the UVec3::element_product bug. Optional: only the `rust-cuda` backend From 1029ddc14750e6d13efb56e186eb7ee1c5532a4d Mon Sep 17 00:00:00 2001 From: haixuantao Date: Mon, 8 Jun 2026 14:30:08 +0200 Subject: [PATCH 05/33] fix(cuda): push element count, not byte length, for slice kernel args MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit cuda-oxide lowers a `&[T]` kernel parameter to a `(ptr, len)` pair where `len` is an ELEMENT count, but the khal CUDA backend was pushing the buffer's BYTE length. So a kernel calling `slice.len()` got `byte_len` (4x too large for `&[u32]`) and read out of bounds — fatal in the batched broad-phase radix sort (`gpu_init_sort_dispatch` derives the workgroup grid from `num_keys_arr.len()`, producing a garbage indirect dispatch and an illegal memory access) and in the LBVH traversal (`*_len.len()`). Fix: in the three CUDA `write_arg` arms (`GpuBuffer`/`GpuBufferSlice`/ `GpuBufferSliceMut`, all generic over `T`), push `byte_len / size_of::()`. Sized arrays (`&[T; N]`), scalars (`&T`) and uniforms reconstruct via the pointer and ignore this value, so they are unaffected; only true `&[T]` slices change — to correct. Latent in vortx (its kernels bound loops by `Shape` uniforms, never `slice.len()`) and in single-env physics; surfaced once N>1 batched physics ran on native CUDA. --- crates/khal/src/backend/any_backend.rs | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/crates/khal/src/backend/any_backend.rs b/crates/khal/src/backend/any_backend.rs index bf720f6..688cb39 100644 --- a/crates/khal/src/backend/any_backend.rs +++ b/crates/khal/src/backend/any_backend.rs @@ -1633,7 +1633,11 @@ impl<'b, T: DeviceValue> crate::ShaderArgs<'b> for GpuBuffer { } #[cfg(feature = "cuda")] (GpuBuffer::Cuda(buffer), GpuDispatch::Cuda(dispatch)) => { - dispatch.set_arg(binding, buffer.device_ptr_raw(), buffer.byte_len()); + // cuda-oxide `&[T]` slice ABI wants an ELEMENT count, but byte_len is bytes; + // push element count so kernel `slice.len()` is correct (off-by-size_of + // otherwise -> OOB reads, e.g. gpu_init_sort_dispatch / lbvh). Arrays, + // scalars and uniforms ignore this value, so they are unaffected. + dispatch.set_arg(binding, buffer.device_ptr_raw(), buffer.byte_len() / std::mem::size_of::() as u64); Ok(()) } #[cfg(feature = "metal")] @@ -1665,7 +1669,7 @@ impl<'b, T: DeviceValue> crate::ShaderArgs<'b> for GpuBufferSlice<'_, T> { } #[cfg(feature = "cuda")] (GpuBufferSlice::Cuda(slice), GpuDispatch::Cuda(dispatch)) => { - dispatch.set_arg(binding, slice.offset_ptr(), slice.byte_len); + dispatch.set_arg(binding, slice.offset_ptr(), slice.byte_len / std::mem::size_of::() as u64); Ok(()) } #[cfg(feature = "metal")] @@ -1702,7 +1706,7 @@ impl<'b, T: DeviceValue> crate::ShaderArgs<'b> for GpuBufferSliceMut<'_, T> { } #[cfg(feature = "cuda")] (GpuBufferSliceMut::Cuda(slice), GpuDispatch::Cuda(dispatch)) => { - dispatch.set_arg(binding, slice.offset_ptr(), slice.byte_len); + dispatch.set_arg(binding, slice.offset_ptr(), slice.byte_len / std::mem::size_of::() as u64); Ok(()) } #[cfg(feature = "metal")] From 79c0ae2cc6410e5604dbc4c77b47144eb10598fb Mon Sep 17 00:00:00 2001 From: haixuantao Date: Tue, 7 Jul 2026 17:56:05 +0200 Subject: [PATCH 06/33] khal-std: drop obsolete glam<0.33 cap (defer to glamx 0.3, matches upstream); khal-builder: KHAL_SKIP_SPIRV skip for cuda-only builds --- crates/khal-builder/src/lib.rs | 14 +++++++++++++- crates/khal-std/Cargo.toml | 7 ------- 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/crates/khal-builder/src/lib.rs b/crates/khal-builder/src/lib.rs index edf8c82..571cd58 100644 --- a/crates/khal-builder/src/lib.rs +++ b/crates/khal-builder/src/lib.rs @@ -93,7 +93,18 @@ impl KhalBuilder { self.setup_change_detection(); - if self.build_spirv { + // Consumers embed this directory at compile time (e.g. vortx's + // `include_dir!("$OUT_DIR/shaders-spirv")`), so it must exist even when the + // SPIR-V build is skipped — otherwise the host crate fails to compile. + std::fs::create_dir_all(output_dir) + .unwrap_or_else(|e| panic!("failed to create shader output dir {output_dir:?}: {e}")); + + // `KHAL_SKIP_SPIRV=1` skips the cargo-gpu SPIR-V compile entirely. This lets + // a CUDA-only build (shaders come from cuda-oxide PTX cubins) avoid the + // cargo-gpu toolchain as a build prerequisite. The WebGPU backend then has + // no embedded shaders and fails hard at runtime — intended on CUDA boxes. + let skip_spirv = std::env::var_os("KHAL_SKIP_SPIRV").is_some(); + if self.build_spirv && !skip_spirv { self.build_spirv(output_dir); } @@ -129,6 +140,7 @@ impl KhalBuilder { println!("cargo:rerun-if-env-changed=CARGO_FEATURE_PUSH_CONSTANTS"); // TODO: currently unused println!("cargo:rerun-if-env-changed=CARGO_FEATURE_CUDA"); + println!("cargo:rerun-if-env-changed=KHAL_SKIP_SPIRV"); } fn build_spirv(&self, output_dir: impl AsRef) { diff --git a/crates/khal-std/Cargo.toml b/crates/khal-std/Cargo.toml index 989b41c..19074b1 100644 --- a/crates/khal-std/Cargo.toml +++ b/crates/khal-std/Cargo.toml @@ -52,13 +52,6 @@ workspace = true [target.'cfg(not(target_arch = "nvptx64"))'.dependencies] spirv-std = "0.10.0-alpha.1" -# Cap glam < 0.33 for the shader build: spirv-std 0.10.0-alpha.1 declares glam -# ">=0.30.8" (open-ended), but glam 0.33 dropped `UVec4`/`UVec3` under -# default-features=false, which breaks spirv-std's compile (112 errors). The -# `cargo gpu build` sub-invocation re-resolves without the workspace lock and -# would otherwise pick the newest 0.33.x. 0.32.1 is the newest pre-0.33 version -# that still works. (Mirrors the same pin in vortx-shaders.) -glam = { version = "=0.32.1", default-features = false } [target.'cfg(target_arch = "nvptx64")'.dependencies] # Fixes the UVec3::element_product bug. Optional: only the `rust-cuda` backend From 784d8d8829498491d4d04b6bebf2c44c0db49411 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Wed, 8 Jul 2026 08:57:21 +0200 Subject: [PATCH 07/33] Port CUDA runtime APIs from fork branches: compute_capability, graph capture (begin/end_capture, CapturedGraph), KHAL_CUDA_PROFILE kernel profiling, prebuilt-PTX escape hatch (CUDA_OXIDE_SHADERS_PTX_*) --- crates/khal-builder/src/lib.rs | 23 ++++++ crates/khal/src/backend/any_backend.rs | 84 ++++++++++++++++++- crates/khal/src/backend/cuda.rs | 107 ++++++++++++++++++++++++- 3 files changed, 210 insertions(+), 4 deletions(-) diff --git a/crates/khal-builder/src/lib.rs b/crates/khal-builder/src/lib.rs index 571cd58..52ffa58 100644 --- a/crates/khal-builder/src/lib.rs +++ b/crates/khal-builder/src/lib.rs @@ -175,9 +175,32 @@ impl KhalBuilder { } /// Compiles the shader crate to PTX for the CUDA backend. + /// + /// `CUDA_OXIDE_SHADERS_PTX_` (upper-snake, e.g. + /// `CUDA_OXIDE_SHADERS_PTX_VORTX_SHADERS`) points at a prebuilt PTX/cubin; + /// when set it is embedded directly as `shaders.ptx` and the `cargo cuda` + /// (cuda-oxide) toolchain is not required. #[cfg(feature = "cuda")] fn build_ptx(&self, output_dir: impl AsRef) { let output_dir = output_dir.as_ref(); + + let crate_name = self + .shader_crate + .file_name() + .and_then(|n| n.to_str()) + .unwrap_or_default() + .to_uppercase() + .replace('-', "_"); + let env_key = format!("CUDA_OXIDE_SHADERS_PTX_{crate_name}"); + println!("cargo:rerun-if-env-changed={env_key}"); + if let Some(prebuilt) = std::env::var_os(&env_key) { + let dst = output_dir.join("shaders.ptx"); + std::fs::copy(&prebuilt, &dst).unwrap_or_else(|e| { + panic!("failed to copy prebuilt PTX {prebuilt:?} ({env_key}) to {dst:?}: {e}") + }); + return; + } + let features_str = self.features.join(","); let mut args = vec![ diff --git a/crates/khal/src/backend/any_backend.rs b/crates/khal/src/backend/any_backend.rs index 688cb39..f85a04b 100644 --- a/crates/khal/src/backend/any_backend.rs +++ b/crates/khal/src/backend/any_backend.rs @@ -57,10 +57,88 @@ impl GpuBackend { } } - /// Returns `true` if this is the CUDA backend. - #[cfg(feature = "cuda")] + /// Returns `true` if this is the CUDA backend. Always available (returns + /// `false` when the `cuda` feature is disabled) so callers need not gate. pub fn is_cuda(&self) -> bool { - matches!(self, Self::Cuda(..)) + #[cfg(feature = "cuda")] + { + matches!(self, Self::Cuda(..)) + } + #[cfg(not(feature = "cuda"))] + { + false + } + } + + /// Auto-selects a backend: native CUDA on a Blackwell (sm_120+) device when + /// the `cuda` feature is compiled, else WebGPU. Override with + /// `KHAL_BACKEND=cuda|webgpu`. + pub async fn auto(features: wgpu::Features, limits: wgpu::Limits) -> anyhow::Result { + // Resolve any explicit override: Some(true)=force CUDA, Some(false)=force + // WebGPU, None=auto-detect. + let want_cuda = match std::env::var("KHAL_BACKEND").ok().as_deref() { + Some("cuda") | Some("CUDA") => Some(true), + Some("webgpu") | Some("WebGpu") | Some("wgpu") => Some(false), + Some(other) => { + eprintln!("[khal] unknown KHAL_BACKEND={other:?}; auto-detecting"); + None + } + None => { + if std::env::var("BIPED_CUDA").as_deref() == Ok("1") { + eprintln!( + "[khal] BIPED_CUDA=1 is deprecated; use KHAL_BACKEND=cuda (treating as such)" + ); + Some(true) + } else { + None + } + } + }; + + if want_cuda == Some(false) { + eprintln!("[khal] backend = WebGPU"); + return Ok(Self::WebGpu(WebGpu::new(features, limits).await?)); + } + + #[cfg(feature = "cuda")] + { + match Cuda::new(0) { + Ok(cuda) => { + let cc = cuda.compute_capability().ok(); + let is_blackwell = matches!(cc, Some((maj, _)) if maj >= 12); + if want_cuda == Some(true) || is_blackwell { + match cc { + Some((maj, min)) => { + eprintln!("[khal] backend = native CUDA (sm_{maj}{min})") + } + None => eprintln!("[khal] backend = native CUDA"), + } + return Ok(Self::Cuda(cuda)); + } + if let Some((maj, min)) = cc { + eprintln!("[khal] CUDA device sm_{maj}{min} < sm_120; using WebGPU"); + } + } + Err(e) => { + // Explicitly requested but unavailable -> surface the error; + // otherwise silently fall back to WebGPU (no CUDA device). + if want_cuda == Some(true) { + return Err(e.into()); + } + } + } + } + #[cfg(not(feature = "cuda"))] + { + if want_cuda == Some(true) { + eprintln!( + "[khal] KHAL_BACKEND=cuda but the `cuda` feature isn't compiled; using WebGPU" + ); + } + } + + eprintln!("[khal] backend = WebGPU"); + Ok(Self::WebGpu(WebGpu::new(features, limits).await?)) } /// Returns `true` if this is the Metal backend. diff --git a/crates/khal/src/backend/cuda.rs b/crates/khal/src/backend/cuda.rs index c9d7a52..c09c03e 100644 --- a/crates/khal/src/backend/cuda.rs +++ b/crates/khal/src/backend/cuda.rs @@ -8,7 +8,38 @@ use cudarc::driver::{self, CudaContext, CudaStream}; use std::collections::HashMap; use std::marker::PhantomData; use std::ops::RangeBounds; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Mutex, OnceLock}; + +static KERNEL_PROFILE: OnceLock>> = OnceLock::new(); +fn kernel_profile() -> &'static Mutex> { + KERNEL_PROFILE.get_or_init(|| Mutex::new(HashMap::new())) +} +/// Print accumulated per-kernel timings (sorted desc) and clear. No-op if unused. +pub fn dump_kernel_profile() { + let mut map = kernel_profile().lock().unwrap(); + if map.is_empty() { + return; + } + let mut rows: Vec<_> = map.iter().map(|(k, (c, ns))| (k.clone(), *c, *ns)).collect(); + rows.sort_by(|a, b| b.2.cmp(&a.2)); + let total: u128 = rows.iter().map(|r| r.2).sum(); + eprintln!( + "\n=== KHAL_CUDA_PROFILE: {} kernels, {:.3} ms total (serialized) ===", + rows.len(), + total as f64 / 1e6 + ); + for (n, c, ns) in &rows { + eprintln!( + "{:>9.3} ms {:>6}x {:>9.2} us {:>6.2}% {}", + *ns as f64 / 1e6, + c, + (*ns as f64 / *c as f64) / 1e3, + *ns as f64 / total as f64 * 100.0, + n + ); + } + map.clear(); +} // ── Core backend ─────────────────────────────────────────────────────── @@ -39,10 +70,68 @@ impl Cuda { &self.ctx } + /// Compute capability `(major, minor)` of the underlying device, e.g. + /// `(12, 0)` for Blackwell `sm_120`. Used to pick the default backend. + pub fn compute_capability(&self) -> Result<(i32, i32), CudaBackendError> { + Ok(self.ctx.compute_capability()?) + } + /// Returns the default stream. pub fn stream(&self) -> &Arc { &self.stream } + + /// Begin CUDA stream capture on the default stream: subsequent kernel + /// launches are *recorded* into a graph instead of executed, until + /// [`Cuda::end_capture`]. Used by the GPU-resident rollout to capture a + /// repeated dispatch sequence (e.g. the physics decimation loop) once and + /// replay it with a single launch — eliminating per-launch host encode/submit + /// overhead. THREAD_LOCAL mode scopes capture to the current thread. + /// + /// The captured sequence must be replay-safe: no allocation/free, no host + /// syncs, and stable buffer addresses across replays. + pub fn begin_capture(&self) -> Result<(), CudaBackendError> { + use cudarc::driver::sys::CUstreamCaptureMode; + self.stream + .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL)?; + Ok(()) + } + + /// End stream capture and instantiate the recorded graph for replay. Errors + /// if nothing was captured (the stream wasn't in capture, or recorded no work). + pub fn end_capture(&self) -> Result { + use cudarc::driver::sys::CUgraphInstantiate_flags; + // No instantiate flags (0). The flags enum has no zero variant; it is + // consumed as `flags as u32 as u64`, so a 0-valued enum means "no flags". + let no_flags: CUgraphInstantiate_flags = unsafe { std::mem::transmute(0u32) }; + let graph = self + .stream + .end_capture(no_flags)? + .ok_or(CudaBackendError::CaptureFailed)?; + Ok(CapturedGraph { graph }) + } +} + +/// A captured, instantiated CUDA graph. Replay the whole recorded kernel +/// sequence with a single [`CapturedGraph::launch`] (one `cuGraphLaunch`), +/// instead of re-encoding/re-submitting each dispatch from the host. +pub struct CapturedGraph { + graph: driver::CudaGraph, +} + +impl CapturedGraph { + /// Replay the captured kernel sequence with a single graph launch. + pub fn launch(&self) -> Result<(), CudaBackendError> { + self.graph.launch()?; + Ok(()) + } + + /// Pre-upload the graph's resources so the first [`CapturedGraph::launch`] + /// doesn't pay instantiation/upload cost. + pub fn upload(&self) -> Result<(), CudaBackendError> { + self.graph.upload()?; + Ok(()) + } } // ── Error ────────────────────────────────────────────────────────────── @@ -55,6 +144,8 @@ pub enum CudaBackendError { Driver(#[from] driver::DriverError), #[error("Invalid PTX module")] InvalidPtx, + #[error("CUDA stream capture produced no graph")] + CaptureFailed, } // ── Buffer ───────────────────────────────────────────────────────────── @@ -621,9 +712,23 @@ impl<'a> Dispatch<'a, Cuda> for CudaDispatch<'a> { self.function.name, param_values.len(), grid_dim, block_dim ); } + let prof_start = if std::env::var_os("KHAL_CUDA_PROFILE").is_some() { + self.stream.synchronize()?; + Some(std::time::Instant::now()) + } else { + None + }; unsafe { builder.launch(cfg)?; } + if let Some(t0) = prof_start { + self.stream.synchronize()?; + let ns = t0.elapsed().as_nanos(); + let mut map = kernel_profile().lock().unwrap(); + let e = map.entry(self.function.name.clone()).or_insert((0, 0)); + e.0 += 1; + e.1 += ns; + } if trace { match self.stream.synchronize() { Ok(()) => eprintln!("[khal-cuda ok ] {}", self.function.name), From 033d51b6a4b1d640ffb882ea5ee6b6d51762045b Mon Sep 17 00:00:00 2001 From: haixuantao Date: Wed, 8 Jul 2026 11:04:09 +0200 Subject: [PATCH 08/33] khal-std: trig Float methods, generic SmemBuf, panic-free atomics for cuda-oxide MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - num_traits: Float gains sin/cos/asin/acos/atan2 via __nv_* libdevice externs (nexus shaders route trig through Float on the device). - SmemBuf is generic over the element type (SmemBuf) — nexus/vortx workgroup tiles are u32/Vec3/i32, not just f32. - atomics: on nvptx64+cuda-oxide, load/store go through cuda-device atomics; core's AtomicU32 load/store keep a reachable AcqRel panic arm that the cuda-oxide backend rejects. - Cargo.toml: wire the cuda-device path dep for the cuda-oxide feature (machine-local path, as documented in the manifest comment). --- crates/khal-std/Cargo.toml | 4 +-- crates/khal-std/src/lib.rs | 12 ++++----- crates/khal-std/src/num_traits.rs | 39 ++++++++++++++++++++++++++--- crates/khal-std/src/sync/atomics.rs | 35 +++++++++++++++++++++++--- 4 files changed, 75 insertions(+), 15 deletions(-) diff --git a/crates/khal-std/Cargo.toml b/crates/khal-std/Cargo.toml index 19074b1..b2e9c6c 100644 --- a/crates/khal-std/Cargo.toml +++ b/crates/khal-std/Cargo.toml @@ -38,7 +38,7 @@ rust-cuda = ["dep:cuda_std"] # `cuda-device` dependency, but that crate lives in the cuda-oxide tree and is # only present on the build box; it is added there. This manifest keeps the # feature host-buildable everywhere else. -cuda-oxide = ["dep:num-traits"] +cuda-oxide = ["dep:num-traits", "dep:cuda-device"] [dependencies] glamx = { version = "0.3", default-features = false, features = ["nostd-libm", "bytemuck", "u32", "i32", "f64"] } @@ -61,4 +61,4 @@ cuda_std = { git = "https://github.com/sebcrozet/Rust-CUDA.git", branch = "glam- # `Float` for the cuda-oxide backend (libm software math, no_std). num-traits = { version = "0.2", default-features = false, features = ["libm"], optional = true } # (cuda-oxide device build only, added on the build box:) -# cuda-device = { path = ".../cuda-oxide/crates/cuda-device", optional = true } +cuda-device = { path = "/home/xavier/cuda-oxide-src/crates/cuda-device", optional = true } diff --git a/crates/khal-std/src/lib.rs b/crates/khal-std/src/lib.rs index 12c9b19..2a3e3b6 100644 --- a/crates/khal-std/src/lib.rs +++ b/crates/khal-std/src/lib.rs @@ -56,16 +56,16 @@ pub mod cuda_oxide_glue { use crate::index::MaybeIndexUnchecked; use cuda_device::SharedArray; - pub struct SmemBuf(pub &'static mut SharedArray); - impl MaybeIndexUnchecked for SmemBuf { + pub struct SmemBuf(pub &'static mut SharedArray); + impl MaybeIndexUnchecked for SmemBuf { #[inline(always)] - fn read(&self, i: usize) -> f32 { self.0[i] } + fn read(&self, i: usize) -> T { self.0[i] } #[inline(always)] - fn write(&mut self, i: usize, v: f32) { self.0[i] = v; } + fn write(&mut self, i: usize, v: T) { self.0[i] = v; } #[inline(always)] - fn at_mut(&mut self, i: usize) -> &mut f32 { &mut self.0[i] } + fn at_mut(&mut self, i: usize) -> &mut T { &mut self.0[i] } #[inline(always)] - fn at(&self, i: usize) -> &f32 { &self.0[i] } + fn at(&self, i: usize) -> &T { &self.0[i] } } } diff --git a/crates/khal-std/src/num_traits.rs b/crates/khal-std/src/num_traits.rs index 6fba8c6..cc43925 100644 --- a/crates/khal-std/src/num_traits.rs +++ b/crates/khal-std/src/num_traits.rs @@ -27,10 +27,16 @@ mod cuda_oxide_float { fn max(self, other: Self) -> Self; fn min(self, other: Self) -> Self; fn atan(self) -> Self; + fn sin(self) -> Self; + fn cos(self) -> Self; + fn asin(self) -> Self; + fn acos(self) -> Self; + fn atan2(self, x: Self) -> Self; } macro_rules! float_impl { - ($ty:ty, $exp:ident, $ln:ident, $sqrt:ident, $pow:ident, $floor:ident, $ceil:ident, $nvatan:ident) => { + ($ty:ty, $exp:ident, $ln:ident, $sqrt:ident, $pow:ident, $floor:ident, $ceil:ident, + $nvatan:ident, $nvsin:ident, $nvcos:ident, $nvasin:ident, $nvacos:ident, $nvatan2:ident) => { impl Float for $ty { #[inline(always)] fn exp(self) -> $ty { unsafe { core::intrinsics::$exp(self) } } @@ -56,9 +62,36 @@ mod cuda_oxide_float { unsafe extern "C" { fn $nvatan(x: $ty) -> $ty; } unsafe { $nvatan(self) } } + #[inline(always)] + fn sin(self) -> $ty { + unsafe extern "C" { fn $nvsin(x: $ty) -> $ty; } + unsafe { $nvsin(self) } + } + #[inline(always)] + fn cos(self) -> $ty { + unsafe extern "C" { fn $nvcos(x: $ty) -> $ty; } + unsafe { $nvcos(self) } + } + #[inline(always)] + fn asin(self) -> $ty { + unsafe extern "C" { fn $nvasin(x: $ty) -> $ty; } + unsafe { $nvasin(self) } + } + #[inline(always)] + fn acos(self) -> $ty { + unsafe extern "C" { fn $nvacos(x: $ty) -> $ty; } + unsafe { $nvacos(self) } + } + #[inline(always)] + fn atan2(self, x: $ty) -> $ty { + unsafe extern "C" { fn $nvatan2(y: $ty, x: $ty) -> $ty; } + unsafe { $nvatan2(self, x) } + } } }; } - float_impl!(f32, expf32, logf32, sqrtf32, powf32, floorf32, ceilf32, __nv_atanf); - float_impl!(f64, expf64, logf64, sqrtf64, powf64, floorf64, ceilf64, __nv_atan); + float_impl!(f32, expf32, logf32, sqrtf32, powf32, floorf32, ceilf32, + __nv_atanf, __nv_sinf, __nv_cosf, __nv_asinf, __nv_acosf, __nv_atan2f); + float_impl!(f64, expf64, logf64, sqrtf64, powf64, floorf64, ceilf64, + __nv_atan, __nv_sin, __nv_cos, __nv_asin, __nv_acos, __nv_atan2); } diff --git a/crates/khal-std/src/sync/atomics.rs b/crates/khal-std/src/sync/atomics.rs index a1629a2..d759f08 100644 --- a/crates/khal-std/src/sync/atomics.rs +++ b/crates/khal-std/src/sync/atomics.rs @@ -112,7 +112,15 @@ pub fn atomic_load_u32(ptr: &mut u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr) } - #[cfg(not(target_arch = "spirv"))] + // core's `AtomicU32::load` keeps a reachable panic arm (`AcqRel` load), + // which the cuda-oxide backend rejects; use the cuda-device atomic, whose + // load lowers directly to `ld.relaxed.gpu`. + #[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] + { + use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; + unsafe { DeviceAtomicU32::from_ptr(ptr) }.load(AtomicOrdering::Relaxed) + } + #[cfg(not(any(target_arch = "spirv", all(target_arch = "nvptx64", feature = "cuda-oxide"))))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -131,7 +139,14 @@ pub fn atomic_load_u32_shared(ptr: &u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr) } - #[cfg(not(target_arch = "spirv"))] + // See `atomic_load_u32`: avoid core's panic-carrying load on cuda-oxide. + #[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] + { + use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; + unsafe { DeviceAtomicU32::from_ptr(ptr as *const u32 as *mut u32) } + .load(AtomicOrdering::Relaxed) + } + #[cfg(not(any(target_arch = "spirv", all(target_arch = "nvptx64", feature = "cuda-oxide"))))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *const u32 as *const AtomicU32) }; @@ -292,7 +307,13 @@ pub fn atomic_store_u32_workgroup(ptr: &mut u32, value: u32) { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value); } - #[cfg(not(target_arch = "spirv"))] + // See `atomic_load_u32`: avoid core's panic-carrying store on cuda-oxide. + #[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] + { + use cuda_device::atomic::{AtomicOrdering, BlockAtomicU32}; + unsafe { BlockAtomicU32::from_ptr(ptr) }.store(value, AtomicOrdering::Relaxed); + } + #[cfg(not(any(target_arch = "spirv", all(target_arch = "nvptx64", feature = "cuda-oxide"))))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -311,7 +332,13 @@ pub fn atomic_load_u32_workgroup(ptr: &mut u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr) } - #[cfg(not(target_arch = "spirv"))] + // See `atomic_load_u32`: avoid core's panic-carrying load on cuda-oxide. + #[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] + { + use cuda_device::atomic::{AtomicOrdering, BlockAtomicU32}; + unsafe { BlockAtomicU32::from_ptr(ptr) }.load(AtomicOrdering::Relaxed) + } + #[cfg(not(any(target_arch = "spirv", all(target_arch = "nvptx64", feature = "cuda-oxide"))))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; From 63af1e9331b727dcbd6375a7e069e022ccaad222 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Wed, 8 Jul 2026 17:03:45 +0200 Subject: [PATCH 09/33] khal-std/derive: unified-mode cuda-oxide path (host-target, no --device) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rework the cuda-oxide backend path to compile through upstream cuda-oxide's supported model — unified interception on the HOST target — instead of a custom --target nvptx64 device build. Gates move from cfg(target_arch = "nvptx64") to cfg(feature = "cuda-oxide"), and every raw LLVM construct is replaced by the cuda-device API the importer recognizes by name on any compilation target: - sregs: `llvm.nvvm.read.ptx.sreg.*` link-name externs -> cuda_device::thread - barrier: `llvm.nvvm.barrier0` extern -> cuda_device::sync_threads - warp shuffles: nvptx inline asm -> cuda_device::warp::shuffle_xor_f32 - cuda-device/num-traits become plain (not target-gated) optional deps khal-derive generates the CUDA entries with cuda-oxide's reserved kernel prefix directly (cuda_oxide_kernel_246e25db_) instead of the heuristic otherwise skips trivial kernels that nothing in-crate calls). khal-std deliberately does NOT re-export cuda-device at the crate root: the khal_std::cuda_device::* visible path breaks the importer's by-name dispatch. Shader crates now build with plain cargo build -p --features "cuda-oxide ... glamx/scalar-math" (host target, glam scalar math instead of host SIMD). Verified: nexus + vortx shader crates compile against upstream cuda-oxide main (+ the two open fix PRs), cubins assemble, and pendulum PPO / biped training on an RTX 5080 match the device-build reward trajectories exactly. Bonus: the vortx obs/reward kernels (previously excluded from device builds) now compile too. --- crates/khal-derive/src/spirv_bindgen/cuda.rs | 34 +++++++++++---- crates/khal-std/Cargo.toml | 16 ++++--- crates/khal-std/src/arch/cuda.rs | 44 ++++++++++---------- crates/khal-std/src/arch/mod.rs | 2 +- crates/khal-std/src/index.rs | 16 +++---- crates/khal-std/src/lib.rs | 18 ++++---- crates/khal-std/src/memory.rs | 8 ++-- crates/khal-std/src/num_traits.rs | 6 +-- crates/khal-std/src/sync/atomics.rs | 16 +++---- crates/khal-std/src/sync/barriers.rs | 32 +++++++------- crates/khal-std/src/sync/subgroup.rs | 31 ++++++++++---- 11 files changed, 129 insertions(+), 94 deletions(-) diff --git a/crates/khal-derive/src/spirv_bindgen/cuda.rs b/crates/khal-derive/src/spirv_bindgen/cuda.rs index 5063f1e..955153a 100644 --- a/crates/khal-derive/src/spirv_bindgen/cuda.rs +++ b/crates/khal-derive/src/spirv_bindgen/cuda.rs @@ -403,8 +403,8 @@ pub(super) fn generate_cuda_oxide_entry_block( syn::Ident::new(&format!("__smembuf_{}", name), name.span()); preludes.push(quote! { static mut #static_name: - khal_std::cuda_device::SharedArray<#elem_ty, { #len_expr }> = - khal_std::cuda_device::SharedArray::UNINIT; + khal_std::cuda_oxide_glue::SharedArray<#elem_ty, { #len_expr }> = + khal_std::cuda_oxide_glue::SharedArray::UNINIT; let mut #sb_name = khal_std::cuda_oxide_glue::SmemBuf( unsafe { &mut *core::ptr::addr_of_mut!(#static_name) } ); @@ -419,9 +419,9 @@ pub(super) fn generate_cuda_oxide_entry_block( syn::Ident::new(&format!("__sref_{}", name), name.span()); preludes.push(quote! { static mut #static_name: - khal_std::cuda_device::SharedArray<#elem_ty, 1> = - khal_std::cuda_device::SharedArray::UNINIT; - let #sref_name: &mut khal_std::cuda_device::SharedArray<#elem_ty, 1> = + khal_std::cuda_oxide_glue::SharedArray<#elem_ty, 1> = + khal_std::cuda_oxide_glue::SharedArray::UNINIT; + let #sref_name: &mut khal_std::cuda_oxide_glue::SharedArray<#elem_ty, 1> = unsafe { &mut *core::ptr::addr_of_mut!(#static_name) }; let #name: &mut #elem_ty = &mut #sref_name[0]; }); @@ -436,10 +436,28 @@ pub(super) fn generate_cuda_oxide_entry_block( let body = &func.block; + // Name the entry with cuda-oxide's reserved kernel prefix directly + // (`cuda_oxide_kernel_246e25db_`) instead of going through the + // `#[kernel]` proc-macro: the collector roots kernels by that def-path + // marker and the PTX entry keeps the unprefixed base name, while the + // macro's host-side glue (a `cuda_host::CudaKernel` marker impl) would + // drag a cuda-host dependency into pure shader crates. Feature-gated, + // not target-gated: with the cuda-oxide backend installed the crate is + // compiled on the HOST target (unified interception) or for nvptx64 + // (device-only builds) — the entry is generated identically for both. + let prefixed_entry_ident = syn::Ident::new( + &format!("cuda_oxide_kernel_246e25db_{}", cuda_entry_ident), + cuda_entry_ident.span(), + ); quote! { - #[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] - #[khal_std::cuda_device::kernel] - pub fn #cuda_entry_ident(#(#kernel_params),*) { + #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] + #[allow(non_snake_case)] + // `no_mangle` forces local codegen: without it rustc's + // cross-crate-inlinable heuristic skips emitting a mono item for + // trivial kernels (nothing in-crate calls an entry point), and the + // collector never sees them. + #[unsafe(no_mangle)] + pub fn #prefixed_entry_ident(#(#kernel_params),*) { #(#preludes)* #(#array_reconstructions)* #body diff --git a/crates/khal-std/Cargo.toml b/crates/khal-std/Cargo.toml index b2e9c6c..880d4cd 100644 --- a/crates/khal-std/Cargo.toml +++ b/crates/khal-std/Cargo.toml @@ -46,6 +46,14 @@ rayon = { version = "1", optional = true } corosensei = { version = "0.3", optional = true } spirv-std-macros = "0.10.0-alpha.1" khal-derive = { version = "0.3.0", path = "../khal-derive" } +# `Float` for the cuda-oxide backend (libm software math, no_std). +num-traits = { version = "0.2", default-features = false, features = ["libm"], optional = true } +# cuda-oxide device API (sregs, barriers, warp shuffles, atomics, SharedArray). +# Plain (not target-gated) dependency: with the cuda-oxide backend the shader +# crate is compiled on the HOST target (unified interception) or for nvptx64 +# (device-only builds); the importer recognizes the API by name either way. +# Lives in the cuda-oxide tree; the path is machine-local. +cuda-device = { path = "/home/xavier/cuda-oxide-src/crates/cuda-device", optional = true } [lints] workspace = true @@ -55,10 +63,6 @@ spirv-std = "0.10.0-alpha.1" [target.'cfg(target_arch = "nvptx64")'.dependencies] # Fixes the UVec3::element_product bug. Optional: only the `rust-cuda` backend -# (rustc_codegen_nvvm) uses it; the `cuda-oxide` backend declares its own -# intrinsics instead (cuda_std does not build under `-Zbuild-std`). +# (rustc_codegen_nvvm) uses it; the `cuda-oxide` backend routes through +# cuda-device instead (cuda_std does not build under `-Zbuild-std`). cuda_std = { git = "https://github.com/sebcrozet/Rust-CUDA.git", branch = "glam-up", optional = true } -# `Float` for the cuda-oxide backend (libm software math, no_std). -num-traits = { version = "0.2", default-features = false, features = ["libm"], optional = true } -# (cuda-oxide device build only, added on the build box:) -cuda-device = { path = "/home/xavier/cuda-oxide-src/crates/cuda-device", optional = true } diff --git a/crates/khal-std/src/arch/cuda.rs b/crates/khal-std/src/arch/cuda.rs index 6bf9070..d56531b 100644 --- a/crates/khal-std/src/arch/cuda.rs +++ b/crates/khal-std/src/arch/cuda.rs @@ -3,43 +3,41 @@ //! Two backends provide the underlying `%tid`/`%ctaid`/`%ntid`/`%nctaid` //! special-register reads: //! - `rust-cuda` (default): `cuda_std::thread` (Rust-CUDA / rustc_codegen_nvvm). -//! - `cuda-oxide`: khal_std's own `llvm.nvvm.read.ptx.sreg.*` externs, lowered -//! by the cuda-oxide PTX backend. `cuda_std` does not build under cuda-oxide. +//! - `cuda-oxide`: `cuda_device::thread::*` accessors, recognized by name by +//! the cuda-oxide MIR importer. `cuda_std` does not build under cuda-oxide. use glamx::UVec3; #[cfg(not(feature = "cuda-oxide"))] use cuda_std::thread; -/// khal_std's own nvptx special-register reads, used by the `cuda-oxide` -/// backend in place of `cuda_std::thread`. Each lowers to a single -/// `mov.u32 %r, %sreg` in PTX. +/// Special-register reads for the `cuda-oxide` backend, routed through +/// `cuda_device::thread::*`: the cuda-oxide MIR importer recognizes those +/// paths by name on ANY compilation target, so this works both for unified +/// (host-target) shader builds and for `--target nvptx64` device builds. +/// Each lowers to a single `mov.u32 %r, %sreg` in PTX. #[cfg(feature = "cuda-oxide")] mod thread { macro_rules! sreg { - ($name:ident, $intr:literal) => { + ($name:ident, $dev:ident) => { #[inline(always)] pub fn $name() -> u32 { - unsafe extern "C" { - #[link_name = $intr] - fn f() -> u32; - } - unsafe { f() } + cuda_device::thread::$dev() } }; } - sreg!(thread_idx_x, "llvm.nvvm.read.ptx.sreg.tid.x"); - sreg!(thread_idx_y, "llvm.nvvm.read.ptx.sreg.tid.y"); - sreg!(thread_idx_z, "llvm.nvvm.read.ptx.sreg.tid.z"); - sreg!(block_idx_x, "llvm.nvvm.read.ptx.sreg.ctaid.x"); - sreg!(block_idx_y, "llvm.nvvm.read.ptx.sreg.ctaid.y"); - sreg!(block_idx_z, "llvm.nvvm.read.ptx.sreg.ctaid.z"); - sreg!(block_dim_x, "llvm.nvvm.read.ptx.sreg.ntid.x"); - sreg!(block_dim_y, "llvm.nvvm.read.ptx.sreg.ntid.y"); - sreg!(block_dim_z, "llvm.nvvm.read.ptx.sreg.ntid.z"); - sreg!(grid_dim_x, "llvm.nvvm.read.ptx.sreg.nctaid.x"); - sreg!(grid_dim_y, "llvm.nvvm.read.ptx.sreg.nctaid.y"); - sreg!(grid_dim_z, "llvm.nvvm.read.ptx.sreg.nctaid.z"); + sreg!(thread_idx_x, threadIdx_x); + sreg!(thread_idx_y, threadIdx_y); + sreg!(thread_idx_z, threadIdx_z); + sreg!(block_idx_x, blockIdx_x); + sreg!(block_idx_y, blockIdx_y); + sreg!(block_idx_z, blockIdx_z); + sreg!(block_dim_x, blockDim_x); + sreg!(block_dim_y, blockDim_y); + sreg!(block_dim_z, blockDim_z); + sreg!(grid_dim_x, gridDim_x); + sreg!(grid_dim_y, gridDim_y); + sreg!(grid_dim_z, gridDim_z); } /// Returns the thread index within the current block as a `UVec3`. diff --git a/crates/khal-std/src/arch/mod.rs b/crates/khal-std/src/arch/mod.rs index 5266bf8..79b040f 100644 --- a/crates/khal-std/src/arch/mod.rs +++ b/crates/khal-std/src/arch/mod.rs @@ -4,5 +4,5 @@ #[cfg(feature = "cpu")] pub mod cpu; /// CUDA intrinsics for thread and block indexing. -#[cfg(target_arch = "nvptx64")] +#[cfg(any(target_arch = "nvptx64", feature = "cuda-oxide"))] pub mod cuda; diff --git a/crates/khal-std/src/index.rs b/crates/khal-std/src/index.rs index d32b3f6..a6bd075 100644 --- a/crates/khal-std/src/index.rs +++ b/crates/khal-std/src/index.rs @@ -17,9 +17,9 @@ pub trait MaybeIndexUnchecked { impl MaybeIndexUnchecked for [T] { #[inline(always)] fn at(&self, id: usize) -> &T { - #[cfg(all(feature = "unsafe_remove_boundchecks", target_arch = "nvptx64"))] + #[cfg(all(feature = "unsafe_remove_boundchecks", any(target_arch = "nvptx64", feature = "cuda-oxide")))] return unsafe { self.get_unchecked(id) }; - #[cfg(all(feature = "unsafe_remove_boundchecks", not(target_arch = "nvptx64")))] + #[cfg(all(feature = "unsafe_remove_boundchecks", not(any(target_arch = "nvptx64", feature = "cuda-oxide"))))] return unsafe { use spirv_std::arch::IndexUnchecked; self.index_unchecked(id) @@ -30,9 +30,9 @@ impl MaybeIndexUnchecked for [T] { #[inline(always)] fn at_mut(&mut self, id: usize) -> &mut T { - #[cfg(all(feature = "unsafe_remove_boundchecks", target_arch = "nvptx64"))] + #[cfg(all(feature = "unsafe_remove_boundchecks", any(target_arch = "nvptx64", feature = "cuda-oxide")))] return unsafe { self.get_unchecked_mut(id) }; - #[cfg(all(feature = "unsafe_remove_boundchecks", not(target_arch = "nvptx64")))] + #[cfg(all(feature = "unsafe_remove_boundchecks", not(any(target_arch = "nvptx64", feature = "cuda-oxide"))))] return unsafe { use spirv_std::arch::IndexUnchecked; self.index_unchecked_mut(id) @@ -55,9 +55,9 @@ impl MaybeIndexUnchecked for [T] { impl MaybeIndexUnchecked for [T; N] { #[inline(always)] fn at(&self, id: usize) -> &T { - #[cfg(all(feature = "unsafe_remove_boundchecks", target_arch = "nvptx64"))] + #[cfg(all(feature = "unsafe_remove_boundchecks", any(target_arch = "nvptx64", feature = "cuda-oxide")))] return unsafe { self.get_unchecked(id) }; - #[cfg(all(feature = "unsafe_remove_boundchecks", not(target_arch = "nvptx64")))] + #[cfg(all(feature = "unsafe_remove_boundchecks", not(any(target_arch = "nvptx64", feature = "cuda-oxide"))))] return unsafe { use spirv_std::arch::IndexUnchecked; self.index_unchecked(id) @@ -68,9 +68,9 @@ impl MaybeIndexUnchecked for [T; N] { #[inline(always)] fn at_mut(&mut self, id: usize) -> &mut T { - #[cfg(all(feature = "unsafe_remove_boundchecks", target_arch = "nvptx64"))] + #[cfg(all(feature = "unsafe_remove_boundchecks", any(target_arch = "nvptx64", feature = "cuda-oxide")))] return unsafe { self.get_unchecked_mut(id) }; - #[cfg(all(feature = "unsafe_remove_boundchecks", not(target_arch = "nvptx64")))] + #[cfg(all(feature = "unsafe_remove_boundchecks", not(any(target_arch = "nvptx64", feature = "cuda-oxide"))))] return unsafe { use spirv_std::arch::IndexUnchecked; self.index_unchecked_mut(id) diff --git a/crates/khal-std/src/lib.rs b/crates/khal-std/src/lib.rs index 2a3e3b6..8a404dd 100644 --- a/crates/khal-std/src/lib.rs +++ b/crates/khal-std/src/lib.rs @@ -10,8 +10,8 @@ #![cfg_attr(target_arch = "nvptx64", feature(asm_experimental_arch))] // `Float` for the cuda-oxide backend calls `core::intrinsics::{expf32,..}` // (lowered to libdevice `__nv_*` by cuda-oxide) instead of software libm. -#![cfg_attr(all(target_arch = "nvptx64", feature = "cuda-oxide"), feature(core_intrinsics))] -#![cfg_attr(all(target_arch = "nvptx64", feature = "cuda-oxide"), allow(internal_features))] +#![cfg_attr(feature = "cuda-oxide", feature(core_intrinsics))] +#![cfg_attr(feature = "cuda-oxide", allow(internal_features))] /// Architecture-specific runtime support (CPU coroutines, CUDA intrinsics). pub mod arch; @@ -41,20 +41,20 @@ pub use glamx; #[cfg(all(target_arch = "nvptx64", not(feature = "cuda-oxide")))] pub use cuda_std; -/// cuda-oxide device crate (re-exported so the `spirv_bindgen`-generated CUDA -/// entry points can name `khal_std::cuda_device::kernel` without the shader -/// crate depending on cuda-device directly). -#[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] -pub use cuda_device; +// NOTE: cuda-device is deliberately NOT re-exported at the crate root: a +// `pub use cuda_device;` gives its functions a `khal_std::cuda_device::*` +// visible path, which breaks the cuda-oxide importer's by-name intrinsic +// dispatch (it matches the canonical `cuda_device::*` paths). The types the +// generated entry code needs are re-exported from `cuda_oxide_glue` instead. /// Glue for the cuda-oxide `spirv_bindgen` entry transform: wraps a /// `&'static mut SharedArray` so a `#[spirv(workgroup)] &mut [T; N]` parameter /// body (which calls `MaybeIndexUnchecked::read/write/at_mut`) lands on shared /// memory. -#[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] +#[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] pub mod cuda_oxide_glue { use crate::index::MaybeIndexUnchecked; - use cuda_device::SharedArray; + pub use cuda_device::SharedArray; pub struct SmemBuf(pub &'static mut SharedArray); impl MaybeIndexUnchecked for SmemBuf { diff --git a/crates/khal-std/src/memory.rs b/crates/khal-std/src/memory.rs index bbd1d16..e4c021f 100644 --- a/crates/khal-std/src/memory.rs +++ b/crates/khal-std/src/memory.rs @@ -1,10 +1,10 @@ -#[cfg(target_arch = "nvptx64")] +#[cfg(any(target_arch = "nvptx64", feature = "cuda-oxide"))] pub use memory_nvptx::*; -#[cfg(not(target_arch = "nvptx64"))] +#[cfg(not(any(target_arch = "nvptx64", feature = "cuda-oxide")))] pub use spirv_std::memory::*; -// On nvptx64, provide the memory scope constants. -#[cfg(target_arch = "nvptx64")] +// On the CUDA backends, provide the memory scope constants. +#[cfg(any(target_arch = "nvptx64", feature = "cuda-oxide"))] pub mod memory_nvptx { /// Memory scope levels matching SPIR-V semantics for use on the CUDA backend. #[derive(Copy, Clone)] diff --git a/crates/khal-std/src/num_traits.rs b/crates/khal-std/src/num_traits.rs index cc43925..80ce764 100644 --- a/crates/khal-std/src/num_traits.rs +++ b/crates/khal-std/src/num_traits.rs @@ -6,13 +6,13 @@ // software math does not resolve under cuda-oxide's PTX backend.) #[cfg(all(target_arch = "nvptx64", not(feature = "cuda-oxide")))] pub use cuda_std::float::GpuFloat as Float; -#[cfg(not(target_arch = "nvptx64"))] +#[cfg(not(any(target_arch = "nvptx64", feature = "cuda-oxide")))] pub use spirv_std::num_traits::Float; -#[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] +#[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] pub use cuda_oxide_float::Float; -#[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] +#[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] mod cuda_oxide_float { /// Minimal `Float` for the cuda-oxide PTX backend. Methods lower to /// libdevice (`__nv_expf`, `__nv_sqrtf`, …) via `core::intrinsics`. diff --git a/crates/khal-std/src/sync/atomics.rs b/crates/khal-std/src/sync/atomics.rs index d759f08..078ce9a 100644 --- a/crates/khal-std/src/sync/atomics.rs +++ b/crates/khal-std/src/sync/atomics.rs @@ -115,12 +115,12 @@ pub fn atomic_load_u32(ptr: &mut u32) -> u32 { // core's `AtomicU32::load` keeps a reachable panic arm (`AcqRel` load), // which the cuda-oxide backend rejects; use the cuda-device atomic, whose // load lowers directly to `ld.relaxed.gpu`. - #[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] + #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] { use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; unsafe { DeviceAtomicU32::from_ptr(ptr) }.load(AtomicOrdering::Relaxed) } - #[cfg(not(any(target_arch = "spirv", all(target_arch = "nvptx64", feature = "cuda-oxide"))))] + #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -140,13 +140,13 @@ pub fn atomic_load_u32_shared(ptr: &u32) -> u32 { >(ptr) } // See `atomic_load_u32`: avoid core's panic-carrying load on cuda-oxide. - #[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] + #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] { use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; unsafe { DeviceAtomicU32::from_ptr(ptr as *const u32 as *mut u32) } .load(AtomicOrdering::Relaxed) } - #[cfg(not(any(target_arch = "spirv", all(target_arch = "nvptx64", feature = "cuda-oxide"))))] + #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *const u32 as *const AtomicU32) }; @@ -308,12 +308,12 @@ pub fn atomic_store_u32_workgroup(ptr: &mut u32, value: u32) { >(ptr, value); } // See `atomic_load_u32`: avoid core's panic-carrying store on cuda-oxide. - #[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] + #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] { use cuda_device::atomic::{AtomicOrdering, BlockAtomicU32}; unsafe { BlockAtomicU32::from_ptr(ptr) }.store(value, AtomicOrdering::Relaxed); } - #[cfg(not(any(target_arch = "spirv", all(target_arch = "nvptx64", feature = "cuda-oxide"))))] + #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -333,12 +333,12 @@ pub fn atomic_load_u32_workgroup(ptr: &mut u32) -> u32 { >(ptr) } // See `atomic_load_u32`: avoid core's panic-carrying load on cuda-oxide. - #[cfg(all(target_arch = "nvptx64", feature = "cuda-oxide"))] + #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] { use cuda_device::atomic::{AtomicOrdering, BlockAtomicU32}; unsafe { BlockAtomicU32::from_ptr(ptr) }.load(AtomicOrdering::Relaxed) } - #[cfg(not(any(target_arch = "spirv", all(target_arch = "nvptx64", feature = "cuda-oxide"))))] + #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; diff --git a/crates/khal-std/src/sync/barriers.rs b/crates/khal-std/src/sync/barriers.rs index de4944f..1b3c920 100644 --- a/crates/khal-std/src/sync/barriers.rs +++ b/crates/khal-std/src/sync/barriers.rs @@ -8,26 +8,24 @@ pub fn workgroup_memory_barrier_with_group_sync() { { spirv_std::arch::workgroup_memory_barrier_with_group_sync(); } - #[cfg(target_arch = "nvptx64")] + // cuda-oxide: `cuda_device::sync_threads` is recognized by name by the + // MIR importer on any compilation target and lowers to the convergent + // NVVM barrier op (`bar.sync`), which optimization passes will not + // tail-duplicate into divergent branches. + #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] { - // Call the LLVM intrinsic directly instead of cuda_std::thread::sync_threads() - // so that LLVM sees the `convergent` attribute during optimization passes. - // Without this, LLVM tail-duplicates the barrier into both sides of divergent - // branches (if/else), causing threads to hit different bar.sync instructions - // and deadlocking the block. - // This fixes kernels with barriers that were otherwise hanging when using - // cuda_std::thread::sync_thread() instead. - unsafe extern "C" { - #[link_name = "llvm.nvvm.barrier0"] - fn nvvm_barrier0(); - } - unsafe { - nvvm_barrier0(); - } - // cuda_std::thread::sync_threads(); + cuda_device::sync_threads(); + } + #[cfg(all(target_arch = "nvptx64", not(feature = "cuda-oxide")))] + { + cuda_std::thread::sync_threads(); } - #[cfg(not(any(target_arch = "spirv", target_arch = "nvptx64")))] + #[cfg(not(any( + target_arch = "spirv", + target_arch = "nvptx64", + feature = "cuda-oxide" + )))] #[cfg(feature = "cpu")] { crate::arch::cpu::barrier_wait(); diff --git a/crates/khal-std/src/sync/subgroup.rs b/crates/khal-std/src/sync/subgroup.rs index cf92a09..8607fce 100644 --- a/crates/khal-std/src/sync/subgroup.rs +++ b/crates/khal-std/src/sync/subgroup.rs @@ -11,11 +11,15 @@ pub fn subgroup_f_add(val: f32) -> f32 { { spirv_std::arch::subgroup_f_add(val) } - #[cfg(target_arch = "nvptx64")] + #[cfg(any(target_arch = "nvptx64", feature = "cuda-oxide"))] { warp_reduce_add(val) } - #[cfg(not(any(target_arch = "spirv", target_arch = "nvptx64")))] + #[cfg(not(any( + target_arch = "spirv", + target_arch = "nvptx64", + feature = "cuda-oxide" + )))] { val } @@ -34,17 +38,30 @@ pub fn subgroup_f_max(val: f32) -> f32 { { spirv_std::arch::subgroup_f_max(val) } - #[cfg(target_arch = "nvptx64")] + #[cfg(any(target_arch = "nvptx64", feature = "cuda-oxide"))] { warp_reduce_max(val) } - #[cfg(not(any(target_arch = "spirv", target_arch = "nvptx64")))] + #[cfg(not(any( + target_arch = "spirv", + target_arch = "nvptx64", + feature = "cuda-oxide" + )))] { val } } -#[cfg(target_arch = "nvptx64")] +// cuda-oxide: warp shuffle through `cuda_device::warp`, recognized by name +// by the MIR importer on any compilation target. +#[cfg(feature = "cuda-oxide")] +#[inline(always)] +fn shfl_xor_sync(val: f32, lane_mask: u32) -> f32 { + cuda_device::warp::shuffle_xor_f32(val, lane_mask) +} + +// rust-cuda (rustc_codegen_nvvm): inline PTX. +#[cfg(all(target_arch = "nvptx64", not(feature = "cuda-oxide")))] #[inline(always)] fn shfl_xor_sync(val: f32, lane_mask: u32) -> f32 { let result: f32; @@ -59,7 +76,7 @@ fn shfl_xor_sync(val: f32, lane_mask: u32) -> f32 { result } -#[cfg(target_arch = "nvptx64")] +#[cfg(any(target_arch = "nvptx64", feature = "cuda-oxide"))] #[inline(always)] fn warp_reduce_add(mut val: f32) -> f32 { val += shfl_xor_sync(val, 16); @@ -70,7 +87,7 @@ fn warp_reduce_add(mut val: f32) -> f32 { val } -#[cfg(target_arch = "nvptx64")] +#[cfg(any(target_arch = "nvptx64", feature = "cuda-oxide"))] #[inline(always)] fn warp_reduce_max(mut val: f32) -> f32 { let other = shfl_xor_sync(val, 16); From d72fd4b97f076ec0ad59f6f9ea5102ed63af581f Mon Sep 17 00:00:00 2001 From: haixuantao Date: Wed, 8 Jul 2026 17:45:51 +0200 Subject: [PATCH 10/33] khal: never degrade the backend selection silently MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three hardenings in KhalGpuBackend::auto: - An explicit KHAL_BACKEND=cuda on a binary built without the `cuda` feature is now a hard error instead of a warning + WebGPU fallback: an explicitly requested backend must not be silently substituted (results and performance would differ from what the user asked to measure). - Auto-detect CUDA init failure now prints the error before falling back to WebGPU; on a machine with an NVIDIA GPU a quiet fallback usually hides real breakage behind a slower working run. - A cuda-less binary auto-detecting on a machine with an NVIDIA GPU (/dev/nvidiactl present) now says so — a stale or misbuilt binary otherwise degrades silently to WebGPU. --- crates/khal/src/backend/any_backend.rs | 40 ++++++++++++++++++++++---- 1 file changed, 34 insertions(+), 6 deletions(-) diff --git a/crates/khal/src/backend/any_backend.rs b/crates/khal/src/backend/any_backend.rs index f85a04b..290f02f 100644 --- a/crates/khal/src/backend/any_backend.rs +++ b/crates/khal/src/backend/any_backend.rs @@ -120,19 +120,35 @@ impl GpuBackend { } } Err(e) => { - // Explicitly requested but unavailable -> surface the error; - // otherwise silently fall back to WebGPU (no CUDA device). + // Explicitly requested but unavailable -> surface the error. if want_cuda == Some(true) { return Err(e.into()); } + // Auto-detect: fall back, but never silently — a CUDA init + // failure on a machine with an NVIDIA GPU is usually real + // breakage (driver mismatch, exhausted contexts), and a + // quiet WebGPU run hides it behind a slower working one. + eprintln!("[khal] CUDA unavailable ({e}); using WebGPU"); } } } #[cfg(not(feature = "cuda"))] { if want_cuda == Some(true) { + // An explicit KHAL_BACKEND=cuda must not silently run on a + // different backend: results/performance would differ from + // what the user asked to measure. + anyhow::bail!( + "KHAL_BACKEND=cuda but this binary was built without the `cuda` feature (rebuild with `--features cuda_backend` / `khal/cuda`)" + ); + } + // Auto-detect on a machine that visibly has an NVIDIA GPU: still + // use WebGPU (only backend compiled in), but say what a cuda-less + // binary is leaving on the table — a stale or misbuilt binary + // otherwise degrades silently. + if std::path::Path::new("/dev/nvidiactl").exists() { eprintln!( - "[khal] KHAL_BACKEND=cuda but the `cuda` feature isn't compiled; using WebGPU" + "[khal] NVIDIA GPU present but this binary was built without the `cuda` feature; using WebGPU" ); } } @@ -1715,7 +1731,11 @@ impl<'b, T: DeviceValue> crate::ShaderArgs<'b> for GpuBuffer { // push element count so kernel `slice.len()` is correct (off-by-size_of // otherwise -> OOB reads, e.g. gpu_init_sort_dispatch / lbvh). Arrays, // scalars and uniforms ignore this value, so they are unaffected. - dispatch.set_arg(binding, buffer.device_ptr_raw(), buffer.byte_len() / std::mem::size_of::() as u64); + dispatch.set_arg( + binding, + buffer.device_ptr_raw(), + buffer.byte_len() / std::mem::size_of::() as u64, + ); Ok(()) } #[cfg(feature = "metal")] @@ -1747,7 +1767,11 @@ impl<'b, T: DeviceValue> crate::ShaderArgs<'b> for GpuBufferSlice<'_, T> { } #[cfg(feature = "cuda")] (GpuBufferSlice::Cuda(slice), GpuDispatch::Cuda(dispatch)) => { - dispatch.set_arg(binding, slice.offset_ptr(), slice.byte_len / std::mem::size_of::() as u64); + dispatch.set_arg( + binding, + slice.offset_ptr(), + slice.byte_len / std::mem::size_of::() as u64, + ); Ok(()) } #[cfg(feature = "metal")] @@ -1784,7 +1808,11 @@ impl<'b, T: DeviceValue> crate::ShaderArgs<'b> for GpuBufferSliceMut<'_, T> { } #[cfg(feature = "cuda")] (GpuBufferSliceMut::Cuda(slice), GpuDispatch::Cuda(dispatch)) => { - dispatch.set_arg(binding, slice.offset_ptr(), slice.byte_len / std::mem::size_of::() as u64); + dispatch.set_arg( + binding, + slice.offset_ptr(), + slice.byte_len / std::mem::size_of::() as u64, + ); Ok(()) } #[cfg(feature = "metal")] From 0991e11a3753db8d55aa2ebff3b270f39b1ccbdd Mon Sep 17 00:00:00 2001 From: haixuantao Date: Thu, 9 Jul 2026 15:05:40 +0200 Subject: [PATCH 11/33] =?UTF-8?q?fix(cuda):=20make=20graph=20capture=20act?= =?UTF-8?q?ually=20work=20=E2=80=94=20owned=20stream=20+=20as=5Fcuda=20on?= =?UTF-8?q?=20the=20backend=20enum?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - The CUDA backend ran on cudarc's default_stream(), i.e. the legacy NULL stream, which CUDA refuses to capture (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). Use an owned non-blocking stream instead; all khal work is ordered on this one stream anyway. - GpuBackend (the enum) never overrode Backend::as_cuda, so the graph-capture API was unreachable through it (always None). Delegate to the Cuda variant. - KHAL_TRACE_INDIRECT=1 prints each indirect dispatch's kernel name: on CUDA every indirect dispatch costs a full stream drain + host count read, and invalidates a graph capture — this is the tool for hunting them down. --- crates/khal/src/backend/any_backend.rs | 8 ++++++++ crates/khal/src/backend/cuda.rs | 23 +++++++++++++++++++---- 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/crates/khal/src/backend/any_backend.rs b/crates/khal/src/backend/any_backend.rs index 290f02f..f652024 100644 --- a/crates/khal/src/backend/any_backend.rs +++ b/crates/khal/src/backend/any_backend.rs @@ -1179,6 +1179,14 @@ impl GpuReadback { } impl Backend for GpuBackend { + #[cfg(feature = "cuda")] + fn as_cuda(&self) -> Option<&super::cuda::Cuda> { + match self { + Self::Cuda(c) => Some(c), + _ => None, + } + } + const NAME: &'static str = "any"; const TARGET: super::CompileTarget = super::CompileTarget::Wgsl; diff --git a/crates/khal/src/backend/cuda.rs b/crates/khal/src/backend/cuda.rs index c09c03e..568fcb2 100644 --- a/crates/khal/src/backend/cuda.rs +++ b/crates/khal/src/backend/cuda.rs @@ -20,7 +20,10 @@ pub fn dump_kernel_profile() { if map.is_empty() { return; } - let mut rows: Vec<_> = map.iter().map(|(k, (c, ns))| (k.clone(), *c, *ns)).collect(); + let mut rows: Vec<_> = map + .iter() + .map(|(k, (c, ns))| (k.clone(), *c, *ns)) + .collect(); rows.sort_by(|a, b| b.2.cmp(&a.2)); let total: u128 = rows.iter().map(|r| r.2).sum(); eprintln!( @@ -57,7 +60,10 @@ impl Cuda { /// Creates a new CUDA backend using the specified device ordinal. pub fn new(device_ordinal: usize) -> Result { let ctx = CudaContext::new(device_ordinal)?; - let stream = ctx.default_stream(); + // An owned non-default stream: the legacy NULL stream cannot be + // captured into a CUDA graph (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED), + // and all khal work is ordered on this one stream anyway. + let stream = ctx.new_stream()?; Ok(Self { ctx, stream, @@ -417,7 +423,10 @@ impl Backend for Cuda { return Err(e.into()); } }; - Ok(CudaFunction { func, name: entry_point.to_string() }) + Ok(CudaFunction { + func, + name: entry_point.to_string(), + }) } fn load_function_with_layouts( @@ -632,6 +641,9 @@ impl<'a> Dispatch<'a, Cuda> for CudaDispatch<'a> { DispatchGrid::Indirect(buffer) => { // CUDA doesn't support indirect dispatch natively. // Read the 12-byte dispatch args from device memory. + if std::env::var("KHAL_TRACE_INDIRECT").is_ok() { + eprintln!("[khal] indirect dispatch: {}", self.function.name); + } self.stream.synchronize()?; if let Some(ref inner) = buffer.inner { let bytes: Vec = self.stream.clone_dtoh(inner)?; @@ -709,7 +721,10 @@ impl<'a> Dispatch<'a, Cuda> for CudaDispatch<'a> { if trace { eprintln!( "[khal-cuda launch] {} nargs={} grid={:?} block={:?}", - self.function.name, param_values.len(), grid_dim, block_dim + self.function.name, + param_values.len(), + grid_dim, + block_dim ); } let prof_start = if std::env::var_os("KHAL_CUDA_PROFILE").is_some() { From 96e3d5fc34eec40ae84181aff7fa0624e622e225 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Thu, 16 Jul 2026 17:26:12 +0200 Subject: [PATCH 12/33] fix(khal-std): resolve cuda-device by git rev, not an absolute local path (#5) khal-std pinned cuda-device at path = "/home/xavier/cuda-oxide-src/crates/cuda-device" an absolute path into one developer's home directory. Cargo reads path dependencies eagerly, even optional ones that no active feature selects, so this did not merely break `--features cuda-oxide` -- it broke `cargo metadata` for the whole graph on every other machine: error: failed to load source for dependency `cuda-device` Caused by: Unable to update /home/xavier/cuda-oxide-src/crates/cuda-device Caused by: No such file or directory (os error 2) Found while building nexus (which path-patches khal) on a fresh box: the build is unreachable off that one machine. This also defeats exactly what the feature comment two lines up promises -- "This manifest keeps the feature host-buildable everywhere else." The intent was clearly that the dependency be added locally on the build box and not committed; the local edit escaped into the branch, and the comment describing the intent stayed behind. cuda-device is not on crates.io, so a git rev is the portable way to say this. Pinned rather than floating on a branch, so the resolve is reproducible. Hacking on both trees still works -- add a [patch] with a local checkout in the consuming workspace. Verified with the local cuda-oxide checkout deleted, so no stale path could satisfy it: cargo metadata resolves clean, and with `-p khal-std --features cuda-oxide` cargo fetches the rev and records source = "git+https://github.com/haixuanTao/cuda-oxide.git?rev=8c3b5e2..." cuda-device then still needs nightly (#![feature(f16)]; cuda-oxide pins nightly-2026-04-03 in rust-toolchain.toml) -- that is pre-existing and identical under the old path dep, not a regression from this change. Co-authored-by: haixuanTao --- crates/khal-std/Cargo.toml | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/crates/khal-std/Cargo.toml b/crates/khal-std/Cargo.toml index 880d4cd..59d047d 100644 --- a/crates/khal-std/Cargo.toml +++ b/crates/khal-std/Cargo.toml @@ -52,8 +52,15 @@ num-traits = { version = "0.2", default-features = false, features = ["libm"], o # Plain (not target-gated) dependency: with the cuda-oxide backend the shader # crate is compiled on the HOST target (unified interception) or for nvptx64 # (device-only builds); the importer recognizes the API by name either way. -# Lives in the cuda-oxide tree; the path is machine-local. -cuda-device = { path = "/home/xavier/cuda-oxide-src/crates/cuda-device", optional = true } +# +# Pinned by git rev, not by path. This was previously +# path = "/home/xavier/cuda-oxide-src/crates/cuda-device" +# an absolute path into one developer's home directory, so `--features +# cuda-oxide` could not resolve on any other machine -- which is precisely the +# portability the feature comment above promises. cuda-device is not published +# on crates.io, so a git rev is the portable way to express it. Override with a +# local checkout via [patch] in the consuming workspace when hacking on both. +cuda-device = { git = "https://github.com/haixuanTao/cuda-oxide.git", rev = "8c3b5e2dabad165f0f82d98d72dcf2cd43c50d2d", optional = true } [lints] workspace = true From b0a33923f70a43dc235923966d40c88dd19ca0d6 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Fri, 17 Jul 2026 16:58:36 +0200 Subject: [PATCH 13/33] khal-std: Index/IndexMut for the cuda-oxide SmemBuf glue The fused solver kernels index workgroup buffers directly (vels_smem[i]); the SPIR-V path gets that via the array type, the cuda-oxide SmemBuf wrapper only had the MaybeIndexUnchecked methods. Claude-Session: https://claude.ai/code/session_01U2n9RqmxTJb8UG5d1Sjw4W --- crates/khal-std/src/lib.rs | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/crates/khal-std/src/lib.rs b/crates/khal-std/src/lib.rs index 8a404dd..0edce3c 100644 --- a/crates/khal-std/src/lib.rs +++ b/crates/khal-std/src/lib.rs @@ -67,6 +67,19 @@ pub mod cuda_oxide_glue { #[inline(always)] fn at(&self, i: usize) -> &T { &self.0[i] } } + impl core::ops::Index for SmemBuf { + type Output = T; + #[inline(always)] + fn index(&self, i: usize) -> &T { + &self.0[i] + } + } + impl core::ops::IndexMut for SmemBuf { + #[inline(always)] + fn index_mut(&mut self, i: usize) -> &mut T { + &mut self.0[i] + } + } } #[cfg(not(any(target_arch = "spirv", target_arch = "nvptx64")))] From 4b45508f318e6cd999ddba5efa32ea08c2aab120 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Fri, 17 Jul 2026 17:04:14 +0200 Subject: [PATCH 14/33] khal-std: cuda-oxide arms for the RMW atomics (fetch_add/max/min/or/and/swap) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The load/store fns already routed through cuda_device atomics on cuda-oxide; the RMW fns still fell back to core::sync::atomic, whose fetch_* lower to AtomicRmwOp — rejected by the legacy-NVVM legalizer. cuda_device's fetch_* lower to atom.* intrinsics. Claude-Session: https://claude.ai/code/session_01U2n9RqmxTJb8UG5d1Sjw4W --- crates/khal-std/src/sync/atomics.rs | 81 +++++++++++++++++++++++++---- 1 file changed, 72 insertions(+), 9 deletions(-) diff --git a/crates/khal-std/src/sync/atomics.rs b/crates/khal-std/src/sync/atomics.rs index 078ce9a..f953034 100644 --- a/crates/khal-std/src/sync/atomics.rs +++ b/crates/khal-std/src/sync/atomics.rs @@ -13,7 +13,14 @@ pub fn atomic_add_i32(ptr: &mut i32, value: i32) -> i32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - #[cfg(not(target_arch = "spirv"))] + // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy + // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. + #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] + { + use cuda_device::atomic::{AtomicOrdering, DeviceAtomicI32}; + unsafe { DeviceAtomicI32::from_ptr(ptr) }.fetch_add(value, AtomicOrdering::Relaxed) + } + #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] { use core::sync::atomic::{AtomicI32, Ordering}; let atomic = unsafe { &*(ptr as *mut i32 as *const AtomicI32) }; @@ -32,7 +39,14 @@ pub fn atomic_add_u32(ptr: &mut u32, value: u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - #[cfg(not(target_arch = "spirv"))] + // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy + // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. + #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] + { + use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; + unsafe { DeviceAtomicU32::from_ptr(ptr) }.fetch_add(value, AtomicOrdering::Relaxed) + } + #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -70,7 +84,14 @@ pub fn atomic_max_u32(ptr: &mut u32, value: u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - #[cfg(not(target_arch = "spirv"))] + // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy + // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. + #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] + { + use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; + unsafe { DeviceAtomicU32::from_ptr(ptr) }.fetch_max(value, AtomicOrdering::Relaxed) + } + #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -89,7 +110,14 @@ pub fn atomic_min_u32(ptr: &mut u32, value: u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - #[cfg(not(target_arch = "spirv"))] + // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy + // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. + #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] + { + use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; + unsafe { DeviceAtomicU32::from_ptr(ptr) }.fetch_min(value, AtomicOrdering::Relaxed) + } + #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -165,7 +193,14 @@ pub fn atomic_exchange_u32(ptr: &mut u32, value: u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - #[cfg(not(target_arch = "spirv"))] + // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy + // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. + #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] + { + use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; + unsafe { DeviceAtomicU32::from_ptr(ptr) }.swap(value, AtomicOrdering::Relaxed) + } + #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -208,7 +243,14 @@ pub fn atomic_or_u32(ptr: &mut u32, value: u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - #[cfg(not(target_arch = "spirv"))] + // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy + // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. + #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] + { + use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; + unsafe { DeviceAtomicU32::from_ptr(ptr) }.fetch_or(value, AtomicOrdering::Relaxed) + } + #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -227,7 +269,14 @@ pub fn atomic_and_u32(ptr: &mut u32, value: u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - #[cfg(not(target_arch = "spirv"))] + // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy + // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. + #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] + { + use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; + unsafe { DeviceAtomicU32::from_ptr(ptr) }.fetch_and(value, AtomicOrdering::Relaxed) + } + #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -269,7 +318,14 @@ pub fn atomic_add_u32_workgroup(ptr: &mut u32, value: u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - #[cfg(not(target_arch = "spirv"))] + // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy + // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. + #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] + { + use cuda_device::atomic::{AtomicOrdering, BlockAtomicU32}; + unsafe { BlockAtomicU32::from_ptr(ptr) }.fetch_add(value, AtomicOrdering::Relaxed) + } + #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -288,7 +344,14 @@ pub fn atomic_max_u32_workgroup(ptr: &mut u32, value: u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - #[cfg(not(target_arch = "spirv"))] + // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy + // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. + #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] + { + use cuda_device::atomic::{AtomicOrdering, BlockAtomicU32}; + unsafe { BlockAtomicU32::from_ptr(ptr) }.fetch_max(value, AtomicOrdering::Relaxed) + } + #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; From 0aa67796ee2d8b5ed0496d44ae270a7312d56197 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Fri, 17 Jul 2026 17:07:26 +0200 Subject: [PATCH 15/33] Revert "khal-std: cuda-oxide arms for the RMW atomics (fetch_add/max/min/or/and/swap)" This reverts commit 9f15a5d4560d06ae3c6813933a9d72eb5a303e17. --- crates/khal-std/src/sync/atomics.rs | 81 ++++------------------------- 1 file changed, 9 insertions(+), 72 deletions(-) diff --git a/crates/khal-std/src/sync/atomics.rs b/crates/khal-std/src/sync/atomics.rs index f953034..078ce9a 100644 --- a/crates/khal-std/src/sync/atomics.rs +++ b/crates/khal-std/src/sync/atomics.rs @@ -13,14 +13,7 @@ pub fn atomic_add_i32(ptr: &mut i32, value: i32) -> i32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy - // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. - #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] - { - use cuda_device::atomic::{AtomicOrdering, DeviceAtomicI32}; - unsafe { DeviceAtomicI32::from_ptr(ptr) }.fetch_add(value, AtomicOrdering::Relaxed) - } - #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] + #[cfg(not(target_arch = "spirv"))] { use core::sync::atomic::{AtomicI32, Ordering}; let atomic = unsafe { &*(ptr as *mut i32 as *const AtomicI32) }; @@ -39,14 +32,7 @@ pub fn atomic_add_u32(ptr: &mut u32, value: u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy - // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. - #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] - { - use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; - unsafe { DeviceAtomicU32::from_ptr(ptr) }.fetch_add(value, AtomicOrdering::Relaxed) - } - #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] + #[cfg(not(target_arch = "spirv"))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -84,14 +70,7 @@ pub fn atomic_max_u32(ptr: &mut u32, value: u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy - // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. - #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] - { - use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; - unsafe { DeviceAtomicU32::from_ptr(ptr) }.fetch_max(value, AtomicOrdering::Relaxed) - } - #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] + #[cfg(not(target_arch = "spirv"))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -110,14 +89,7 @@ pub fn atomic_min_u32(ptr: &mut u32, value: u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy - // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. - #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] - { - use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; - unsafe { DeviceAtomicU32::from_ptr(ptr) }.fetch_min(value, AtomicOrdering::Relaxed) - } - #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] + #[cfg(not(target_arch = "spirv"))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -193,14 +165,7 @@ pub fn atomic_exchange_u32(ptr: &mut u32, value: u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy - // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. - #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] - { - use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; - unsafe { DeviceAtomicU32::from_ptr(ptr) }.swap(value, AtomicOrdering::Relaxed) - } - #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] + #[cfg(not(target_arch = "spirv"))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -243,14 +208,7 @@ pub fn atomic_or_u32(ptr: &mut u32, value: u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy - // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. - #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] - { - use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; - unsafe { DeviceAtomicU32::from_ptr(ptr) }.fetch_or(value, AtomicOrdering::Relaxed) - } - #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] + #[cfg(not(target_arch = "spirv"))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -269,14 +227,7 @@ pub fn atomic_and_u32(ptr: &mut u32, value: u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy - // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. - #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] - { - use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; - unsafe { DeviceAtomicU32::from_ptr(ptr) }.fetch_and(value, AtomicOrdering::Relaxed) - } - #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] + #[cfg(not(target_arch = "spirv"))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -318,14 +269,7 @@ pub fn atomic_add_u32_workgroup(ptr: &mut u32, value: u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy - // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. - #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] - { - use cuda_device::atomic::{AtomicOrdering, BlockAtomicU32}; - unsafe { BlockAtomicU32::from_ptr(ptr) }.fetch_add(value, AtomicOrdering::Relaxed) - } - #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] + #[cfg(not(target_arch = "spirv"))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -344,14 +288,7 @@ pub fn atomic_max_u32_workgroup(ptr: &mut u32, value: u32) -> u32 { { spirv_std::memory::Semantics::NONE.bits() }, >(ptr, value) } - // cuda-oxide: core RMW atomics lower to AtomicRmwOp which the legacy - // NVVM path rejects; cuda_device atomics lower to atom.* intrinsics. - #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] - { - use cuda_device::atomic::{AtomicOrdering, BlockAtomicU32}; - unsafe { BlockAtomicU32::from_ptr(ptr) }.fetch_max(value, AtomicOrdering::Relaxed) - } - #[cfg(not(any(target_arch = "spirv", feature = "cuda-oxide")))] + #[cfg(not(target_arch = "spirv"))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; From 876afad09300e49fc0091e29f750feb02e7abf89 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Fri, 17 Jul 2026 17:09:31 +0200 Subject: [PATCH 16/33] khal-std: legacy-backend-compatible barrier (llvm.nvvm.barrier0 extern) Claude-Session: https://claude.ai/code/session_01U2n9RqmxTJb8UG5d1Sjw4W --- crates/khal-std/src/sync/barriers.rs | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/crates/khal-std/src/sync/barriers.rs b/crates/khal-std/src/sync/barriers.rs index 1b3c920..52f03d8 100644 --- a/crates/khal-std/src/sync/barriers.rs +++ b/crates/khal-std/src/sync/barriers.rs @@ -12,7 +12,26 @@ pub fn workgroup_memory_barrier_with_group_sync() { // MIR importer on any compilation target and lowers to the convergent // NVVM barrier op (`bar.sync`), which optimization passes will not // tail-duplicate into divergent branches. - #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] + // Device build (cuda-oxide, real nvptx64 target): raw intrinsic extern — + // matched by name by both backend lines ("llvm.nvvm.barrier0"). + #[cfg(all(feature = "cuda-oxide", target_arch = "nvptx64"))] + { + unsafe extern "C" { + #[link_name = "llvm.nvvm.barrier0"] + fn nvvm_barrier0(); + } + unsafe { + nvvm_barrier0(); + } + } + // Unified-compilation host build (cuda-oxide feature on a host target): + // the MIR importer intercepts this call by name; the host stub is never + // executed natively. + #[cfg(all( + feature = "cuda-oxide", + not(target_arch = "spirv"), + not(target_arch = "nvptx64") + ))] { cuda_device::sync_threads(); } From 12b7f7d28e9cdc42bd0f124cdde69edbb628632a Mon Sep 17 00:00:00 2001 From: haixuantao Date: Tue, 21 Jul 2026 23:36:39 +0200 Subject: [PATCH 17/33] fix(builder): rerun-if-changed on prebuilt PTX/cubin path The env short-circuit declared rerun-if-env-changed only, so replacing the cubin file at the same path silently kept the stale embedded shaders. Claude-Session: https://claude.ai/code/session_01B7NC7U2wDx2tSmT9mkeF9h --- crates/khal-builder/src/lib.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/crates/khal-builder/src/lib.rs b/crates/khal-builder/src/lib.rs index 52ffa58..4b4615f 100644 --- a/crates/khal-builder/src/lib.rs +++ b/crates/khal-builder/src/lib.rs @@ -194,6 +194,8 @@ impl KhalBuilder { let env_key = format!("CUDA_OXIDE_SHADERS_PTX_{crate_name}"); println!("cargo:rerun-if-env-changed={env_key}"); if let Some(prebuilt) = std::env::var_os(&env_key) { + // Re-embed when the prebuilt file itself changes, not only its path. + println!("cargo:rerun-if-changed={}", Path::new(&prebuilt).display()); let dst = output_dir.join("shaders.ptx"); std::fs::copy(&prebuilt, &dst).unwrap_or_else(|e| { panic!("failed to copy prebuilt PTX {prebuilt:?} ({env_key}) to {dst:?}: {e}") From cc7682cad284169622d1b57fa03ec7278704d004 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Wed, 22 Jul 2026 11:58:38 +0200 Subject: [PATCH 18/33] debug(cuda): env-gated capture diagnostics KHAL_CUDA_GRAPH_DOT= dumps the captured graph via cuGraphDebugDotPrint (node-type census: MEM_ALLOC nodes explain relaunch INVALID_VALUE). KHAL_CUDA_ALLOC_TRACE=1 backtraces any buffer allocation recorded while the stream is capturing. Both no-ops unless the env vars are set. Claude-Session: https://claude.ai/code/session_01B7NC7U2wDx2tSmT9mkeF9h --- crates/khal/src/backend/cuda.rs | 39 +++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/crates/khal/src/backend/cuda.rs b/crates/khal/src/backend/cuda.rs index 568fcb2..3c353e0 100644 --- a/crates/khal/src/backend/cuda.rs +++ b/crates/khal/src/backend/cuda.rs @@ -114,6 +114,18 @@ impl Cuda { .stream .end_capture(no_flags)? .ok_or(CudaBackendError::CaptureFailed)?; + // Debug: dump the captured graph's node structure to a DOT file + // (KHAL_CUDA_GRAPH_DOT=/path/out.dot). Invaluable for diagnosing + // relaunch failures — e.g. a stray mem-alloc node makes cuGraphLaunch + // return CUDA_ERROR_INVALID_VALUE on every launch after the first. + if let Ok(path) = std::env::var("KHAL_CUDA_GRAPH_DOT") { + let cpath = std::ffi::CString::new(path).unwrap(); + // 1 = CU_GRAPH_DEBUG_DOT_FLAGS_VERBOSE + let res = unsafe { + cudarc::driver::sys::cuGraphDebugDotPrint(graph.cu_graph(), cpath.as_ptr(), 1) + }; + eprintln!("[khal-cuda] graph dot dump: {res:?}"); + } Ok(CapturedGraph { graph }) } } @@ -491,6 +503,18 @@ impl Backend for Cuda { }); } let bytes: &[u8] = bytemuck::cast_slice(data); + if std::env::var_os("KHAL_CUDA_ALLOC_TRACE").is_some() + && matches!( + self.stream.capture_status(), + Ok(cudarc::driver::sys::CUstreamCaptureStatus::CU_STREAM_CAPTURE_STATUS_ACTIVE) + ) + { + eprintln!( + "[khal-cuda] BUFFER_FROM_DATA DURING CAPTURE: {} bytes\n{}", + bytes.len(), + std::backtrace::Backtrace::force_capture() + ); + } let slice = self.stream.clone_htod(bytes)?; let raw_ptr = extract_raw_ptr(&slice, &self.stream); Ok(CudaBuffer { @@ -517,6 +541,21 @@ impl Backend for Cuda { }); } let byte_len = len * std::mem::size_of::(); + // Debug: catch capture-illegal allocations (KHAL_CUDA_ALLOC_TRACE=1). + // An allocation recorded while the stream is capturing becomes a + // MEM_ALLOC graph node, and a graph with un-freed alloc nodes cannot + // be relaunched (cuGraphLaunch -> CUDA_ERROR_INVALID_VALUE). + if std::env::var_os("KHAL_CUDA_ALLOC_TRACE").is_some() + && matches!( + self.stream.capture_status(), + Ok(cudarc::driver::sys::CUstreamCaptureStatus::CU_STREAM_CAPTURE_STATUS_ACTIVE) + ) + { + eprintln!( + "[khal-cuda] ALLOC DURING CAPTURE: {byte_len} bytes\n{}", + std::backtrace::Backtrace::force_capture() + ); + } let slice = self.stream.alloc_zeros::(byte_len)?; let raw_ptr = extract_raw_ptr(&slice, &self.stream); Ok(CudaBuffer { From 19a58acf01b1fce3f505cfa0cdf718c5c98dfb78 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Mon, 27 Jul 2026 19:17:31 +0200 Subject: [PATCH 19/33] khal-std: pin cuda-device to upstream NVlabs cuda-oxide The fork pin predates the unified conversion; the API surface consumed here (kernel macro, SharedArray, sync_threads) is upstream. Consuming workspaces still [patch] to a local checkout while two fixes (NVlabs #518, #519) are pending upstream. --- crates/khal-std/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/khal-std/Cargo.toml b/crates/khal-std/Cargo.toml index 59d047d..c517c5a 100644 --- a/crates/khal-std/Cargo.toml +++ b/crates/khal-std/Cargo.toml @@ -60,7 +60,7 @@ num-traits = { version = "0.2", default-features = false, features = ["libm"], o # portability the feature comment above promises. cuda-device is not published # on crates.io, so a git rev is the portable way to express it. Override with a # local checkout via [patch] in the consuming workspace when hacking on both. -cuda-device = { git = "https://github.com/haixuanTao/cuda-oxide.git", rev = "8c3b5e2dabad165f0f82d98d72dcf2cd43c50d2d", optional = true } +cuda-device = { git = "https://github.com/NVlabs/cuda-oxide.git", rev = "3655fa1246fb6c55c1d1ad0f3d2ce82e719ed8c3", optional = true } [lints] workspace = true From 3187c3f8ca528d209da70ca6f31e0d856d81c725 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Mon, 27 Jul 2026 20:48:54 +0200 Subject: [PATCH 20/33] docs: the unified host-target path is production; cuda-device comes from NVlabs --- CUDA_OXIDE_AUTOPORT.md | 12 ++++++++++++ crates/khal-std/Cargo.toml | 18 +++++++++--------- 2 files changed, 21 insertions(+), 9 deletions(-) diff --git a/CUDA_OXIDE_AUTOPORT.md b/CUDA_OXIDE_AUTOPORT.md index 2f4efea..779663f 100644 --- a/CUDA_OXIDE_AUTOPORT.md +++ b/CUDA_OXIDE_AUTOPORT.md @@ -1,5 +1,17 @@ # Running rust-gpu `#[spirv]` shaders on native CUDA via cuda-oxide +> **Status update (2026-07-27): the unified host-target path IS production.** +> The device-target mode described in the update below lived on a cuda-oxide +> fork; the fork is retired and shader crates now compile through **upstream +> NVlabs cuda-oxide** on the host target (standard interception): the importer +> recognizes the `cuda_device` API by name on any target, so the barrier and +> sreg shims are gone — the remaining host-target requirement is building +> glamx with `scalar-math` (no SSE paths in device code). Kernel entries are +> emitted by `spirv_bindgen` with the reserved kernel-prefix marker. The +> production chain (nexus + vortx -> sm_120 cubins) is +> `zealot/scripts/full_unified_chain.sh`; validated physics-equivalent to the +> device-mode builds it replaces (G1 standing gate + wall-time parity). +> > Status: working prototype (2026-06-05). `gpu_add` and `gemm_tiled` run **bit-exact** > on the RTX 5090 (sm_120) through khal's CUDA backend, compiled by cuda-oxide. > A `#[spirv_kernel]` proc-macro auto-lowers the *verbatim* shader source. diff --git a/crates/khal-std/Cargo.toml b/crates/khal-std/Cargo.toml index c517c5a..a8add2e 100644 --- a/crates/khal-std/Cargo.toml +++ b/crates/khal-std/Cargo.toml @@ -28,16 +28,16 @@ unsafe_remove_boundchecks = [] # --- nvptx64 device backend selection --- # `rust-cuda` (default): compile device code with Rust-CUDA's rustc_codegen_nvvm, # pulling `cuda_std` for thread/grid intrinsics and `GpuFloat` math. -# `cuda-oxide`: compile device code with the cuda-oxide PTX backend instead. -# cuda_std does not build under cuda-oxide's `-Zbuild-std`, so this path uses -# khal_std's own `llvm.nvvm.*` intrinsic declarations (lowered by cuda-oxide) -# and a libm-backed `Float`. Build shader crates with -# `default-features = false, features = ["cuda-oxide", ...]`. +# `cuda-oxide`: compile device code with the cuda-oxide backend instead. +# Thread/grid/barrier primitives route through the `cuda_device` API, which +# the cuda-oxide importer recognizes by name on ANY compilation target - so +# shader crates build as ordinary host-target crates (unified interception) +# with no nvptx64 cross-target step; `Float` is libm-backed. Build shader +# crates with `default-features = false, features = ["cuda-oxide", ...]`. rust-cuda = ["dep:cuda_std"] -# NOTE: the cuda-oxide backend (device-only nvptx build) also needs a -# `cuda-device` dependency, but that crate lives in the cuda-oxide tree and is -# only present on the build box; it is added there. This manifest keeps the -# feature host-buildable everywhere else. +# The `cuda-device` API-surface crate resolves from upstream NVlabs +# cuda-oxide (git pin below); consuming workspaces may [patch] it to a local +# checkout when hacking on both. cuda-oxide = ["dep:num-traits", "dep:cuda-device"] [dependencies] From a70c8f70d33679b5b7bb760755f0506a6c947a0e Mon Sep 17 00:00:00 2001 From: haixuantao Date: Tue, 28 Jul 2026 10:23:53 +0200 Subject: [PATCH 21/33] khal-std: bump cuda-device to NVlabs main with our merged fixes (#518, #520) --- crates/khal-std/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/khal-std/Cargo.toml b/crates/khal-std/Cargo.toml index a8add2e..c52742b 100644 --- a/crates/khal-std/Cargo.toml +++ b/crates/khal-std/Cargo.toml @@ -60,7 +60,7 @@ num-traits = { version = "0.2", default-features = false, features = ["libm"], o # portability the feature comment above promises. cuda-device is not published # on crates.io, so a git rev is the portable way to express it. Override with a # local checkout via [patch] in the consuming workspace when hacking on both. -cuda-device = { git = "https://github.com/NVlabs/cuda-oxide.git", rev = "3655fa1246fb6c55c1d1ad0f3d2ce82e719ed8c3", optional = true } +cuda-device = { git = "https://github.com/NVlabs/cuda-oxide.git", rev = "62472763", optional = true } [lints] workspace = true From bf7ba72f992e527fc9b2e408451a158afe135519 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Tue, 28 Jul 2026 11:42:48 +0200 Subject: [PATCH 22/33] khal-builder: create the shader output dir on the prebuilt-cubin path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The early-return prebuilt branch copied into OUT_DIR/shaders-spirv without creating it — masked whenever the SPIR-V build ran first, breaks cuda-feature builds that hit the prebuilt path before any spirv step. --- crates/khal-builder/src/lib.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/crates/khal-builder/src/lib.rs b/crates/khal-builder/src/lib.rs index 4b4615f..b47a4d0 100644 --- a/crates/khal-builder/src/lib.rs +++ b/crates/khal-builder/src/lib.rs @@ -196,6 +196,8 @@ impl KhalBuilder { if let Some(prebuilt) = std::env::var_os(&env_key) { // Re-embed when the prebuilt file itself changes, not only its path. println!("cargo:rerun-if-changed={}", Path::new(&prebuilt).display()); + std::fs::create_dir_all(output_dir) + .expect("failed to create shader output dir for the prebuilt cubin"); let dst = output_dir.join("shaders.ptx"); std::fs::copy(&prebuilt, &dst).unwrap_or_else(|e| { panic!("failed to copy prebuilt PTX {prebuilt:?} ({env_key}) to {dst:?}: {e}") From 523cea51ba69d609a43f739a136ecb5f8bd011ce Mon Sep 17 00:00:00 2001 From: haixuantao Date: Tue, 4 Aug 2026 15:31:38 +0200 Subject: [PATCH 23/33] webgpu: wasm perf counters at the browser boundary (+ pub webgpu module) Counts submits, real pass opens (merged-cache hits excluded), encoder-level copies, write_buffer uploads and map requests, behind take-and-zero accessors. This is how the zealot demo's "is there a hidden readback" question got answered with numbers (steady state: copies=0 writes=0 maps=0, ~450 submits/s) instead of a guess. Rides on the existing (uncommitted until now) wasm pass-merging in this file. --- crates/khal/src/backend/mod.rs | 3 +- crates/khal/src/backend/webgpu.rs | 112 +++++++++++++++++++++++++++++- 2 files changed, 111 insertions(+), 4 deletions(-) diff --git a/crates/khal/src/backend/mod.rs b/crates/khal/src/backend/mod.rs index a7e051d..d3b07bc 100644 --- a/crates/khal/src/backend/mod.rs +++ b/crates/khal/src/backend/mod.rs @@ -17,7 +17,8 @@ pub enum CompileTarget { #[cfg(feature = "webgpu")] pub use webgpu::WebGpu; #[cfg(feature = "webgpu")] -mod webgpu; +#[cfg_attr(target_arch = "wasm32", allow(unused))] +pub mod webgpu; #[cfg(feature = "cuda")] pub use cuda::Cuda; diff --git a/crates/khal/src/backend/webgpu.rs b/crates/khal/src/backend/webgpu.rs index fae5ba1..13f88b8 100644 --- a/crates/khal/src/backend/webgpu.rs +++ b/crates/khal/src/backend/webgpu.rs @@ -78,15 +78,32 @@ pub struct WebGpuFunction { } /// A WebGPU compute pass that carries its own device reference. +/// +/// On wasm the pass may be a MERGED pass shared by successive `begin_pass` +/// calls (see [`WebGpuEncoder::open_pass`]); dropping the handle then returns +/// the underlying pass to the encoder's cache instead of ending it. pub struct WebGpuPass { - pub(crate) pass: ComputePass<'static>, + pub(crate) pass: Option>, pub(crate) device: Device, + /// wasm pass-merging: the encoder cache slot this pass returns to on drop + /// (`None` = plain pass, ends on drop as before). + #[cfg(target_arch = "wasm32")] + pub(crate) home: Option>>>>, } impl WebGpuPass { /// Begins a compute dispatch within this pass, binding the given function. pub fn begin_dispatch<'a>(&'a mut self, function: &'a WebGpuFunction) -> WebGpuDispatch<'a> { - WebGpuDispatch::new(&self.device, &mut self.pass, function) + WebGpuDispatch::new(&self.device, self.pass.as_mut().expect("pass ended"), function) + } +} + +#[cfg(target_arch = "wasm32")] +impl Drop for WebGpuPass { + fn drop(&mut self) { + if let Some(home) = self.home.take() { + *home.borrow_mut() = self.pass.take(); + } } } @@ -94,6 +111,51 @@ impl WebGpuPass { pub struct WebGpuEncoder { pub(crate) encoder: CommandEncoder, pub(crate) device: Device, + /// wasm pass-merging: the open compute pass shared by successive + /// `begin_pass` calls. The browser's WebGPU spec guarantees storage-write + /// visibility between dispatches within one pass, and a Metal pass switch + /// costs ~60 µs — per-kernel passes put a ~15 ms/step floor under the + /// browser physics demo (~206 passes/step); merging removes it. Closed + /// (ended) by buffer copies and at submit. + #[cfg(target_arch = "wasm32")] + pub(crate) open_pass: std::rc::Rc>>>, +} + +#[cfg(target_arch = "wasm32")] +impl WebGpuEncoder { + /// End the cached merged pass (if any) — required before encoder-level + /// commands (copies) and before `finish`. + fn close_open_pass(&mut self) { + drop(self.open_pass.borrow_mut().take()); + } +} + +/// wasm perf counters: every operation that crosses the wasm→browser GPU +/// boundary, for the zealot demo's per-second HUD. Reading them answers +/// "where do the crossings go" without a profiler build. +#[cfg(target_arch = "wasm32")] +pub mod perf_counters { + use core::sync::atomic::{AtomicU32, Ordering::Relaxed}; + pub static SUBMITS: AtomicU32 = AtomicU32::new(0); + /// Compute passes actually OPENED (merged-cache hits don't count). + pub static PASSES: AtomicU32 = AtomicU32::new(0); + /// Encoder-level buffer copies — each one closes the merged pass. + pub static COPIES: AtomicU32 = AtomicU32::new(0); + /// queue.write_buffer uploads. + pub static WRITES: AtomicU32 = AtomicU32::new(0); + /// Buffer map requests (readbacks). + pub static MAPS: AtomicU32 = AtomicU32::new(0); + + /// Take-and-zero all counters: (submits, passes, copies, writes, maps). + pub fn take() -> (u32, u32, u32, u32, u32) { + ( + SUBMITS.swap(0, Relaxed), + PASSES.swap(0, Relaxed), + COPIES.swap(0, Relaxed), + WRITES.swap(0, Relaxed), + MAPS.swap(0, Relaxed), + ) + } } /// Helper struct to initialize a device and its queue. @@ -619,6 +681,8 @@ impl Backend for WebGpu { .device .create_command_encoder(&CommandEncoderDescriptor::default()), device: self.device.clone(), + #[cfg(target_arch = "wasm32")] + open_pass: Default::default(), } } @@ -631,6 +695,12 @@ impl Backend for WebGpu { } fn submit(&self, encoder: Self::Encoder) -> Result<(), Self::Error> { + #[cfg(target_arch = "wasm32")] + let mut encoder = encoder; + #[cfg(target_arch = "wasm32")] + encoder.close_open_pass(); + #[cfg(target_arch = "wasm32")] + perf_counters::SUBMITS.fetch_add(1, core::sync::atomic::Ordering::Relaxed); let _ = self.queue.submit(Some(encoder.encoder.finish())); Ok(()) } @@ -686,6 +756,8 @@ impl Backend for WebGpu { offset: u64, data: &[T], ) -> Result<(), Self::Error> { + #[cfg(target_arch = "wasm32")] + perf_counters::WRITES.fetch_add(1, core::sync::atomic::Ordering::Relaxed); let elt_sz = std::mem::size_of::() as u64; self.queue .write_buffer(buffer, offset * elt_sz, bytemuck::cast_slice(data)); @@ -767,6 +839,7 @@ impl Encoder for WebGpuEncoder { timestamp_writes: None, }; + let mut wants_timestamps = false; if let Some(timestamps) = timestamps && let Some((begin_idx, end_idx)) = timestamps.alloc_timestamp_pair(label.to_string()) { @@ -775,11 +848,35 @@ impl Encoder for WebGpuEncoder { beginning_of_pass_write_index: Some(begin_idx), end_of_pass_write_index: Some(end_idx), }); + wants_timestamps = true; + } + + // wasm pass-merging: reuse the cached open pass (per-pass timestamps + // force a dedicated pass, so those bypass the cache). + #[cfg(target_arch = "wasm32")] + if !wants_timestamps { + let cached = self.open_pass.borrow_mut().take(); + let pass = cached.unwrap_or_else(|| { + perf_counters::PASSES.fetch_add(1, core::sync::atomic::Ordering::Relaxed); + self.encoder.begin_compute_pass(&desc).forget_lifetime() + }); + return WebGpuPass { + pass: Some(pass), + device: self.device.clone(), + home: Some(self.open_pass.clone()), + }; + } + #[cfg(target_arch = "wasm32")] + if wants_timestamps { + self.close_open_pass(); } + let _ = wants_timestamps; WebGpuPass { - pass: self.encoder.begin_compute_pass(&desc).forget_lifetime(), + pass: Some(self.encoder.begin_compute_pass(&desc).forget_lifetime()), device: self.device.clone(), + #[cfg(target_arch = "wasm32")] + home: None, } } @@ -791,6 +888,11 @@ impl Encoder for WebGpuEncoder { target_offset: usize, copy_len: usize, ) -> Result<(), WebGpuBackendError> { + // Copies are encoder-level commands: the merged pass must end first. + #[cfg(target_arch = "wasm32")] + self.close_open_pass(); + #[cfg(target_arch = "wasm32")] + perf_counters::COPIES.fetch_add(1, core::sync::atomic::Ordering::Relaxed); wgpu::CommandEncoder::copy_buffer_to_buffer( &mut self.encoder, source, @@ -976,6 +1078,8 @@ async fn read_bytes(device: &Device, buffer: &Buffer) -> Result Result Date: Mon, 31 Aug 2026 18:36:33 +0800 Subject: [PATCH 24/33] metal: fix write_buffer/GPU race + WebGPU-parity shader semantics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The native Metal backend produced wrong physics on multi-env scenes: the zealot biped ran every env with the LAST env's robot state. Root cause (proven by bisection down to a 4-env repro): `write_buffer` memcpys into shared memory immediately, racing in-flight kernel reads — wgpu's `queue.write_buffer` is ordered after submitted work. Staged per-env uploads corrupted accordingly. - write_buffer: while work may be in flight (`gpu_busy`, set on submit), route bytes through a staging buffer + ONE lazily-committed blit command buffer per write burst (`pending_writes`, flushed by the next submit/synchronize). Queue FIFO order = wgpu's semantics, no CPU stall. Direct-memcpy fast path only before any submit. - MSL translation now matches wgpu's WebGPU semantics: bounds checks ReadZeroSkipWrite (kernels rely on zero-read/skip-write for by-design OOB accesses), zero-initialized workgroup memory, force_loop_bounding, fast math OFF, MSL 3.0. - uninit_buffer zero-fills (WebGPU guarantees zeroed buffers). - KHAL_BACKEND=metal|msl now selects the backend in GpuBackend::auto (never auto-picked; errors when built without the feature). - Warn when a pipeline's maxTotalThreadsPerThreadgroup is below the shader's declared workgroup size (dispatching more is Metal UB). - KHAL_METAL_WGSL_ROUNDTRIP=1 debug lever: normalize IR through a WGSL round-trip before MSL generation (entry-point names restored). - khal-example: grid_probe/batch_probe 2D-dispatch validation kernels. Verified against WebGPU on the zealot biped (M4 Pro): bare-physics and multi-env terrain contact sets now match; zero-action survival curves identical (61% @1s). Perf at 1024 envs: 8.7 s/iter vs WebGPU 8.2 — the previously observed 6.5 was the race illegally overlapping uploads with compute. gpuwait dominates (270 vs 76 ms/step); closing that (concurrent dispatch + explicit hazards) is future work. Claude-Session: https://claude.ai/code/session_01EvbTzx6ZJNtcAKwVQ1B68i --- crates/khal-example-shaders/src/lib.rs | 41 +++++ crates/khal-example/Cargo.toml | 1 + .../shaders-spirv/batch_probe.spv | Bin 0 -> 2244 bytes .../khal-example/shaders-spirv/grid_probe.spv | Bin 0 -> 1204 bytes .../khal-example/shaders-spirv/manifest.json | 10 ++ crates/khal-example/src/main.rs | 115 ++++++++++---- crates/khal/Cargo.toml | 2 +- crates/khal/src/backend/any_backend.rs | 19 +++ crates/khal/src/backend/metal.rs | 142 +++++++++++++++++- 9 files changed, 296 insertions(+), 34 deletions(-) create mode 100644 crates/khal-example/shaders-spirv/batch_probe.spv create mode 100644 crates/khal-example/shaders-spirv/grid_probe.spv diff --git a/crates/khal-example-shaders/src/lib.rs b/crates/khal-example-shaders/src/lib.rs index 167a116..6edb7b6 100644 --- a/crates/khal-example-shaders/src/lib.rs +++ b/crates/khal-example-shaders/src/lib.rs @@ -16,3 +16,44 @@ pub fn add_assign( *a.at_mut(thread_id) += b.read(thread_id); } } + +/// 2D-grid probe: writes a tag encoding (batch=y, x, src[y]) per slot. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn grid_probe( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] src: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] out: &mut [u32], +) { + let x = invocation_id.x; + let y = invocation_id.y; + let cap = 64u32; + if x >= cap || y as usize >= src.len() { + return; + } + out.write((y * cap + x) as usize, (y << 20) | (src.read(y as usize) << 10) | x); +} + +/// Batched StepRng probe mimicking the nexus narrow-phase pattern: +/// per-batch len + strided loop + batch-sliced reads. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn batch_probe( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(num_workgroups)] num_workgroups: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] lens: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] src: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] out: &mut [u32], +) { + let num_threads = num_workgroups.x * 64; + let batch = invocation_id.y; + if batch as usize >= lens.len() { + return; + } + let cap = 64u32; + let len = lens.read(batch as usize).min(cap); + for i in khal_std::iter::StepRng::new(invocation_id.x..len, num_threads) { + let v = src.read((batch * cap + i) as usize); + out.write((batch * cap + i) as usize, (batch << 20) | v); + } +} diff --git a/crates/khal-example/Cargo.toml b/crates/khal-example/Cargo.toml index 0fc1cb2..09bcc2d 100644 --- a/crates/khal-example/Cargo.toml +++ b/crates/khal-example/Cargo.toml @@ -7,6 +7,7 @@ publish = false [features] default = ["webgpu"] webgpu = [ "khal/webgpu" ] +metal = [ "khal/metal" ] cpu = ["khal/cpu", "khal-example-shaders/cpu"] cpu-parallel = ["cpu", "khal-example-shaders/cpu-parallel"] cuda = ["khal/cuda", "khal-builder/cuda", "khal-example-shaders/cuda"] diff --git a/crates/khal-example/shaders-spirv/batch_probe.spv b/crates/khal-example/shaders-spirv/batch_probe.spv new file mode 100644 index 0000000000000000000000000000000000000000..083223f8c45ae4a4137d9ecdbcd6d7774b84de51 GIT binary patch literal 2244 zcmZvc>q?bT5Qdk!)k1M=YIf4LNrZ<^6efhRGfO*>+WA1DVJCv36b$;YgQ#nXAi9XU zkS@ZG1U=8UzoG9(1GDFySu?ZVHEVslsd1<^H8-S`%IO`yN>^%3MzmjhQtC)eS=`@p zYD>*n*O!-8?>$_6_WbGc3YjU(Dq&r3K|A@a_erUWFIHz#Da&sI@5=oStToop(5Z~I z=PcdC%23^8^>VR(8*09FbzI8uAim;WuIi=&!#Mm4o%(v-S24 zKJqtT-~P_ZobbJiqVMeW2gr2p_EH<*s^5;oOHtyX--WXUOp5r5Vo^?NJ-2FYv0j%;1Hm)5+522n%{1a=A zzVT0apTpo{pCkB+eO%K&idrY0;uOAxany4=H|w2`^EigO7X2r{p2hfkYR!3z7UOZw zN#dSUnIpFyTSJXo<7_-f6)PWk>3fg*#y3)HR(BFyL+hMVItBg)$7!rsvGz39T5a$+ z1GZj?*bLZM@eXG3&2`N?IEz)BoonaN^QisB3!KS!egSoEzIR0))|zL2?0FIFo{eZN ze;=2??&+HSoY5R=JUGvh>pt_S-;DXiJ=`niGzGR!e)e&H*Q{SaJ=oK^Y8@-B}`yqdC*Ya2e>&&iW6=&w!4fH1JEX46$+yalGzNMJOZLo3ojrVp3 zY}{|q{&%s8{aw>vMC~t*_q+u5oAq9-kN3O`cI|UM{}r%2{BQh{$L@oTn;SeHfbHXV zFONr9MINr{KSt#tAMeMvsqb3wdID}igVzaq8FNo-oS!xJD)w-O#dY)aef#EXW9C(G x?5*8Mo3sCz*Wh(jep>ffL#_2pFR_Y!%=J4oW`C{w#ea@g-yGMqe|RYq=wA{7oA3Yt literal 0 HcmV?d00001 diff --git a/crates/khal-example/shaders-spirv/grid_probe.spv b/crates/khal-example/shaders-spirv/grid_probe.spv new file mode 100644 index 0000000000000000000000000000000000000000..05299593b7c290bf7c405f42610f69f688c635ef GIT binary patch literal 1204 zcmYk5Sx!_@5QYnAMwM=v+8AkX+aQAw8BBwsh8Pokbp<3uLU@Ao3V2{W#6`eDScMEB z@%t|4kXuQ0jpwhw&bd9%-x-dgp0V>YzOlTGGyVkrF6s`G*aFS-jd`+C#h=Fatg}6D zeFk->E_ZCrvx!YS){M7MxyjM{ q@%u2gmbre$_pRkB-gmIfIkS4FVpTD`YvTedrnXvl)&pP0Hu@ikBSnz_ literal 0 HcmV?d00001 diff --git a/crates/khal-example/shaders-spirv/manifest.json b/crates/khal-example/shaders-spirv/manifest.json index f61dc1a..f3577ed 100644 --- a/crates/khal-example/shaders-spirv/manifest.json +++ b/crates/khal-example/shaders-spirv/manifest.json @@ -3,5 +3,15 @@ "source_path": "../khal-example/shaders-spirv/add_assign.spv", "entry_point": "add_assign", "wgsl_entry_point": "add_assign" + }, + { + "source_path": "../khal-example/shaders-spirv/batch_probe.spv", + "entry_point": "batch_probe", + "wgsl_entry_point": "batch_probe" + }, + { + "source_path": "../khal-example/shaders-spirv/grid_probe.spv", + "entry_point": "grid_probe", + "wgsl_entry_point": "grid_probe" } ] \ No newline at end of file diff --git a/crates/khal-example/src/main.rs b/crates/khal-example/src/main.rs index 70d36fe..8616a0a 100644 --- a/crates/khal-example/src/main.rs +++ b/crates/khal-example/src/main.rs @@ -1,44 +1,103 @@ -use khal::backend::{Backend, Buffer, Encoder, GpuBackend, GpuBackendError, WebGpu}; +use khal::backend::{Backend, DispatchGrid, Encoder, GpuBackend}; use khal::re_exports::include_dir::{Dir, include_dir}; use khal::{BufferUsages, Shader}; -use khal_example_shaders::AddAssign; static SPIRV_DIR: Dir<'static> = include_dir!("$CARGO_MANIFEST_DIR/shaders-spirv"); +use khal_example_shaders::{AddAssign, BatchProbe, GridProbe}; #[derive(Shader)] pub struct GpuKernels { + #[allow(dead_code)] add_assign: AddAssign, + grid_probe: GridProbe, + batch_probe: BatchProbe, } #[async_std::main] async fn main() { - let webgpu = WebGpu::default().await.unwrap(); - let backend = GpuBackend::WebGpu(webgpu); + let backend = GpuBackend::auto(Default::default(), Default::default()) + .await + .unwrap(); + let kernels = GpuKernels::from_backend(&backend).unwrap(); + let nb = 4u32; + let cap = 64u32; - // Run the operation and display the result. - let a = (0..10000).map(|i| i as f32).collect::>(); - let b = (0..10000).map(|i| i as f32 * 10.0).collect::>(); - let result = compute_sum(&backend, &a, &b).await.unwrap(); - println!("Computed sum: {result:?}"); -} - -async fn compute_sum( - backend: &GpuBackend, - a: &[f32], - b: &[f32], -) -> Result, GpuBackendError> { - // Generate the GPU buffers. - let mut a = backend.init_buffer(a, BufferUsages::STORAGE | BufferUsages::COPY_SRC)?; - let b = backend.init_buffer(b, BufferUsages::STORAGE)?; - - // Dispatch the operation on the gpu. - let kernels = GpuKernels::from_backend(backend)?; - let mut encoder = backend.begin_encoding(); - let mut pass = encoder.begin_pass("add_assign", None); - kernels.add_assign.call(&mut pass, a.len(), &mut a, &b)?; + // grid_probe: bare thread-count dispatch [cap, nb, 1]. + let src: Vec = (0..nb).map(|b| 100 + b).collect(); + let src_buf = backend.init_buffer(&src, BufferUsages::STORAGE).unwrap(); + let mut out = backend + .init_buffer( + &vec![0u32; (cap * nb) as usize], + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + ) + .unwrap(); + let mut enc = backend.begin_encoding(); + let mut pass = enc.begin_pass("probe", None); + kernels + .grid_probe + .call(&mut pass, [cap, nb, 1], &src_buf, &mut out) + .unwrap(); drop(pass); - backend.submit(encoder)?; + backend.submit(enc).unwrap(); + let got: Vec = backend.slow_read_vec(&out).await.unwrap(); + let mut bad = 0; + for y in 0..nb { + for x in 0..cap { + let want = (y << 20) | ((100 + y) << 10) | x; + if got[(y * cap + x) as usize] != want { + if bad < 5 { + println!( + "grid_probe MISMATCH y={y} x={x} got={:#x} want={want:#x}", + got[(y * cap + x) as usize] + ); + } + bad += 1; + } + } + } + println!("grid_probe (threads form): {} mismatches", bad); - // Read the result (slower but convenient version). - backend.slow_read_vec(&a).await + // batch_probe: raw workgroup Grid dispatch [1, nb, 1] like nexus fixed grids. + let lens: Vec = (0..nb).map(|b| 10 + b).collect(); + let lens_buf = backend.init_buffer(&lens, BufferUsages::STORAGE).unwrap(); + let src2: Vec = (0..cap * nb).collect(); + let src2_buf = backend.init_buffer(&src2, BufferUsages::STORAGE).unwrap(); + let mut out2 = backend + .init_buffer( + &vec![u32::MAX; (cap * nb) as usize], + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + ) + .unwrap(); + let mut enc = backend.begin_encoding(); + let mut pass = enc.begin_pass("probe2", None); + kernels + .batch_probe + .call( + &mut pass, + DispatchGrid::Grid([1, nb, 1]), + &lens_buf, + &src2_buf, + &mut out2, + ) + .unwrap(); + drop(pass); + backend.submit(enc).unwrap(); + let got2: Vec = backend.slow_read_vec(&out2).await.unwrap(); + let mut bad2 = 0; + for b in 0..nb { + for i in 0..lens[b as usize] { + let idx = (b * cap + i) as usize; + let want = (b << 20) | src2[idx]; + if got2[idx] != want { + if bad2 < 5 { + println!( + "batch_probe MISMATCH b={b} i={i} got={:#x} want={want:#x}", + got2[idx] + ); + } + bad2 += 1; + } + } + } + println!("batch_probe (Grid form): {} mismatches", bad2); } diff --git a/crates/khal/Cargo.toml b/crates/khal/Cargo.toml index 9caddf5..738d101 100644 --- a/crates/khal/Cargo.toml +++ b/crates/khal/Cargo.toml @@ -39,7 +39,7 @@ paste = "1" [target.'cfg(target_os = "macos")'.dependencies] metal = { version = "0.33", optional = true } -naga = { version = "30", optional = true, features = ["spv-in", "msl-out"] } +naga = { version = "30", optional = true, features = ["spv-in", "msl-out", "wgsl-in", "wgsl-out"] } # Objective-C block support for Metal command-buffer completion handlers # (non-blocking readback completion). Matches metal-rs's own `block` dep. block = { version = "0.1.6", optional = true } diff --git a/crates/khal/src/backend/any_backend.rs b/crates/khal/src/backend/any_backend.rs index f652024..908323e 100644 --- a/crates/khal/src/backend/any_backend.rs +++ b/crates/khal/src/backend/any_backend.rs @@ -74,6 +74,25 @@ impl GpuBackend { /// the `cuda` feature is compiled, else WebGPU. Override with /// `KHAL_BACKEND=cuda|webgpu`. pub async fn auto(features: wgpu::Features, limits: wgpu::Limits) -> anyhow::Result { + // Explicit native-Metal request (macOS): bypasses wgpu entirely — + // same SPIR-V, same naga MSL translation, but no wgpu state + // tracking / implicit-barrier layer. Never auto-selected: the + // WebGPU path stays the default on Metal targets. + if matches!( + std::env::var("KHAL_BACKEND").ok().as_deref(), + Some("metal" | "Metal" | "msl") + ) { + #[cfg(feature = "metal")] + { + eprintln!("[khal] backend = native Metal"); + return Ok(Self::Metal(Metal::new()?)); + } + #[cfg(not(feature = "metal"))] + anyhow::bail!( + "KHAL_BACKEND=metal requested but khal was built without the \ + 'metal' feature" + ); + } // Resolve any explicit override: Some(true)=force CUDA, Some(false)=force // WebGPU, None=auto-detect. let want_cuda = match std::env::var("KHAL_BACKEND").ok().as_deref() { diff --git a/crates/khal/src/backend/metal.rs b/crates/khal/src/backend/metal.rs index 1efeb67..100813f 100644 --- a/crates/khal/src/backend/metal.rs +++ b/crates/khal/src/backend/metal.rs @@ -42,6 +42,22 @@ pub struct Metal { /// Capabilities for GPU timestamp queries; `None` if the device or /// driver doesn't expose stage-boundary timestamp sampling. timing_caps: Option>, + /// Set by `submit`, cleared by `synchronize`. While set, `write_buffer` + /// must not memcpy into shared memory directly — wgpu's + /// `queue.write_buffer` is ordered AFTER already-submitted work, and a + /// raw memcpy is not: host writes then race in-flight kernel reads of + /// the same buffer (observed: the zealot biped's staged per-env uploads + /// corrupting so every env ran the last env's robot). Instead, writes + /// while work may be in flight go through a staging buffer + a blit + /// command buffer committed to the queue — FIFO queue order gives + /// exactly wgpu's semantics with no CPU stall. + gpu_busy: Arc, + /// Open blit command buffer collecting the current write burst's staged + /// copies (see `gpu_busy`). Committed lazily by the next `submit` / + /// `synchronize`, so a burst of writes costs ONE command buffer, and + /// queue FIFO order still sequences it after in-flight work and before + /// the next compute submission. + pending_writes: Arc>>, } // SAFETY: metal::Device and CommandQueue are thread-safe (MTLDevice/MTLCommandQueue @@ -51,6 +67,15 @@ unsafe impl Send for Metal {} unsafe impl Sync for Metal {} impl Metal { + /// Commits the open write-burst blit command buffer, if any (see + /// `pending_writes`). + fn flush_pending_writes(&self) { + if let Some((cmd, blit)) = self.pending_writes.lock().unwrap().take() { + blit.end_encoding(); + cmd.commit(); + } + } + /// Commits an empty command buffer with a completion handler that flips the /// returned flag once the GPU finishes it. /// @@ -94,6 +119,8 @@ impl Metal { let timing_caps = detect_timing_caps(&device); Ok(Self { device, + gpu_busy: Arc::new(AtomicBool::new(false)), + pending_writes: Arc::new(Mutex::new(None)), queue, module_cache: Arc::new(Mutex::new(HashMap::new())), timing_caps, @@ -537,6 +564,35 @@ impl Backend for Metal { ) .map_err(|e| MetalBackendError::SpirVParse(format!("{e}")))?; + // Optional (KHAL_METAL_WGSL_ROUNDTRIP=1): normalize the IR through a + // WGSL round-trip before MSL generation — the shape the wgpu path + // effectively feeds the MSL writer. Debug lever for spv-in-shaped-IR + // miscompiles. + let module = if std::env::var("KHAL_METAL_WGSL_ROUNDTRIP").as_deref() == Ok("1") { + let info = naga::valid::Validator::new( + naga::valid::ValidationFlags::all(), + naga::valid::Capabilities::all(), + ) + .validate(&module) + .map_err(|e| MetalBackendError::SpirVParse(format!("pre-wgsl validate: {e:?}")))?; + let wgsl = naga::back::wgsl::write_string( + &module, + &info, + naga::back::wgsl::WriterFlags::empty(), + ) + .map_err(|e| MetalBackendError::SpirVParse(format!("wgsl out: {e:?}")))?; + let mut rt = naga::front::wgsl::parse_str(&wgsl) + .map_err(|e| MetalBackendError::SpirVParse(format!("wgsl re-parse: {e}")))?; + // WGSL sanitizes `::` out of entry-point names — restore the + // originals so the per-entry-point resource map still keys. + for (ep, orig) in rt.entry_points.iter_mut().zip(module.entry_points.iter()) { + ep.name = orig.name.clone(); + } + rt + } else { + module + }; + // Validate so the MSL backend has the type info it needs. let info = naga::valid::Validator::new( naga::valid::ValidationFlags::all(), @@ -666,14 +722,34 @@ impl Backend for Metal { per_entry_point.insert(entry_point.to_string(), entry_point_resources); let options = naga::back::msl::Options { - lang_version: (2, 4), + // Metal 3: matches what wgpu targets on modern macOS. The 2.4 + // target visibly miscompiled the fused multibody dynamics kernel + // (lane decode collapsed to the last slot) on M-series. + lang_version: (3, 0), per_entry_point_map: per_entry_point, inline_samplers: vec![], spirv_cross_compatibility: false, fake_missing_bindings: false, - bounds_check_policies: naga::proc::BoundsCheckPolicies::default(), - zero_initialize_workgroup_memory: false, - force_loop_bounding: false, + // Match wgpu's WebGPU semantics EXACTLY: OOB reads return zero + // and OOB writes are skipped (`ReadZeroSkipWrite`), and workgroup + // memory starts zeroed. This is load-bearing, not defensive: the + // physics kernels were written against wgpu and rely on + // zero-read/skip-write semantics for by-design out-of-bounds + // accesses (inactive lane slots, StepRng trailer threads). naga's + // default (`Unchecked`) is UB there, and `Restrict` deterministically + // clamps reads to the END of each buffer — every batch then + // computes from the LAST batch's slice (observed: all envs ran + // batch N-1's robot on Metal while WebGPU was correct). + bounds_check_policies: naga::proc::BoundsCheckPolicies { + index: naga::proc::BoundsCheckPolicy::ReadZeroSkipWrite, + buffer: naga::proc::BoundsCheckPolicy::ReadZeroSkipWrite, + ..naga::proc::BoundsCheckPolicies::default() + }, + zero_initialize_workgroup_memory: true, + // wgpu enables loop bounding on Metal: Apple's compiler + // miscompiles potentially-unbounded loops (the trimesh BVH + // traversal visibly lost contacts without this). + force_loop_bounding: true, task_dispatch_limits: None, mesh_shader_primitive_indices_clamp: true, ray_query_initialization_tracking: true, @@ -721,8 +797,12 @@ impl Backend for Metal { .map(|ep| ep.workgroup_size) .ok_or_else(|| MetalBackendError::EntryPointNotFound(entry_point.into()))?; - // Compile MSL. + // Compile MSL. Metal's compiler defaults to fast math, which drops + // NaN/Inf semantics and reassociates float ops — the physics kernels + // (and wgpu's Metal path, which this backend must agree with) need + // IEEE behavior, so disable it. let compile_options = metal::CompileOptions::new(); + compile_options.set_fast_math_enabled(false); let library: Library = self .device .new_library_with_source(&msl, &compile_options) @@ -741,6 +821,18 @@ impl Backend for Metal { .new_compute_pipeline_state(&descriptor) .map_err(MetalBackendError::PipelineCreate)?; + // Dispatching more threads per threadgroup than the pipeline's + // register-pressure-derived maximum is UNDEFINED BEHAVIOR on Metal — + // surface it loudly instead of silently corrupting. + let max_threads = pipeline.max_total_threads_per_threadgroup() as u32; + let wanted = workgroup_size[0] * workgroup_size[1] * workgroup_size[2]; + if wanted > max_threads { + eprintln!( + "[khal-metal] WARNING: `{entry_point}` needs {wanted} threads/threadgroup \ + but the pipeline supports only {max_threads} — dispatches will misbehave" + ); + } + Ok(MetalFunction { pipeline, slot_map: Arc::new(slot_map), @@ -779,6 +871,7 @@ impl Backend for Metal { } fn synchronize(&self) -> Result<(), Self::Error> { + self.flush_pending_writes(); // Submit and wait on a fresh empty command buffer to flush the queue. let cb = self.queue.new_command_buffer(); cb.commit(); @@ -787,7 +880,10 @@ impl Backend for Metal { } fn submit(&self, encoder: Self::Encoder) -> Result<(), Self::Error> { + // Write-burst blits must land BEFORE this submission's compute work. + self.flush_pending_writes(); encoder.command_buffer.commit(); + self.gpu_busy.store(true, Ordering::Release); Ok(()) } @@ -831,6 +927,14 @@ impl Backend for Metal { let inner = self .device .new_buffer(byte_len as NSUInteger, resource_options(usage)); + // WebGPU (and thus the wgpu backend) guarantees zero-initialized + // buffers; Metal leaves new-buffer contents undefined. The compute + // pipelines are written against wgpu and some state buffers rely on + // starting zeroed — match the semantics. Shared storage: memset on + // the CPU side is cheapest. + unsafe { + std::ptr::write_bytes(inner.contents() as *mut u8, 0, byte_len); + } Ok(MetalBuffer { inner, len, @@ -852,6 +956,34 @@ impl Backend for Metal { if bytes.is_empty() { return Ok(()); } + // Order this write after all submitted GPU work (see `gpu_busy`): + // while kernels may be in flight, route the bytes through a staging + // buffer + a blit command buffer — the queue's FIFO order sequences + // it after the in-flight work and before later submits, exactly like + // wgpu's `queue.write_buffer`, with no CPU stall. The direct memcpy + // fast path only runs when nothing has been submitted since the last + // drain (scene setup, post-readback staging). + if self.gpu_busy.load(Ordering::Acquire) { + let staging = self.device.new_buffer_with_data( + bytes.as_ptr() as *const _, + bytes.len() as NSUInteger, + MTLResourceOptions::StorageModeShared, + ); + let mut pending = self.pending_writes.lock().unwrap(); + let (_, blit) = pending.get_or_insert_with(|| { + let cmd = self.queue.new_command_buffer().to_owned(); + let blit = cmd.new_blit_command_encoder().to_owned(); + (cmd, blit) + }); + blit.copy_from_buffer( + &staging, + 0, + &buffer.inner, + byte_offset as NSUInteger, + bytes.len() as NSUInteger, + ); + return Ok(()); + } // SAFETY: contents() is valid for the buffer's lifetime; we copy non-overlapping bytes. unsafe { let dst = (buffer.inner.contents() as *mut u8).add(byte_offset); From 4ecf4175bc985bba3574025a90392c72c34301fd Mon Sep 17 00:00:00 2001 From: haixuantao Date: Tue, 1 Sep 2026 16:51:36 +0800 Subject: [PATCH 25/33] metal: set threadGroupSizeIsMultipleOfThreadExecutionWidth on PSOs All kernels use 64-thread (2x SIMD-width) groups dispatched at exactly that size; the hint lets Metal drop per-dispatch edge masking. Measured ~3% on zealot biped training @1024 envs (8.7-8.8 -> 8.4-8.6 s/iter). Claude-Session: https://claude.ai/code/session_01EvbTzx6ZJNtcAKwVQ1B68i --- crates/khal/src/backend/metal.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/crates/khal/src/backend/metal.rs b/crates/khal/src/backend/metal.rs index 100813f..4906c19 100644 --- a/crates/khal/src/backend/metal.rs +++ b/crates/khal/src/backend/metal.rs @@ -815,6 +815,11 @@ impl Backend for Metal { let descriptor = ComputePipelineDescriptor::new(); descriptor.set_compute_function(Some(&function)); descriptor.set_label(entry_point); + // All kernels declare 64-thread (2x SIMD-width) or wider groups and + // are dispatched at exactly that size — let Metal skip edge masking. + if workgroup_size[0] * workgroup_size[1] * workgroup_size[2] % 32 == 0 { + descriptor.set_thread_group_size_is_multiple_of_thread_execution_width(true); + } let pipeline = self .device From 930da9995cee4931d702fbccfe0a614027814a27 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Wed, 2 Sep 2026 16:01:35 +0800 Subject: [PATCH 26/33] webgpu: name the oversize dispatch in the 65535-cap panic wgpu's validation error doesn't say which dispatch overflowed; panic at the khal layer with the grid so the kernel is identifiable from the backtrace. Claude-Session: https://claude.ai/code/session_01EvbTzx6ZJNtcAKwVQ1B68i --- crates/khal/src/backend/webgpu.rs | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/crates/khal/src/backend/webgpu.rs b/crates/khal/src/backend/webgpu.rs index 13f88b8..872ab89 100644 --- a/crates/khal/src/backend/webgpu.rs +++ b/crates/khal/src/backend/webgpu.rs @@ -987,10 +987,21 @@ impl<'a> Dispatch<'a, WebGpu> for WebGpuDispatch<'a> { } } + // WebGPU caps each workgroup-grid dimension at 65535; wgpu's own + // validation error doesn't say WHICH dispatch — name it here. + let check_grid = |grid_dim: [u32; 3]| { + if grid_dim.iter().any(|&d| d > 65535) { + panic!( + "dispatch grid {grid_dim:?} exceeds WebGPU's 65535 per-dimension \ + workgroup cap (a flattened dispatch that needs a 2D split?)" + ); + } + }; match grid.into() { DispatchGrid::Grid(grid_dim) => { // NOTE: we don't need to queue if the workgroup is empty. if grid_dim[0] * grid_dim[1] * grid_dim[2] > 0 { + check_grid(grid_dim); self.pass .dispatch_workgroups(grid_dim[0], grid_dim[1], grid_dim[2]); } @@ -1002,6 +1013,7 @@ impl<'a> Dispatch<'a, WebGpu> for WebGpuDispatch<'a> { threads[2].div_ceil(_block_dim[2]), ]; if grid_dim[0] * grid_dim[1] * grid_dim[2] > 0 { + check_grid(grid_dim); self.pass .dispatch_workgroups(grid_dim[0], grid_dim[1], grid_dim[2]); } From 1cffa2450269fda62f489a961f6f1e7b70567078 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 13 Sep 2026 12:22:41 +0200 Subject: [PATCH 27/33] cuda-oxide + rust-cuda: end-to-end PTX/cubin build path and device fixes --- .gitignore | 2 + README.md | 70 ++++- crates/cargo-cuda/src/main.rs | 7 + crates/khal-builder/Cargo.toml | 5 + crates/khal-builder/src/lib.rs | 310 ++++++++++++++++++- crates/khal-derive/src/spirv_bindgen/cuda.rs | 80 +++-- crates/khal-example-shaders/Cargo.toml | 3 + crates/khal-example/Cargo.toml | 3 + crates/khal-std/Cargo.toml | 4 +- crates/khal-std/src/lib.rs | 2 +- crates/khal-std/src/num_traits.rs | 27 +- crates/khal-std/src/sync/barriers.rs | 16 +- 12 files changed, 490 insertions(+), 39 deletions(-) diff --git a/.gitignore b/.gitignore index df59c6c..0634360 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,8 @@ debug/ target/ autogen/ +# PTX generated by khal-builder for the local GPU (cargo cuda / cargo oxide) +crates/khal-example/shaders-spirv/*.ptx # Remove Cargo.lock from gitignore if creating an executable, leave it for libraries # More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html diff --git a/README.md b/README.md index fc3a270..995dcc3 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,73 @@ This requires the **CUDA toolkit** to be installed and the `CUDA_PATH` environme point to it (e.g. `/usr/local/cuda`). The install step downloads a pinned Rust nightly, adds the `nvptx64-nvidia-cuda` target, and compiles the codegen backend. +Compiling that backend (`rustc_codegen_nvvm`) needs a few system packages besides the toolkit: +`pkg-config` and `libssl-dev` (for its `openssl-sys` dependency), and libclang with its resource +headers for `bindgen` (`libclang-common-21-dev`, or `clang-21`, on Debian/Ubuntu; otherwise the +libNVVM bindings fail with `'stddef.h' file not found`). + +It also needs **LLVM 7.1.0** (libNVVM only accepts LLVM 7 bitcode). Rust-CUDA ships no prebuilt +LLVM for Linux, so build it once from source and point `LLVM_CONFIG` at it. Distro packages no +longer exist for LLVM 7; this is the recipe from Rust-CUDA's own Dockerfiles (cmake 3.x is required, +cmake 4 rejects LLVM 7's old policy settings, and GCC 15 builds it cleanly): + +```bash +curl -sSfLO https://github.com/llvm/llvm-project/releases/download/llvmorg-7.1.0/llvm-7.1.0.src.tar.xz && tar -xf llvm-7.1.0.src.tar.xz && mkdir llvm-7.1.0.src/build && cd llvm-7.1.0.src/build && cmake -G Ninja -DCMAKE_BUILD_TYPE=Release -DLLVM_TARGETS_TO_BUILD="X86;NVPTX" -DLLVM_BUILD_LLVM_DYLIB=ON -DLLVM_LINK_LLVM_DYLIB=ON -DLLVM_ENABLE_ASSERTIONS=OFF -DLLVM_ENABLE_BINDINGS=OFF -DLLVM_INCLUDE_EXAMPLES=OFF -DLLVM_INCLUDE_TESTS=OFF -DLLVM_INCLUDE_BENCHMARKS=OFF -DLLVM_INCLUDE_DOCS=OFF -DLLVM_ENABLE_ZLIB=OFF -DLLVM_ENABLE_TERMINFO=OFF -DLLVM_ENABLE_LIBXML2=OFF -DLLVM_ENABLE_LIBEDIT=OFF -DCMAKE_INSTALL_PREFIX=$HOME/llvm-7 .. && ninja && ninja install +``` + +```bash +LLVM_CONFIG=$HOME/llvm-7/bin/llvm-config cargo cuda install +``` + +### cargo-oxide (alternative CUDA / PTX compiler) + +The `cuda-oxide` feature compiles the kernels with [cuda-oxide](https://github.com/NVlabs/cuda-oxide) +(NVIDIA's LLVM 21 based Rust → PTX backend) instead of `cargo-cuda`. The shader crate is then built +as an ordinary host-target crate and the cuda-oxide codegen backend intercepts the kernel entries. + +Requirements: + +- The **CUDA toolkit** (12.8+) on `PATH`, with `CUDA_HOME` pointing at it (e.g. `/usr/local/cuda`). + `nvcc` must be on `PATH` for `cudarc`'s build script, and `CUDA_HOME` is how cuda-oxide finds + libdevice, libnvvm and nvJitLink. +- `libffi-dev` (Debian/Ubuntu package name). cuda-oxide's codegen backend is a dylib linked against + rustc's `librustc_driver`, which needs libffi at link time. +- The Rust nightly pinned by cuda-oxide's `rust-toolchain.toml` (currently `nightly-2026-04-03`), + with the `rust-src`, `rustc-dev` and `llvm-tools` components. cuda-oxide is a rustc codegen + backend linked against that exact toolchain, so build everything with `cargo +`. +- `cargo-oxide`, installed from the same cuda-oxide revision that `khal-std` pins its `cuda-device` + dependency to (see `crates/khal-std/Cargo.toml`): + +```bash +cargo +nightly-2026-04-03 install --git https://github.com/NVlabs/cuda-oxide.git --rev 62472763 cargo-oxide +``` + +On first use `cargo-oxide` clones cuda-oxide into `~/.cargo/cuda-oxide/src` and builds the codegen +backend from it (several minutes). To keep that clone on the pinned revision as well: + +```bash +git clone https://github.com/NVlabs/cuda-oxide.git ~/.cargo/cuda-oxide/src && git -C ~/.cargo/cuda-oxide/src checkout 62472763 +``` + +Then, for the example: + +```bash +cargo +nightly-2026-04-03 run --bin khal-example --features cuda-oxide +``` + +`khal-builder` targets the local GPU's compute capability (via `nvidia-smi`); set `KHAL_CUDA_ARCH=sm_XX` +(or `CUDA_OXIDE_TARGET=sm_XX`) to override it. A prebuilt PTX/cubin can be embedded instead of compiling +by setting `CUDA_OXIDE_SHADERS_PTX_` to its path. + +### PTX vs. cubin + +With either compiler, `khal-builder` then assembles the PTX into a **cubin** for the local GPU using the +toolkit's `ptxas` when it can find it (`PATH`, `CUDA_HOME`, `CUDA_PATH`, `/usr/local/cuda`). The driver +otherwise JIT-compiles the PTX at load time and rejects PTX whose ISA version is newer than the driver +supports (`CUDA_ERROR_UNSUPPORTED_PTX_VERSION`, typically a toolkit newer than the driver, e.g. 13.3 vs +13.2). A cubin sidesteps that and skips the JIT. Set `KHAL_CUDA_KEEP_PTX=1` to keep forward-compatible +PTX text instead (e.g. when shipping a binary to machines with other GPUs). + ## Crates | Crate | Description | @@ -66,7 +133,8 @@ point to it (e.g. `/usr/local/cuda`). The install step downloads a pinned Rust n | Backend | Feature flag | Shader format | Notes | |---------|-------------|---------------|-------| | WebGPU | `webgpu` (default) | SPIR-V | Cross-platform via wgpu | -| CUDA | `cuda` | PTX | NVIDIA GPUs, requires CUDA toolkit | +| CUDA | `cuda` | PTX | NVIDIA GPUs, requires CUDA toolkit; kernels compiled by `cargo-cuda` | +| CUDA | `cuda-oxide` | PTX | Same runtime as `cuda`; kernels compiled by `cargo-oxide` instead | | CPU | `cpu` | Native | Single-threaded; use `cpu-parallel` for rayon-based dispatch | diff --git a/crates/cargo-cuda/src/main.rs b/crates/cargo-cuda/src/main.rs index 5156022..812a7ee 100644 --- a/crates/cargo-cuda/src/main.rs +++ b/crates/cargo-cuda/src/main.rs @@ -190,6 +190,13 @@ components = ["llvm-tools-preview", "rust-src", "rustc-dev"] if !status.success() { eprintln!("Failed to build rustc_codegen_nvvm"); + eprintln!(); + eprintln!("rustc_codegen_nvvm has native build requirements beyond the CUDA toolkit:"); + eprintln!(" - LLVM 7.1.0 (libNVVM only accepts LLVM 7 bitcode). Rust-CUDA ships no"); + eprintln!(" prebuilt LLVM for Linux, so build it from source and set"); + eprintln!(" LLVM_CONFIG=/path/to/llvm-7/bin/llvm-config (see the khal README)."); + eprintln!(" - pkg-config + OpenSSL headers (openssl-sys), e.g. `pkg-config libssl-dev`."); + eprintln!(" - libclang with its resource headers for bindgen, e.g. `libclang-common--dev`."); std::process::exit(1); } diff --git a/crates/khal-builder/Cargo.toml b/crates/khal-builder/Cargo.toml index 7219733..5513725 100644 --- a/crates/khal-builder/Cargo.toml +++ b/crates/khal-builder/Cargo.toml @@ -8,7 +8,12 @@ edition = "2024" license = "MIT OR Apache-2.0" [features] +# Build PTX for the CUDA backend with rust-cuda's `cargo cuda`. cuda = [] +# Build PTX for the CUDA backend with cuda-oxide's `cargo oxide` instead +# (unified host-target interception; the shader crate must expose a +# `cuda-oxide` feature forwarding to `khal-std/cuda-oxide`). Implies `cuda`. +cuda-oxide = ["cuda"] unsafe_remove_boundchecks = [] [dependencies] diff --git a/crates/khal-builder/src/lib.rs b/crates/khal-builder/src/lib.rs index b47a4d0..902cc43 100644 --- a/crates/khal-builder/src/lib.rs +++ b/crates/khal-builder/src/lib.rs @@ -14,7 +14,8 @@ pub struct KhalBuilder { features: Vec, // The `RUST_MIN_STACK` given to the shader builders. rust_min_stack: u32, - /// If the `cuda` feature is enabled and this is `true`, then cuda PTX kernels will be built with cargo-cuda. + /// If the `cuda` feature is enabled and this is `true`, then cuda PTX kernels will be built + /// with cargo-cuda (or cargo-oxide when the `cuda-oxide` feature is enabled). /// Default: `true` #[allow(dead_code)] build_cuda: bool, @@ -140,6 +141,7 @@ impl KhalBuilder { println!("cargo:rerun-if-env-changed=CARGO_FEATURE_PUSH_CONSTANTS"); // TODO: currently unused println!("cargo:rerun-if-env-changed=CARGO_FEATURE_CUDA"); + println!("cargo:rerun-if-env-changed=CARGO_FEATURE_CUDA_OXIDE"); println!("cargo:rerun-if-env-changed=KHAL_SKIP_SPIRV"); } @@ -178,8 +180,12 @@ impl KhalBuilder { /// /// `CUDA_OXIDE_SHADERS_PTX_` (upper-snake, e.g. /// `CUDA_OXIDE_SHADERS_PTX_VORTX_SHADERS`) points at a prebuilt PTX/cubin; - /// when set it is embedded directly as `shaders.ptx` and the `cargo cuda` - /// (cuda-oxide) toolchain is not required. + /// when set it is embedded directly as `shaders.ptx` and no PTX compiler + /// (`cargo cuda` / `cargo oxide`) is required. + /// + /// Otherwise the PTX is produced by `cargo oxide` (cuda-oxide) when the + /// `cuda-oxide` feature is enabled, and by `cargo cuda` (rust-cuda) when + /// it is not. #[cfg(feature = "cuda")] fn build_ptx(&self, output_dir: impl AsRef) { let output_dir = output_dir.as_ref(); @@ -205,6 +211,21 @@ impl KhalBuilder { return; } + #[cfg(feature = "cuda-oxide")] + { + self.build_ptx_cuda_oxide(output_dir); + } + #[cfg(not(feature = "cuda-oxide"))] + { + self.build_ptx_cargo_cuda(output_dir); + } + + assemble_cubin(&output_dir.join("shaders.ptx")); + } + + /// PTX via rust-cuda's `cargo cuda build` (rustc_codegen_nvvm). + #[cfg(all(feature = "cuda", not(feature = "cuda-oxide")))] + fn build_ptx_cargo_cuda(&self, output_dir: &Path) { let features_str = self.features.join(","); let mut args = vec![ @@ -233,4 +254,287 @@ impl KhalBuilder { panic!("cargo cuda build failed"); } } + + /// PTX via cuda-oxide's `cargo oxide build` (unified host-target interception). + /// + /// The shader crate is compiled as an ordinary host-target crate with its + /// `cuda-oxide` feature enabled; the cuda-oxide codegen backend intercepts + /// the `#[spirv_bindgen]`-generated kernel entries and writes + /// `.ptx` into `CUDA_OXIDE_PTX_DIR`, which we then copy to + /// `/shaders.ptx` (the file `from_dir_ptx` loads). + /// + /// Environment knobs (all optional): + /// - `KHAL_CUDA_ARCH=sm_XX` / `CUDA_OXIDE_TARGET=sm_XX`: pin the PTX + /// target. Otherwise the local GPU's compute capability is detected with + /// `nvidia-smi`; if that fails, cuda-oxide's own default target is used. + /// - `KHAL_CUDA_OXIDE_TARGET_DIR`: cargo target dir for the nested build. + /// Defaults to `/khal-cuda-oxide`. It must differ + /// from the outer build's target dir, which the outer cargo holds locked + /// while this build script runs. + /// - `CUDA_OXIDE_BACKEND`, `CUDA_OXIDE_LLC`, `CUDA_HOME`: forwarded to + /// cargo-oxide as-is (see its docs). + #[cfg(feature = "cuda-oxide")] + fn build_ptx_cuda_oxide(&self, output_dir: &Path) { + println!("cargo:rerun-if-env-changed=CUDA_OXIDE_TARGET"); + println!("cargo:rerun-if-env-changed=CUDA_OXIDE_BACKEND"); + println!("cargo:rerun-if-env-changed=KHAL_CUDA_OXIDE_TARGET_DIR"); + + // Keep every intermediate (.ll/.ptx, cargo target dir) OUT of + // `output_dir`: consumers embed that whole directory with + // `include_dir!`, so anything left there ends up in the binary. + let target_dir = std::env::var_os("KHAL_CUDA_OXIDE_TARGET_DIR") + .map(PathBuf::from) + .unwrap_or_else(|| cargo_target_dir_from_out_dir().join("khal-cuda-oxide")); + let ptx_dir = target_dir.join("ptx"); + std::fs::create_dir_all(&ptx_dir) + .unwrap_or_else(|e| panic!("failed to create cuda-oxide PTX dir {ptx_dir:?}: {e}")); + + let mut features = self.features.clone(); + if !features.iter().any(|f| f == "cuda-oxide") { + features.push("cuda-oxide".to_string()); + } + let features_str = features.join(","); + + // `cargo oxide build [opts] -- ` is cargo-oxide's + // "passthrough" mode: it runs `cargo build ` in the current + // directory (our shader crate) with the cuda-oxide codegen backend + // injected, without touching sources or cleaning artifacts. + let mut cmd = Command::new("cargo"); + cmd.args([ + "oxide", + "build", + "--features", + &features_str, + "--cargo-target-dir", + ]) + .arg(&target_dir); + + // cargo-oxide reads CUDA_OXIDE_TARGET itself; only pass `--arch` when + // the target comes from elsewhere (KHAL_CUDA_ARCH or GPU detection). + let arch = resolve_sm_arch(); + if std::env::var_os("CUDA_OXIDE_TARGET").is_none() + && let Some(arch) = &arch + { + cmd.args(["--arch", arch]); + } + + cmd.args(["--", "--release"]) + .current_dir(&self.shader_crate) + .env("CUDA_OXIDE_PTX_DIR", &ptx_dir) + .env("RUST_MIN_STACK", self.rust_min_stack.to_string()); + + println!( + "cargo:warning=khal-builder: compiling {} to PTX with cargo-oxide (target {})", + self.shader_crate.display(), + arch.as_deref().unwrap_or("cuda-oxide default") + ); + + let status = cmd.status().unwrap_or_else(|e| { + panic!( + "failed to run `cargo oxide` ({e}). Install it with \ + `cargo install --git https://github.com/NVlabs/cuda-oxide.git cargo-oxide` \ + using the nightly pinned by cuda-oxide, or set \ + CUDA_OXIDE_SHADERS_PTX_ to a prebuilt PTX/cubin." + ) + }); + if !status.success() { + panic!("cargo oxide build failed"); + } + + let ptx = find_cuda_oxide_ptx(&ptx_dir, &self.shader_crate); + let dst = output_dir.join("shaders.ptx"); + std::fs::copy(&ptx, &dst) + .unwrap_or_else(|e| panic!("failed to copy {ptx:?} to {dst:?}: {e}")); + println!("cargo:rerun-if-changed={}", ptx.display()); + } +} + +/// The cargo target directory of the build that is running this build +/// script, derived from `OUT_DIR` (`//build/-/out`). +/// Falls back to a directory under `OUT_DIR` when the layout is unexpected. +#[cfg(feature = "cuda-oxide")] +fn cargo_target_dir_from_out_dir() -> PathBuf { + let out_dir = PathBuf::from(std::env::var_os("OUT_DIR").expect("OUT_DIR is set by cargo")); + let looks_like_cargo_layout = out_dir + .ancestors() + .nth(2) + .is_some_and(|b| b.ends_with("build")); + out_dir + .ancestors() + .nth(4) + .filter(|_| looks_like_cargo_layout) + .map(Path::to_path_buf) + .unwrap_or_else(|| out_dir.join("target")) +} + +/// The `sm_XY` target for PTX/cubin generation: `KHAL_CUDA_ARCH`, else +/// `CUDA_OXIDE_TARGET`, else the local GPU's compute capability. +#[cfg(feature = "cuda")] +fn resolve_sm_arch() -> Option { + println!("cargo:rerun-if-env-changed=KHAL_CUDA_ARCH"); + std::env::var("KHAL_CUDA_ARCH") + .or_else(|_| std::env::var("CUDA_OXIDE_TARGET")) + .ok() + .filter(|a| !a.is_empty()) + .or_else(detect_local_sm_arch) +} + +/// Assembles `/shaders.ptx` in place into a cubin for the local +/// GPU with the CUDA toolkit's `ptxas`, when possible. +/// +/// Why: the CUDA driver JIT-compiles PTX text at load time and rejects PTX +/// whose ISA version is newer than the driver supports +/// (`CUDA_ERROR_UNSUPPORTED_PTX_VERSION`, e.g. PTX 9.3 from a 13.3 toolkit on +/// a driver that only supports CUDA 13.2). A cubin assembled by the toolkit's +/// own `ptxas` loads regardless of that version skew, and skips the JIT at +/// startup. khal's CUDA loader detects cubins by their ELF magic, so the file +/// keeps its `shaders.ptx` name. +/// +/// Best effort: when `ptxas` or the target arch cannot be determined, or +/// `ptxas` fails, the PTX text is left as-is (it may still JIT fine). Set +/// `KHAL_CUDA_KEEP_PTX=1` to always keep forward-compatible PTX text. +#[cfg(feature = "cuda")] +fn assemble_cubin(ptx_path: &Path) { + println!("cargo:rerun-if-env-changed=KHAL_CUDA_KEEP_PTX"); + println!("cargo:rerun-if-env-changed=CUDA_HOME"); + println!("cargo:rerun-if-env-changed=CUDA_PATH"); + if std::env::var_os("KHAL_CUDA_KEEP_PTX").is_some() { + return; + } + // Already a cubin (e.g. a prebuilt one supplied via CUDA_OXIDE_SHADERS_PTX_*). + match std::fs::read(ptx_path) { + Ok(bytes) if bytes.starts_with(b"\x7fELF") => return, + Ok(_) => {} + Err(e) => panic!("cannot read {ptx_path:?}: {e}"), + } + + let Some(arch) = resolve_sm_arch() else { + println!( + "cargo:warning=khal-builder: no GPU arch known (set KHAL_CUDA_ARCH=sm_XX); \ + keeping PTX text for driver JIT" + ); + return; + }; + let Some(ptxas) = find_cuda_tool("ptxas") else { + println!( + "cargo:warning=khal-builder: ptxas not found (PATH, CUDA_HOME, CUDA_PATH, \ + /usr/local/cuda); keeping PTX text for driver JIT" + ); + return; + }; + + let cubin_path = ptx_path.with_extension("cubin.tmp"); + let status = Command::new(&ptxas) + .args(["-arch", &arch, "-O3"]) + .arg(ptx_path) + .arg("-o") + .arg(&cubin_path) + .status(); + match status { + Ok(st) if st.success() => { + std::fs::rename(&cubin_path, ptx_path) + .unwrap_or_else(|e| panic!("failed to move {cubin_path:?} over {ptx_path:?}: {e}")); + println!( + "cargo:warning=khal-builder: assembled {} to a {arch} cubin with {}", + ptx_path.display(), + ptxas.display() + ); + } + Ok(st) => { + let _ = std::fs::remove_file(&cubin_path); + println!( + "cargo:warning=khal-builder: ptxas exited with {st}; keeping PTX text for driver JIT" + ); + } + Err(e) => { + let _ = std::fs::remove_file(&cubin_path); + println!( + "cargo:warning=khal-builder: failed to run {}: {e}; keeping PTX text for driver JIT", + ptxas.display() + ); + } + } +} + +/// A CUDA toolkit binary: from `PATH`, else `$CUDA_HOME/bin`, `$CUDA_PATH/bin`, +/// `/usr/local/cuda/bin`. +#[cfg(feature = "cuda")] +fn find_cuda_tool(name: &str) -> Option { + if Command::new(name) + .arg("--version") + .output() + .is_ok_and(|o| o.status.success()) + { + return Some(PathBuf::from(name)); + } + ["CUDA_HOME", "CUDA_PATH"] + .iter() + .filter_map(|var| std::env::var_os(var)) + .map(PathBuf::from) + .chain(std::iter::once(PathBuf::from("/usr/local/cuda"))) + .map(|root| root.join("bin").join(name)) + .find(|p| p.is_file()) +} + +/// `sm_XY` of the first local GPU as reported by `nvidia-smi`, if available. +#[cfg(feature = "cuda")] +fn detect_local_sm_arch() -> Option { + let out = Command::new("nvidia-smi") + .args(["--query-gpu=compute_cap", "--format=csv,noheader"]) + .output() + .ok()?; + if !out.status.success() { + return None; + } + let cap = String::from_utf8_lossy(&out.stdout); + let cap = cap.lines().next()?.trim(); + let (major, minor) = cap.split_once('.')?; + let major: u32 = major.parse().ok()?; + let minor: u32 = minor.parse().ok()?; + Some(format!("sm_{major}{minor}")) +} + +/// The `.ptx` cargo-oxide produced for the shader crate: `.ptx` +/// (hyphens normalised to underscores), or the only `.ptx` in the directory. +#[cfg(feature = "cuda-oxide")] +fn find_cuda_oxide_ptx(ptx_dir: &Path, shader_crate: &Path) -> PathBuf { + let package_name = std::fs::read_to_string(shader_crate.join("Cargo.toml")) + .ok() + .and_then(|manifest| { + manifest.lines().find_map(|line| { + let line = line.trim(); + let value = line.strip_prefix("name")?.trim_start().strip_prefix('=')?; + Some(value.trim().trim_matches('"').replace('-', "_")) + }) + }) + .or_else(|| { + shader_crate + .file_name() + .and_then(|n| n.to_str()) + .map(|n| n.replace('-', "_")) + }); + + if let Some(name) = &package_name { + let candidate = ptx_dir.join(format!("{name}.ptx")); + if candidate.is_file() { + return candidate; + } + } + + let mut ptx_files: Vec = std::fs::read_dir(ptx_dir) + .unwrap_or_else(|e| panic!("failed to read {ptx_dir:?}: {e}")) + .filter_map(|e| e.ok().map(|e| e.path())) + .filter(|p| p.extension().is_some_and(|ext| ext == "ptx")) + .collect(); + match ptx_files.len() { + 1 => ptx_files.remove(0), + 0 => panic!( + "cargo oxide produced no .ptx in {ptx_dir:?}. Does the shader crate enable its \ + `cuda-oxide` feature and contain `#[spirv_bindgen]` kernels?" + ), + _ => panic!( + "cargo oxide produced several .ptx files in {ptx_dir:?} ({ptx_files:?}) and none is \ + named after the shader package {package_name:?}" + ), + } } diff --git a/crates/khal-derive/src/spirv_bindgen/cuda.rs b/crates/khal-derive/src/spirv_bindgen/cuda.rs index 955153a..31c04d2 100644 --- a/crates/khal-derive/src/spirv_bindgen/cuda.rs +++ b/crates/khal-derive/src/spirv_bindgen/cuda.rs @@ -5,7 +5,9 @@ use super::{BuiltinKind, OriginalParam, OriginalParamKind, ShaderBinding, is_sli /// Generate the CUDA (nvptx64) kernel entry point. /// /// This function: -/// 1. Receives raw device pointers + byte lengths as u64 parameters +/// 1. Receives raw device pointers + element counts as u64 parameters +/// (the host's `CudaDispatch` pushes `byte_len / size_of::()` for slice +/// bindings, the same `(ptr, len)` ABI cuda-oxide's `&[T]` params use) /// 2. Computes builtin values from CUDA thread/block indices /// 3. Reconstructs slices from raw pointers /// 4. Calls the original shader function @@ -20,7 +22,7 @@ pub(super) fn generate_cuda_entry_block( let wg_y = workgroup_size[1]; // Generate kernel parameters: for each binding (sorted by set/index), - // storage buffers become (ptr: u64, byte_len: u64), uniforms become (ptr: u64). + // storage buffers become (ptr: u64, len: u64), uniforms become (ptr: u64). let mut cuda_params = Vec::new(); let mut cuda_body = Vec::new(); let mut cuda_call_args = Vec::new(); @@ -59,7 +61,7 @@ pub(super) fn generate_cuda_entry_block( for param in &sorted_bindings { let name = ¶m.name; let ptr_name = syn::Ident::new(&format!("{}_ptr", name), name.span()); - let len_name = syn::Ident::new(&format!("{}_byte_len", name), name.span()); + let len_name = syn::Ident::new(&format!("{}_len", name), name.span()); if let OriginalParamKind::Binding { is_uniform, @@ -125,7 +127,7 @@ pub(super) fn generate_cuda_entry_block( is_mutable, } => { let ptr_name = syn::Ident::new(&format!("{}_ptr", name), name.span()); - let len_name = syn::Ident::new(&format!("{}_byte_len", name), name.span()); + let len_name = syn::Ident::new(&format!("{}_len", name), name.span()); let is_slice = is_slice_reference(¶m.ty); let elem_ty = if let syn::Type::Reference(ref_type) = ¶m.ty { @@ -149,7 +151,7 @@ pub(super) fn generate_cuda_entry_block( let #name = unsafe { core::slice::from_raw_parts_mut( #ptr_name as *mut #elem_ty, - #len_name as usize / core::mem::size_of::<#elem_ty>(), + #len_name as usize, ) }; }, @@ -166,7 +168,7 @@ pub(super) fn generate_cuda_entry_block( let #name = unsafe { core::slice::from_raw_parts( #ptr_name as *const #elem_ty, - #len_name as usize / core::mem::size_of::<#elem_ty>(), + #len_name as usize, ) }; }, @@ -275,7 +277,7 @@ pub(super) fn generate_cuda_oxide_entry_block( .map(|b| (b.descriptor_set, b.binding)) .unwrap_or((0, 0)) }); - // The host (khal CudaDispatch) pushes a `(ptr, byte_len)` pair for EVERY + // The host (khal CudaDispatch) pushes a `(ptr, len)` pair for EVERY // storage binding. A `&[T]`/`&mut [T]` typed param already lowers to that // 2-arg ABI, but a sized-array ref `&[T; N]`/`&mut [T; N]` lowers to a single // thin pointer — so the kernel would consume one fewer arg than the host @@ -291,7 +293,7 @@ pub(super) fn generate_cuda_oxide_entry_block( let ty = &p.ty; // Uniform bindings are pushed by the host as a single pointer, so keep // them typed (thin pointer ABI). Storage bindings are pushed as a - // `(ptr, byte_len)` pair: slice refs already match that 2-arg ABI, but + // `(ptr, len)` pair: slice refs already match that 2-arg ABI, but // sized-array refs (`&[T; N]`) and scalar refs (`&T`) lower to a single // thin pointer and would under-consume the host's args, shifting every // later binding. Receive those as slices and reconstruct the original @@ -392,23 +394,25 @@ pub(super) fn generate_cuda_oxide_entry_block( syn::Ident::new(&format!("__smem_{}", name), name.span()); match &p.ty { syn::Type::Reference(r) => match &*r.elem { - // `&mut [f32; N]` shared tile -> `&mut SmemBuf` (direct - // `SharedArray` indexing -> st.shared). Body indexes it via - // MaybeIndexUnchecked and can reborrow it across helpers - // generic over `&mut impl MaybeIndexUnchecked`. + // `&mut [f32; N]` shared tile -> a real `&mut [f32; N]` into + // the `SharedArray`'s storage, so the body can index it, + // call MaybeIndexUnchecked on it AND pass it verbatim to + // helpers typed `&mut [T; N]` (e.g. vortx's tree reductions). + // `SharedArray::as_mut_ptr` is intercepted by the cuda-oxide + // importer and yields a generic pointer to shared memory + // (cvta.shared), so accesses through the reference are + // ordinary generic loads/stores on shared memory. syn::Type::Array(arr) => { let elem_ty = &*arr.elem; let len_expr = &arr.len; - let sb_name = - syn::Ident::new(&format!("__smembuf_{}", name), name.span()); preludes.push(quote! { static mut #static_name: khal_std::cuda_oxide_glue::SharedArray<#elem_ty, { #len_expr }> = khal_std::cuda_oxide_glue::SharedArray::UNINIT; - let mut #sb_name = khal_std::cuda_oxide_glue::SmemBuf( - unsafe { &mut *core::ptr::addr_of_mut!(#static_name) } - ); - let #name = &mut #sb_name; + let #name: &mut [#elem_ty; #len_expr] = unsafe { + &mut *((&mut *core::ptr::addr_of_mut!(#static_name)).as_mut_ptr() + as *mut [#elem_ty; #len_expr]) + }; }); } // `&mut T` scalar broadcast slot -> `SharedArray` -> @@ -437,19 +441,35 @@ pub(super) fn generate_cuda_oxide_entry_block( let body = &func.block; // Name the entry with cuda-oxide's reserved kernel prefix directly - // (`cuda_oxide_kernel_246e25db_`) instead of going through the - // `#[kernel]` proc-macro: the collector roots kernels by that def-path - // marker and the PTX entry keeps the unprefixed base name, while the - // macro's host-side glue (a `cuda_host::CudaKernel` marker impl) would - // drag a cuda-host dependency into pure shader crates. Feature-gated, - // not target-gated: with the cuda-oxide backend installed the crate is - // compiled on the HOST target (unified interception) or for nvptx64 - // (device-only builds) — the entry is generated identically for both. + // (`KERNEL_PREFIX` in cuda-oxide's `reserved-oxide-symbols`, currently + // `cuda_oxide_codegen_v1_cuda_oxide_kernel_246e25db_`) instead of + // going through the `#[kernel]` proc-macro: the collector roots kernels by + // that def-path marker and the PTX entry keeps the unprefixed base name, + // while the macro's host-side glue (a `cuda_host::CudaKernel` marker impl) + // would drag a cuda-host dependency into pure shader crates. The prefix + // must track cuda-oxide's: the backend rejects roots using an older + // prefix (e.g. the pre-"scoped cache protocol" `cuda_oxide_kernel_246e25db_`) + // with an explicit diagnostic. Feature-gated, not target-gated: with the + // cuda-oxide backend installed the crate is compiled on the HOST target + // (unified interception) or for nvptx64 (device-only builds) — the entry + // is generated identically for both. let prefixed_entry_ident = syn::Ident::new( - &format!("cuda_oxide_kernel_246e25db_{}", cuda_entry_ident), + &format!("{CUDA_OXIDE_KERNEL_PREFIX}{cuda_entry_ident}"), cuda_entry_ident.span(), ); quote! { + // The "scoped Cargo cache protocol": cuda-oxide's `#[kernel]` records + // the backend's codegen identity (output mode, arch, tool provenance) + // in the *device crate's* dep-info so that only crates owning device + // code are rebuilt when it changes. `option_env!` does the same from + // stable Rust (rustc tracks env vars read by `env!`/`option_env!`). + #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] + const _: (Option<&str>, Option<&str>, Option<&str>) = ( + option_env!("CUDA_OXIDE_INTERNAL_CODEGEN_FINGERPRINT"), + option_env!("CUDA_OXIDE_MATERIALIZE_CUBIN"), + option_env!("CUDA_OXIDE_INTERNAL_MATERIALIZER_PROVENANCE"), + ); + #[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] #[allow(non_snake_case)] // `no_mangle` forces local codegen: without it rustc's @@ -464,3 +484,9 @@ pub(super) fn generate_cuda_oxide_entry_block( } } } + +/// cuda-oxide's reserved kernel-entry prefix (`reserved_oxide_symbols::KERNEL_PREFIX`). +/// The backend strips it to obtain the PTX `.entry` name, so the PTX entry is +/// exactly `CUDA_ENTRY_POINT`. Keep in sync with the `cuda-device` revision +/// pinned in `khal-std/Cargo.toml`. +const CUDA_OXIDE_KERNEL_PREFIX: &str = "cuda_oxide_codegen_v1_cuda_oxide_kernel_246e25db_"; diff --git a/crates/khal-example-shaders/Cargo.toml b/crates/khal-example-shaders/Cargo.toml index 896a150..46dc788 100644 --- a/crates/khal-example-shaders/Cargo.toml +++ b/crates/khal-example-shaders/Cargo.toml @@ -16,6 +16,9 @@ build = "build.rs" cpu = ["khal-std/cpu"] cpu-parallel = ["khal-std/cpu-parallel"] cuda = ["khal-std/cuda"] +# Compile the CUDA kernels with the cuda-oxide backend (host-target unified +# interception) instead of rust-cuda; see khal-std's feature docs. +cuda-oxide = ["khal-std/cuda-oxide"] [dependencies] khal-std = { path = "../khal-std" } diff --git a/crates/khal-example/Cargo.toml b/crates/khal-example/Cargo.toml index 09bcc2d..c50fb62 100644 --- a/crates/khal-example/Cargo.toml +++ b/crates/khal-example/Cargo.toml @@ -11,6 +11,9 @@ metal = [ "khal/metal" ] cpu = ["khal/cpu", "khal-example-shaders/cpu"] cpu-parallel = ["cpu", "khal-example-shaders/cpu-parallel"] cuda = ["khal/cuda", "khal-builder/cuda", "khal-example-shaders/cuda"] +# Native CUDA with kernels compiled by cuda-oxide (`cargo oxide`) rather than +# rust-cuda's `cargo cuda`. Implies `cuda`. +cuda-oxide = ["cuda", "khal-builder/cuda-oxide", "khal-example-shaders/cuda-oxide"] [dependencies] khal = { version = "0.3.0", path = "../khal", features = ["derive"] } diff --git a/crates/khal-std/Cargo.toml b/crates/khal-std/Cargo.toml index c52742b..726ac25 100644 --- a/crates/khal-std/Cargo.toml +++ b/crates/khal-std/Cargo.toml @@ -38,7 +38,9 @@ rust-cuda = ["dep:cuda_std"] # The `cuda-device` API-surface crate resolves from upstream NVlabs # cuda-oxide (git pin below); consuming workspaces may [patch] it to a local # checkout when hacking on both. -cuda-oxide = ["dep:num-traits", "dep:cuda-device"] +# `glamx/scalar-math`: the unified path compiles shader crates for the HOST +# target, so glam must not take its SSE/SIMD paths (those don't lower to PTX). +cuda-oxide = ["dep:num-traits", "dep:cuda-device", "glamx/scalar-math"] [dependencies] glamx = { version = "0.3", default-features = false, features = ["nostd-libm", "bytemuck", "u32", "i32", "f64"] } diff --git a/crates/khal-std/src/lib.rs b/crates/khal-std/src/lib.rs index 0edce3c..707aec3 100644 --- a/crates/khal-std/src/lib.rs +++ b/crates/khal-std/src/lib.rs @@ -8,7 +8,7 @@ // nvptx inline asm (warp shuffle in `sync::subgroup`) is unstable; the // cuda-oxide backend compiles for the real nvptx64 target so it needs the gate. #![cfg_attr(target_arch = "nvptx64", feature(asm_experimental_arch))] -// `Float` for the cuda-oxide backend calls `core::intrinsics::{expf32,..}` +// `Float` for the cuda-oxide backend calls `core::intrinsics::{sqrtf32,..}` // (lowered to libdevice `__nv_*` by cuda-oxide) instead of software libm. #![cfg_attr(feature = "cuda-oxide", feature(core_intrinsics))] #![cfg_attr(feature = "cuda-oxide", allow(internal_features))] diff --git a/crates/khal-std/src/num_traits.rs b/crates/khal-std/src/num_traits.rs index 80ce764..ef5dc02 100644 --- a/crates/khal-std/src/num_traits.rs +++ b/crates/khal-std/src/num_traits.rs @@ -35,20 +35,37 @@ mod cuda_oxide_float { } macro_rules! float_impl { - ($ty:ty, $exp:ident, $ln:ident, $sqrt:ident, $pow:ident, $floor:ident, $ceil:ident, + ($ty:ty, $sqrt:ident, $pow:ident, $floor:ident, $ceil:ident, + $nvexp:ident, $nvlog:ident, $nvatan:ident, $nvsin:ident, $nvcos:ident, $nvasin:ident, $nvacos:ident, $nvatan2:ident) => { impl Float for $ty { + // `exp`/`ln` go through libdevice externs (like the trig + // methods below) rather than `core::intrinsics::expf32` &co: + // those intrinsic names were removed from newer nightlies, + // while the `__nv_*` symbols are stable across toolchains. #[inline(always)] - fn exp(self) -> $ty { unsafe { core::intrinsics::$exp(self) } } + fn exp(self) -> $ty { + unsafe extern "C" { fn $nvexp(x: $ty) -> $ty; } + unsafe { $nvexp(self) } + } #[inline(always)] - fn ln(self) -> $ty { unsafe { core::intrinsics::$ln(self) } } + fn ln(self) -> $ty { + unsafe extern "C" { fn $nvlog(x: $ty) -> $ty; } + unsafe { $nvlog(self) } + } + // These intrinsics still exist on every nightly cuda-oxide + // supports; newer toolchains made them safe, hence the allow. #[inline(always)] + #[allow(unused_unsafe)] fn sqrt(self) -> $ty { unsafe { core::intrinsics::$sqrt(self) } } #[inline(always)] + #[allow(unused_unsafe)] fn powf(self, n: $ty) -> $ty { unsafe { core::intrinsics::$pow(self, n) } } #[inline(always)] + #[allow(unused_unsafe)] fn floor(self) -> $ty { unsafe { core::intrinsics::$floor(self) } } #[inline(always)] + #[allow(unused_unsafe)] fn ceil(self) -> $ty { unsafe { core::intrinsics::$ceil(self) } } // abs/max/min via plain ops (no stable intrinsic names here). #[inline(always)] @@ -90,8 +107,8 @@ mod cuda_oxide_float { } }; } - float_impl!(f32, expf32, logf32, sqrtf32, powf32, floorf32, ceilf32, + float_impl!(f32, sqrtf32, powf32, floorf32, ceilf32, __nv_expf, __nv_logf, __nv_atanf, __nv_sinf, __nv_cosf, __nv_asinf, __nv_acosf, __nv_atan2f); - float_impl!(f64, expf64, logf64, sqrtf64, powf64, floorf64, ceilf64, + float_impl!(f64, sqrtf64, powf64, floorf64, ceilf64, __nv_exp, __nv_log, __nv_atan, __nv_sin, __nv_cos, __nv_asin, __nv_acos, __nv_atan2); } diff --git a/crates/khal-std/src/sync/barriers.rs b/crates/khal-std/src/sync/barriers.rs index 52f03d8..a7bedc1 100644 --- a/crates/khal-std/src/sync/barriers.rs +++ b/crates/khal-std/src/sync/barriers.rs @@ -35,9 +35,23 @@ pub fn workgroup_memory_barrier_with_group_sync() { { cuda_device::sync_threads(); } + // rust-cuda: link the LLVM intrinsic directly rather than calling + // `cuda_std::thread::sync_threads()`. That helper calls an external + // `__nvvm_block_barrier` from a bitcode library, so rustc's LLVM sees a + // plain (non-`convergent`) call and freely jump-threads / duplicates it. + // For `if tid < stride { .. } barrier()` tree reductions this yields + // `bar.sync` at different PCs for threads of the same warp, which + // deadlocks on sm_70+ (`bar.sync` is `.aligned`). The named intrinsic + // carries LLVM's `convergent` attribute, so it stays at the join point. #[cfg(all(target_arch = "nvptx64", not(feature = "cuda-oxide")))] { - cuda_std::thread::sync_threads(); + unsafe extern "C" { + #[link_name = "llvm.nvvm.barrier0"] + fn nvvm_barrier0(); + } + unsafe { + nvvm_barrier0(); + } } #[cfg(not(any( From e5cc8b221f9ae120268ae5179879ba1821d8443a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 18 Sep 2026 14:37:08 +0200 Subject: [PATCH 28/33] Fix MPM hang (rust-cuda) and NaNs (cuda-oxide) in multi-dim workgroups --- crates/cargo-cuda/src/main.rs | 8 ++++++++ crates/khal-derive/src/spirv_bindgen/cuda.rs | 8 +++++++- crates/khal-derive/src/spirv_bindgen/mod.rs | 1 + 3 files changed, 16 insertions(+), 1 deletion(-) diff --git a/crates/cargo-cuda/src/main.rs b/crates/cargo-cuda/src/main.rs index 812a7ee..c710811 100644 --- a/crates/cargo-cuda/src/main.rs +++ b/crates/cargo-cuda/src/main.rs @@ -285,6 +285,14 @@ fn build(args: BuildArgs) { "-Zcrate-attr=no_std", "-Zsaturating_float_casts=false", "-Cllvm-args=--override-libm", + // MIR JumpThreading knows nothing about convergence: when a kernel + // tests the same (uniform) condition twice with a workgroup barrier in + // between, it threads the second test and DUPLICATES the barrier onto + // the threaded path. Threads of one warp then reach different + // `bar.sync` instructions, which deadlocks on sm_70+ (nexus MPM P2G). + // LLVM's own passes respect `convergent`; MIR's does not, so turn it + // off (cuda-oxide's cargo-oxide does the same). + "-Zmir-enable-passes=-JumpThreading", ]; let rustflags = nvvm_flags.join("\x1f"); // Use unit separator for CARGO_ENCODED_RUSTFLAGS diff --git a/crates/khal-derive/src/spirv_bindgen/cuda.rs b/crates/khal-derive/src/spirv_bindgen/cuda.rs index 31c04d2..409e460 100644 --- a/crates/khal-derive/src/spirv_bindgen/cuda.rs +++ b/crates/khal-derive/src/spirv_bindgen/cuda.rs @@ -243,8 +243,11 @@ pub(super) fn generate_cuda_oxide_entry_block( func: &syn::ItemFn, original_params: &[OriginalParam], bindings: &[ShaderBinding], + workgroup_size: [u32; 3], cuda_entry_ident: &syn::Ident, ) -> proc_macro2::TokenStream { + let wg_x = workgroup_size[0]; + let wg_y = workgroup_size[1]; let all_params: Vec<&OriginalParam> = original_params .iter() .filter(|p| { @@ -373,10 +376,13 @@ pub(super) fn generate_cuda_oxide_entry_block( BuiltinKind::NumWorkgroups => { quote! { khal_std::arch::cuda::num_workgroups() } } + // Flattened like SPIR-V's LocalInvocationIndex + // (x + y * size_x + z * size_x * size_y): kernels declared + // with 2D/3D `threads(..)` are launched with 2D/3D blocks. BuiltinKind::LocalInvocationIndex => { quote! {{ let __tid = khal_std::arch::cuda::thread_idx(); - __tid.x + __tid.z * #wg_x * #wg_y + __tid.y * #wg_x + __tid.x }} } _ => quote! { Default::default() }, diff --git a/crates/khal-derive/src/spirv_bindgen/mod.rs b/crates/khal-derive/src/spirv_bindgen/mod.rs index 0b850cf..abc22ef 100644 --- a/crates/khal-derive/src/spirv_bindgen/mod.rs +++ b/crates/khal-derive/src/spirv_bindgen/mod.rs @@ -666,6 +666,7 @@ pub(crate) fn spirv_bindgen(attr: TokenStream, item: TokenStream) -> TokenStream &func, &original_params, &bindings, + workgroup_size, &cuda_entry_ident, ); From 2f2c1c737691f1928c470280534b88a68000d00c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 20 Sep 2026 10:07:53 +0200 Subject: [PATCH 29/33] GpuReadback: implement the CUDA arm (async device->pinned copy + event) --- crates/khal/src/backend/any_backend.rs | 95 +++++++++++++++++++++++++- 1 file changed, 94 insertions(+), 1 deletion(-) diff --git a/crates/khal/src/backend/any_backend.rs b/crates/khal/src/backend/any_backend.rs index 908323e..9618e7f 100644 --- a/crates/khal/src/backend/any_backend.rs +++ b/crates/khal/src/backend/any_backend.rs @@ -992,6 +992,11 @@ enum ReadbackState { /// handler once the copy has landed in the shared staging buffer. #[cfg(feature = "metal")] Metal(std::sync::Arc), + /// A CUDA event recorded right after the async device→host copy of the + /// staging buffer into the readback's pinned host mirror; complete once + /// the data is on the host. + #[cfg(feature = "cuda")] + Cuda(cudarc::driver::CudaEvent), } /// Non-blocking GPU→CPU readback of a small buffer region. @@ -1008,6 +1013,12 @@ pub struct GpuReadback { staging: GpuBuffer, len: usize, state: ReadbackState, + /// CUDA: page-locked host mirror of `staging`. Device memory is never + /// host-visible, so each request enqueues an async device→host copy into + /// this buffer (pinned memory keeps that copy asynchronous) and + /// [`try_take`](Self::try_take) polls the event recorded after it. + #[cfg(feature = "cuda")] + cuda_host: Option>, } // SAFETY: mirrors `MetalTimestamps` — the only non-`Send`/`Sync` payload is a @@ -1021,10 +1032,24 @@ impl GpuReadback { pub fn new(backend: &GpuBackend, len: usize) -> Result { let staging = backend.uninit_buffer::(len, BufferUsages::MAP_READ | BufferUsages::COPY_DST)?; + #[cfg(feature = "cuda")] + let cuda_host = match (&staging, backend) { + (GpuBuffer::Cuda(_), GpuBackend::Cuda(cuda)) if len > 0 => { + let nbytes = len * std::mem::size_of::(); + // SAFETY: the pinned memory is only ever read after a + // device→host copy of `nbytes` bytes has completed into it. + let host = unsafe { cuda.context().alloc_pinned::(nbytes) } + .map_err(|e| GpuBackendError::Cuda(e.into()))?; + Some(host) + } + _ => None, + }; Ok(Self { staging, len, state: ReadbackState::Idle, + #[cfg(feature = "cuda")] + cuda_host, }) } @@ -1094,8 +1119,33 @@ impl GpuReadback { } /// Sets up backend-specific completion tracking after the copy is submitted. - fn begin_completion(&self, backend: &GpuBackend) -> ReadbackState { + fn begin_completion(&mut self, backend: &GpuBackend) -> ReadbackState { match (&self.staging, backend) { + #[cfg(feature = "cuda")] + (GpuBuffer::Cuda(buffer), GpuBackend::Cuda(cuda)) => { + // Enqueue the device→host copy into the pinned mirror right + // behind the staging copy, then record an event to poll. Both + // are asynchronous (pinned destination), so nothing here waits + // on the GPU. If anything fails, fall back to a synchronous + // read of the staging buffer at take time. + let nbytes = self.len * std::mem::size_of::(); + match (buffer.inner(), self.cuda_host.as_mut()) { + (Some(inner), Some(host)) if nbytes > 0 => { + let stream = cuda.stream(); + let copied = stream.memcpy_dtoh(&inner.slice(..nbytes), host); + let event = copied.and_then(|()| { + stream.record_event(Some( + cudarc::driver::sys::CUevent_flags::CU_EVENT_DISABLE_TIMING, + )) + }); + match event { + Ok(done) => ReadbackState::Cuda(done), + Err(_) => ReadbackState::Ready, + } + } + _ => ReadbackState::Ready, + } + } #[cfg(feature = "webgpu")] (GpuBuffer::WebGpu(buffer), _) => { let (sender, receiver) = async_channel::bounded(1); @@ -1150,6 +1200,36 @@ impl GpuReadback { false } } + #[cfg(feature = "cuda")] + ReadbackState::Cuda(done) => { + if done.is_complete() { + self.read_cuda_host_into(out); + true + } else { + self.state = ReadbackState::Cuda(done); + false + } + } + } + } + + /// CUDA: copies the pinned host mirror (filled by the request's async + /// device→host copy, whose event has completed) into `out`. + #[cfg(feature = "cuda")] + fn read_cuda_host_into(&self, out: &mut [T]) { + let n = out.len().min(self.len); + if n == 0 { + return; + } + let Some(host) = self.cuda_host.as_ref() else { + return self.read_staging_into(out); + }; + // `as_slice` waits on the pinned slice's own transfer event, which is + // already complete here, so this does not block. + if let Ok(bytes) = host.as_slice() { + let dst: &mut [u8] = bytemuck::cast_slice_mut(&mut out[..n]); + let len = dst.len().min(bytes.len()); + dst[..len].copy_from_slice(&bytes[..len]); } } @@ -1174,6 +1254,19 @@ impl GpuReadback { ); } } + #[cfg(feature = "cuda")] + GpuBuffer::Cuda(buffer) => { + // Synchronous fallback (no pinned mirror / event): copy straight + // from the device staging buffer. Blocks until the stream has + // run the staging copy. + if let Some(inner) = buffer.inner() { + let dst: &mut [u8] = bytemuck::cast_slice_mut(&mut out[..n]); + let nbytes = dst.len().min(inner.len()); + let _ = inner + .stream() + .memcpy_dtoh(&inner.slice(..nbytes), &mut dst[..nbytes]); + } + } #[cfg(feature = "webgpu")] GpuBuffer::WebGpu(buffer) => { let slice = buffer.slice(..); From 6347dcddfdd46bd9daaa79fd854414c2e00c087c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 20 Sep 2026 16:52:19 +0200 Subject: [PATCH 30/33] CUDA: indirect dispatch without host sync (persistent grid + entry loop) --- README.md | 10 ++ crates/khal-derive/src/spirv_bindgen/cuda.rs | 164 +++++++++++++++---- crates/khal-example/Cargo.toml | 8 +- crates/khal/src/backend/any_backend.rs | 1 + crates/khal/src/backend/cuda.rs | 72 ++++++-- 5 files changed, 207 insertions(+), 48 deletions(-) diff --git a/README.md b/README.md index 995dcc3..5fd061c 100644 --- a/README.md +++ b/README.md @@ -109,6 +109,16 @@ cargo +nightly-2026-04-03 run --bin khal-example --features cuda-oxide (or `CUDA_OXIDE_TARGET=sm_XX`) to override it. A prebuilt PTX/cubin can be embedded instead of compiling by setting `CUDA_OXIDE_SHADERS_PTX_` to its path. +### Indirect dispatch on CUDA + +CUDA has no device-side indirect launch. Rather than reading the workgroup count back to the host +before every indirect dispatch (a full stream drain each time), the CUDA backend launches a fixed +number of resident blocks (`multiprocessor count × 16`, override with `KHAL_CUDA_PERSISTENT_BLOCKS`) +and the kernel entries generated by `#[spirv_bindgen]` loop over the *virtual* workgroups read from +the indirect-args buffer on the device. Direct dispatches run exactly one iteration per block, and the +loop count is uniform per block so workgroup barriers stay valid. `KHAL_CUDA_INDIRECT_SYNC=1` restores +the synchronous host readback for debugging. + ### PTX vs. cubin With either compiler, `khal-builder` then assembles the PTX into a **cubin** for the local GPU using the diff --git a/crates/khal-derive/src/spirv_bindgen/cuda.rs b/crates/khal-derive/src/spirv_bindgen/cuda.rs index 409e460..5941b92 100644 --- a/crates/khal-derive/src/spirv_bindgen/cuda.rs +++ b/crates/khal-derive/src/spirv_bindgen/cuda.rs @@ -96,25 +96,20 @@ pub(super) fn generate_cuda_entry_block( for param in &all_params { let name = ¶m.name; match ¶m.kind { + // Builtins derive from the VIRTUAL workgroup of the current loop + // iteration (see `virtual_workgroup_loop`), not the hardware block. OriginalParamKind::Builtin(kind) => { let arg = match kind { BuiltinKind::GlobalInvocationId => { - quote! { khal_std::arch::cuda::global_invocation_id() } - } - BuiltinKind::LocalInvocationId => { - quote! { khal_std::arch::cuda::local_invocation_id() } - } - BuiltinKind::WorkgroupId => { - quote! { khal_std::arch::cuda::workgroup_id() } - } - BuiltinKind::NumWorkgroups => { - quote! { khal_std::arch::cuda::num_workgroups() } + quote! { __khal_wg * __khal_block_dim + __khal_tid } } + BuiltinKind::LocalInvocationId => quote! { __khal_tid }, + BuiltinKind::WorkgroupId => quote! { __khal_wg }, + BuiltinKind::NumWorkgroups => quote! { __khal_num_wg }, BuiltinKind::LocalInvocationIndex => { - quote! {{ - let __tid = khal_std::arch::cuda::thread_idx(); - __tid.z * #wg_x * #wg_y + __tid.y * #wg_x + __tid.x - }} + quote! { + (__khal_tid.z * #wg_x * #wg_y + __khal_tid.y * #wg_x + __khal_tid.x) + } } _ => { quote! { Default::default() } @@ -216,14 +211,32 @@ pub(super) fn generate_cuda_entry_block( // Only generate if there are bindings (otherwise it's not a real kernel) if !sorted_bindings.is_empty() || !sorted_push_constants.is_empty() { + let (loop_prologue, wg_decl, iteration_barrier) = virtual_workgroup_loop(quote! { + if __khal_indirect_len >= 3 { + let __a = unsafe { + core::slice::from_raw_parts(__khal_indirect_ptr as *const u32, 3) + }; + Some([__a[0], __a[1], __a[2]]) + } else { + None + } + }); quote! { #[cfg(all(target_arch = "nvptx64", not(feature = "cuda-oxide")))] #[khal_std::cuda_std::kernel] pub unsafe fn #cuda_entry_ident( - #(#cuda_params),* + #(#cuda_params,)* + __khal_indirect_ptr: u64, + __khal_indirect_len: u64, ) { #(#cuda_body)* - #func_ident(#(#cuda_call_args),*); + #loop_prologue + while __khal_lin < __khal_total { + #wg_decl + #func_ident(#(#cuda_call_args),*); + __khal_lin += __khal_stride; + #iteration_barrier + } } } } else { @@ -357,37 +370,34 @@ pub(super) fn generate_cuda_oxide_entry_block( }) .collect(); - // Preludes for builtins and workgroup shared memory, in original order. + // Preludes for workgroup shared memory (once per kernel) and for builtins + // (once per virtual-workgroup iteration), in original order. let mut preludes: Vec = Vec::new(); + let mut builtin_preludes: Vec = Vec::new(); for p in &all_params { let name = &p.name; match &p.kind { + // Builtins derive from the VIRTUAL workgroup of the current loop + // iteration (see `virtual_workgroup_loop`), not the hardware block. OriginalParamKind::Builtin(kind) => { let expr = match kind { BuiltinKind::GlobalInvocationId => { - quote! { khal_std::arch::cuda::global_invocation_id() } - } - BuiltinKind::LocalInvocationId => { - quote! { khal_std::arch::cuda::local_invocation_id() } - } - BuiltinKind::WorkgroupId => { - quote! { khal_std::arch::cuda::workgroup_id() } - } - BuiltinKind::NumWorkgroups => { - quote! { khal_std::arch::cuda::num_workgroups() } + quote! { __khal_wg * __khal_block_dim + __khal_tid } } + BuiltinKind::LocalInvocationId => quote! { __khal_tid }, + BuiltinKind::WorkgroupId => quote! { __khal_wg }, + BuiltinKind::NumWorkgroups => quote! { __khal_num_wg }, // Flattened like SPIR-V's LocalInvocationIndex // (x + y * size_x + z * size_x * size_y): kernels declared // with 2D/3D `threads(..)` are launched with 2D/3D blocks. BuiltinKind::LocalInvocationIndex => { - quote! {{ - let __tid = khal_std::arch::cuda::thread_idx(); - __tid.z * #wg_x * #wg_y + __tid.y * #wg_x + __tid.x - }} + quote! { + (__khal_tid.z * #wg_x * #wg_y + __khal_tid.y * #wg_x + __khal_tid.x) + } } _ => quote! { Default::default() }, }; - preludes.push(quote! { let #name = #expr; }); + builtin_preludes.push(quote! { let #name = #expr; }); } OriginalParamKind::Workgroup => { // `#[spirv(workgroup)] x: &mut [T; N]` -> a function-local @@ -463,6 +473,13 @@ pub(super) fn generate_cuda_oxide_entry_block( &format!("{CUDA_OXIDE_KERNEL_PREFIX}{cuda_entry_ident}"), cuda_entry_ident.span(), ); + let (loop_prologue, wg_decl, iteration_barrier) = virtual_workgroup_loop(quote! { + if __khal_indirect_args.len() >= 3 { + Some([__khal_indirect_args[0], __khal_indirect_args[1], __khal_indirect_args[2]]) + } else { + None + } + }); quote! { // The "scoped Cargo cache protocol": cuda-oxide's `#[kernel]` records // the backend's codegen identity (output mode, arch, tool provenance) @@ -483,10 +500,26 @@ pub(super) fn generate_cuda_oxide_entry_block( // trivial kernels (nothing in-crate calls an entry point), and the // collector never sees them. #[unsafe(no_mangle)] - pub fn #prefixed_entry_ident(#(#kernel_params),*) { + pub fn #prefixed_entry_ident( + #(#kernel_params,)* + __khal_indirect_args: &[u32], + ) { #(#preludes)* #(#array_reconstructions)* - #body + #loop_prologue + // The shader body runs once per virtual workgroup. It lives in a + // closure so that an early `return` in the body ends the current + // iteration, not the whole loop. + let mut __khal_body = |__khal_wg: khal_std::glamx::UVec3| { + #(#builtin_preludes)* + #body + }; + while __khal_lin < __khal_total { + #wg_decl + __khal_body(__khal_wg); + __khal_lin += __khal_stride; + #iteration_barrier + } } } } @@ -496,3 +529,66 @@ pub(super) fn generate_cuda_oxide_entry_block( /// exactly `CUDA_ENTRY_POINT`. Keep in sync with the `cuda-device` revision /// pinned in `khal-std/Cargo.toml`. const CUDA_OXIDE_KERNEL_PREFIX: &str = "cuda_oxide_codegen_v1_cuda_oxide_kernel_246e25db_"; + +/// Virtual-workgroup loop shared by both CUDA entry flavours. +/// +/// CUDA has no device-side indirect dispatch. khal's CUDA backend therefore +/// launches indirect dispatches with a fixed number of resident blocks and +/// hands the entry the indirect-args `[u32; 3]` as a trailing `(ptr, len)` +/// slice (`(0, 0)` for direct dispatches). The entry loops over the virtual +/// workgroups `lin = block_id; lin < x*y*z; lin += real_grid_size`, so: +/// - a direct dispatch runs exactly one iteration per block (no overhead +/// beyond a few integer ops), +/// - an indirect dispatch runs the exact virtual workgroup count without any +/// host round trip, +/// - the trip count is uniform per block, so workgroup barriers in the body +/// remain valid. +/// +/// `read_args` must evaluate to `Option<[u32; 3]>` (the indirect args when +/// present). Returns `(prologue, wg_decl, iteration_barrier)`: the prologue +/// defines `__khal_tid`, `__khal_block_dim`, `__khal_num_wg` and the loop +/// bounds; the per-iteration declaration defines `__khal_wg` from +/// `__khal_lin`; the barrier separates consecutive virtual workgroups. +fn virtual_workgroup_loop( + read_args: proc_macro2::TokenStream, +) -> ( + proc_macro2::TokenStream, + proc_macro2::TokenStream, + proc_macro2::TokenStream, +) { + let prologue = quote! { + let __khal_tid = khal_std::arch::cuda::thread_idx(); + let __khal_block_dim = khal_std::arch::cuda::block_dim(); + let __khal_real_grid = khal_std::arch::cuda::num_workgroups(); + let __khal_num_wg: khal_std::glamx::UVec3 = match #read_args { + Some(a) => khal_std::glamx::UVec3::new(a[0], a[1], a[2]), + None => __khal_real_grid, + }; + let __khal_total = __khal_num_wg.x * __khal_num_wg.y * __khal_num_wg.z; + let __khal_stride = __khal_real_grid.x * __khal_real_grid.y * __khal_real_grid.z; + let __khal_bid = khal_std::arch::cuda::block_idx(); + let mut __khal_lin = __khal_bid.x + + __khal_bid.y * __khal_real_grid.x + + __khal_bid.z * __khal_real_grid.x * __khal_real_grid.y; + }; + let wg_decl = quote! { + let __khal_wg = khal_std::glamx::UVec3::new( + __khal_lin % __khal_num_wg.x, + (__khal_lin / __khal_num_wg.x) % __khal_num_wg.y, + __khal_lin / (__khal_num_wg.x * __khal_num_wg.y), + ); + }; + // A real workgroup starts with all of its threads synchronized and its + // shared memory untouched by anyone else. Between two virtual workgroups + // of the same block, fast threads would otherwise start writing the next + // iteration's shared-memory staging while slow threads still read the + // previous one's (nexus MPM G2P halo nodes: particles in the top rows of + // a block went astray once blocks looped). The condition is uniform per + // block, so the barrier is well-formed and skipped on the last iteration. + let iteration_barrier = quote! { + if __khal_lin < __khal_total { + khal_std::sync::workgroup_memory_barrier_with_group_sync(); + } + }; + (prologue, wg_decl, iteration_barrier) +} diff --git a/crates/khal-example/Cargo.toml b/crates/khal-example/Cargo.toml index c50fb62..6075b1a 100644 --- a/crates/khal-example/Cargo.toml +++ b/crates/khal-example/Cargo.toml @@ -12,8 +12,12 @@ cpu = ["khal/cpu", "khal-example-shaders/cpu"] cpu-parallel = ["cpu", "khal-example-shaders/cpu-parallel"] cuda = ["khal/cuda", "khal-builder/cuda", "khal-example-shaders/cuda"] # Native CUDA with kernels compiled by cuda-oxide (`cargo oxide`) rather than -# rust-cuda's `cargo cuda`. Implies `cuda`. -cuda-oxide = ["cuda", "khal-builder/cuda-oxide", "khal-example-shaders/cuda-oxide"] +# rust-cuda's `cargo cuda`. Implies `cuda`. Deliberately does NOT forward to +# the shader crate's `cuda-oxide` feature: khal-builder enables that one in its +# separate `cargo oxide` build only. In the host build it would swap khal-std's +# `Float`/atomics for GPU-only libdevice externs, which the CPU backend (`cpu`) +# cannot link. +cuda-oxide = ["cuda", "khal-builder/cuda-oxide"] [dependencies] khal = { version = "0.3.0", path = "../khal", features = ["derive"] } diff --git a/crates/khal/src/backend/any_backend.rs b/crates/khal/src/backend/any_backend.rs index 9618e7f..c475fbb 100644 --- a/crates/khal/src/backend/any_backend.rs +++ b/crates/khal/src/backend/any_backend.rs @@ -668,6 +668,7 @@ impl GpuPass { args: Vec::new(), #[cfg(feature = "push_constants")] push_constants: Vec::new(), + persistent_blocks: pass.persistent_blocks, }), #[cfg(feature = "metal")] (Self::Metal(pass), InnerGpuFunction::Metal(f)) => { diff --git a/crates/khal/src/backend/cuda.rs b/crates/khal/src/backend/cuda.rs index 3c353e0..8454be6 100644 --- a/crates/khal/src/backend/cuda.rs +++ b/crates/khal/src/backend/cuda.rs @@ -54,6 +54,11 @@ pub struct Cuda { /// Cache of loaded PTX modules keyed by content hash. /// Avoids re-loading the same PTX for every kernel in a crate. module_cache: Arc>>>, + /// Grid size (in blocks) used for indirect dispatches: CUDA has no + /// device-side indirect launch, so the generated kernel entries loop over + /// the *virtual* workgroups read from the indirect-args buffer with this + /// many resident blocks. See [`CudaDispatch::launch`]. + persistent_blocks: u32, } impl Cuda { @@ -64,10 +69,23 @@ impl Cuda { // captured into a CUDA graph (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED), // and all khal work is ordered on this one stream anyway. let stream = ctx.new_stream()?; + // Enough blocks to fill the device several times over for the small + // (64-thread) workgroups khal kernels use; `KHAL_CUDA_PERSISTENT_BLOCKS` + // overrides it. + let sms = ctx + .attribute(driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) + .unwrap_or(16) + .max(1) as u32; + let persistent_blocks = std::env::var("KHAL_CUDA_PERSISTENT_BLOCKS") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|n| *n > 0) + .unwrap_or(sms * 16); Ok(Self { ctx, stream, module_cache: Arc::new(Mutex::new(HashMap::new())), + persistent_blocks, }) } @@ -239,6 +257,8 @@ pub struct CudaFunction { /// is essentially a thin wrapper that holds a reference to the stream. pub struct CudaEncoder { pub(crate) stream: Arc, + /// See [`Cuda::persistent_blocks`]; forwarded to passes and dispatches. + pub(crate) persistent_blocks: u32, } /// CUDA compute pass. CUDA doesn't have explicit compute pass boundaries. @@ -246,6 +266,8 @@ pub struct CudaPass { pub(crate) stream: Arc, /// Optional timing: records an end event on drop and pushes to the shared pending list. timing: Option, + /// See [`Cuda::persistent_blocks`]. + pub(crate) persistent_blocks: u32, } impl Drop for CudaPass { @@ -284,6 +306,8 @@ pub struct CudaDispatch<'a> { pub(crate) args: Vec<(ShaderBinding, u64, u64)>, #[cfg(feature = "push_constants")] pub(crate) push_constants: Vec, + /// See [`Cuda::persistent_blocks`]. + pub(crate) persistent_blocks: u32, } impl<'a> CudaDispatch<'a> { @@ -458,6 +482,7 @@ impl Backend for Cuda { fn begin_encoding(&self) -> Self::Encoder { CudaEncoder { stream: self.stream.clone(), + persistent_blocks: self.persistent_blocks, } } @@ -472,6 +497,7 @@ impl Backend for Cuda { args: Vec::new(), #[cfg(feature = "push_constants")] push_constants: Vec::new(), + persistent_blocks: self.persistent_blocks, } } @@ -632,6 +658,7 @@ impl Encoder for CudaEncoder { CudaPass { stream: self.stream.clone(), timing, + persistent_blocks: self.persistent_blocks, } } @@ -670,6 +697,10 @@ impl<'a> Dispatch<'a, Cuda> for CudaDispatch<'a> { grid: impl Into>, block_dim: [u32; 3], ) -> Result<(), CudaBackendError> { + // Indirect-args slice handed to the kernel entry as a trailing + // `(ptr, len)` pair: `(0, 0)` for direct dispatches (the entry then + // uses the real grid), the device `[u32; 3]` for indirect ones. + let mut indirect_args: (u64, u64) = (0, 0); let grid_dim = match grid.into() { DispatchGrid::Grid(g) => g, DispatchGrid::ThreadCount(t) => [ @@ -678,25 +709,37 @@ impl<'a> Dispatch<'a, Cuda> for CudaDispatch<'a> { t[2].div_ceil(block_dim[2]), ], DispatchGrid::Indirect(buffer) => { - // CUDA doesn't support indirect dispatch natively. - // Read the 12-byte dispatch args from device memory. if std::env::var("KHAL_TRACE_INDIRECT").is_ok() { eprintln!("[khal] indirect dispatch: {}", self.function.name); } - self.stream.synchronize()?; - if let Some(ref inner) = buffer.inner { + let Some(ref inner) = buffer.inner else { + return Ok(()); + }; + if std::env::var_os("KHAL_CUDA_INDIRECT_SYNC").is_some() { + // Debug fallback: read the 12-byte dispatch args back to + // the host and launch the exact grid. Drains the stream on + // every indirect dispatch. + self.stream.synchronize()?; let bytes: Vec = self.stream.clone_dtoh(inner)?; - if bytes.len() >= 12 { - [ - u32::from_ne_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]), - u32::from_ne_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]), - u32::from_ne_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]), - ] - } else { + if bytes.len() < 12 { return Ok(()); } + [ + u32::from_ne_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]), + u32::from_ne_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]), + u32::from_ne_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]), + ] } else { - return Ok(()); + // CUDA has no device-side indirect launch. Instead of + // reading the workgroup count back (a full stream drain per + // dispatch), launch a fixed number of resident blocks: the + // generated entry reads the `[u32; 3]` on the device and + // loops over the virtual workgroups + // (`lin = block_id; lin < x*y*z; lin += grid_size`), so + // every block runs a uniform trip count and workgroup + // barriers stay valid. + indirect_args = (buffer.device_ptr_raw(), 3); + [self.persistent_blocks, 1, 1] } } }; @@ -743,6 +786,11 @@ impl<'a> Dispatch<'a, Cuda> for CudaDispatch<'a> { } } + // Trailing indirect-args `(ptr, len)` pair (see above); the generated + // CUDA entries always declare it. + param_values.push(indirect_args.0); + param_values.push(indirect_args.1); + // Launch using cudarc's launch_builder API. let mut builder = self.stream.launch_builder(&self.function.func); for value in ¶m_values { From 173cea5f02340024e3ec9d6c0b86472e89176959 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 26 Sep 2026 10:07:36 +0200 Subject: [PATCH 31/33] khal-builder: fall back to __nvcc_device_query and a cached arch when nvidia-smi fails --- crates/khal-builder/src/lib.rs | 46 ++++++++++++++++++++++++++++++++-- 1 file changed, 44 insertions(+), 2 deletions(-) diff --git a/crates/khal-builder/src/lib.rs b/crates/khal-builder/src/lib.rs index 902cc43..1f288da 100644 --- a/crates/khal-builder/src/lib.rs +++ b/crates/khal-builder/src/lib.rs @@ -352,7 +352,7 @@ impl KhalBuilder { /// The cargo target directory of the build that is running this build /// script, derived from `OUT_DIR` (`//build/-/out`). /// Falls back to a directory under `OUT_DIR` when the layout is unexpected. -#[cfg(feature = "cuda-oxide")] +#[cfg(feature = "cuda")] fn cargo_target_dir_from_out_dir() -> PathBuf { let out_dir = PathBuf::from(std::env::var_os("OUT_DIR").expect("OUT_DIR is set by cargo")); let looks_like_cargo_layout = out_dir @@ -476,9 +476,37 @@ fn find_cuda_tool(name: &str) -> Option { .find(|p| p.is_file()) } -/// `sm_XY` of the first local GPU as reported by `nvidia-smi`, if available. +/// `sm_XY` of the first local GPU. +/// +/// Tries, in order: `nvidia-smi` (NVML; breaks after a driver package update +/// until reboot), the toolkit's `__nvcc_device_query` (driver API, unaffected +/// by NVML), and finally the arch cached by a previous successful detection in +/// the cargo target dir, so a transient failure keeps building for the same +/// GPU instead of falling back to driver-JIT'd PTX. #[cfg(feature = "cuda")] fn detect_local_sm_arch() -> Option { + let cache = cargo_target_dir_from_out_dir().join("khal-cuda-arch"); + let detected = detect_sm_arch_nvidia_smi().or_else(detect_sm_arch_nvcc_device_query); + match detected { + Some(arch) => { + let _ = std::fs::write(&cache, &arch); + Some(arch) + } + None => { + let cached = std::fs::read_to_string(&cache).ok()?; + let cached = cached.trim(); + (!cached.is_empty()).then(|| { + println!( + "cargo:warning=khal-builder: GPU arch detection failed; using cached {cached}" + ); + cached.to_string() + }) + } + } +} + +#[cfg(feature = "cuda")] +fn detect_sm_arch_nvidia_smi() -> Option { let out = Command::new("nvidia-smi") .args(["--query-gpu=compute_cap", "--format=csv,noheader"]) .output() @@ -494,6 +522,20 @@ fn detect_local_sm_arch() -> Option { Some(format!("sm_{major}{minor}")) } +/// `__nvcc_device_query` (shipped with the toolkit next to `nvcc`) prints the +/// compute capability of the first device as e.g. `120`. +#[cfg(feature = "cuda")] +fn detect_sm_arch_nvcc_device_query() -> Option { + let tool = find_cuda_tool("__nvcc_device_query")?; + let out = Command::new(tool).output().ok()?; + if !out.status.success() { + return None; + } + let cap = String::from_utf8_lossy(&out.stdout); + let cap = cap.trim(); + (!cap.is_empty() && cap.chars().all(|c| c.is_ascii_digit())).then(|| format!("sm_{cap}")) +} + /// The `.ptx` cargo-oxide produced for the shader crate: `.ptx` /// (hyphens normalised to underscores), or the only `.ptx` in the directory. #[cfg(feature = "cuda-oxide")] From f130c14479f6403a1dd3d299a9bf1c8b05ed0835 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 26 Sep 2026 16:52:02 +0200 Subject: [PATCH 32/33] Backend-level CUDA graph capture (GpuGraph, begin/end_capture) --- README.md | 9 ++++ crates/khal/src/backend/any_backend.rs | 62 +++++++++++++++++++++++++- 2 files changed, 70 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 5fd061c..b8bbdd7 100644 --- a/README.md +++ b/README.md @@ -119,6 +119,15 @@ the indirect-args buffer on the device. Direct dispatches run exactly one iterat loop count is uniform per block so workgroup barriers stay valid. `KHAL_CUDA_INDIRECT_SYNC=1` restores the synchronous host readback for debugging. +### CUDA graphs + +`GpuBackend::begin_capture()` / `end_capture()` record every dispatch and copy issued in between into a +`GpuGraph` that `launch()` replays with a single driver call (CUDA only; other backends return +`GpuBackendError::Unsupported`). The captured region must be replay-safe: no buffer allocation, no host +readback or synchronization, no upload from pageable host memory, and host-side control flow is frozen +at capture time. nexus uses this to replay a whole physics frame (`NEXUS_CUDA_GRAPHS=1`, or the +"CUDA graphs" checkbox / `--cuda-graphs` flag of its testbed). + ### PTX vs. cubin With either compiler, `khal-builder` then assembles the PTX into a **cubin** for the local GPU using the diff --git a/crates/khal/src/backend/any_backend.rs b/crates/khal/src/backend/any_backend.rs index c475fbb..bd82521 100644 --- a/crates/khal/src/backend/any_backend.rs +++ b/crates/khal/src/backend/any_backend.rs @@ -2,7 +2,8 @@ use crate::backend::WebGpu; #[cfg(feature = "cuda")] use crate::backend::cuda::{ - Cuda, CudaBackendError, CudaBuffer, CudaBufferSlice, CudaDispatch as CudaDispatchInner, + CapturedGraph, Cuda, CudaBackendError, CudaBuffer, CudaBufferSlice, + CudaDispatch as CudaDispatchInner, CudaEncoder as CudaEncoderInner, CudaFunction as CudaFunctionInner, CudaModule as CudaModuleInner, CudaPass as CudaPassInner, CudaTimestamps, }; @@ -57,6 +58,43 @@ impl GpuBackend { } } + /// Whether [`Self::begin_capture`] / [`Self::end_capture`] are available + /// (CUDA graphs). WebGPU, Metal and CPU have no equivalent. + pub fn supports_graphs(&self) -> bool { + self.is_cuda() + } + + /// Starts recording every dispatch, buffer copy and event submitted to + /// this backend into a graph instead of executing it, until + /// [`Self::end_capture`]. The recorded work is NOT executed by the capture + /// itself: launch the returned [`GpuGraph`] to run it. + /// + /// The captured region must be replay-safe: no buffer allocation or + /// release, no host readback or synchronization, no upload from pageable + /// host memory, and every buffer address used must stay valid for as long + /// as the graph is replayed. Host-side control flow (dispatch counts, loop + /// trip counts) is frozen at capture time. + pub fn begin_capture(&self) -> Result<(), GpuBackendError> { + match self { + #[cfg(feature = "cuda")] + Self::Cuda(cuda) => Ok(cuda.begin_capture()?), + #[allow(unreachable_patterns)] + _ => Err(GpuBackendError::Unsupported("graph capture")), + } + } + + /// Ends the recording started by [`Self::begin_capture`] and returns the + /// instantiated graph. Errors if the capture was invalidated (an operation + /// that cannot be captured was issued) or recorded nothing. + pub fn end_capture(&self) -> Result { + match self { + #[cfg(feature = "cuda")] + Self::Cuda(cuda) => Ok(GpuGraph::Cuda(cuda.end_capture()?)), + #[allow(unreachable_patterns)] + _ => Err(GpuBackendError::Unsupported("graph capture")), + } + } + /// Returns `true` if this is the CUDA backend. Always available (returns /// `false` when the `cuda` feature is disabled) so callers need not gate. pub fn is_cuda(&self) -> bool { @@ -788,6 +826,28 @@ pub enum GpuBackendError { ContextNotFound, #[error("Noop backend error")] Noop, + #[error("operation not supported by this backend: {0}")] + Unsupported(&'static str), +} + +/// A captured sequence of GPU work that can be replayed with a single launch +/// (see [`GpuBackend::begin_capture`]). Only the CUDA backend can produce one. +pub enum GpuGraph { + #[cfg(feature = "cuda")] + Cuda(CapturedGraph), +} + +impl GpuGraph { + /// Replays the captured work on `backend`'s queue. Ordered with respect to + /// everything else submitted to that queue. + pub fn launch(&self, backend: &GpuBackend) -> Result<(), GpuBackendError> { + match (self, backend) { + #[cfg(feature = "cuda")] + (Self::Cuda(graph), GpuBackend::Cuda(_)) => Ok(graph.launch()?), + #[allow(unreachable_patterns)] + _ => Err(GpuBackendError::Unsupported("graph launch")), + } + } } /// Result of a single timed compute pass. From 78724f3ff2a7dc6e903e11f288b404417f6d8669 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Wed, 30 Sep 2026 18:54:49 +0200 Subject: [PATCH 33/33] chore: CI fixes --- crates/cargo-cuda/src/main.rs | 4 +- crates/khal-derive/src/spirv_bindgen/cuda.rs | 15 ++-- crates/khal-example-shaders/src/lib.rs | 7 +- crates/khal-std/src/index.rs | 40 ++++++++-- crates/khal-std/src/lib.rs | 16 +++- crates/khal-std/src/num_traits.rs | 83 +++++++++++++++----- crates/khal-std/src/sync/barriers.rs | 6 +- crates/khal-std/src/sync/subgroup.rs | 12 +-- crates/khal/src/backend/any_backend.rs | 6 +- crates/khal/src/backend/webgpu.rs | 18 +++-- 10 files changed, 143 insertions(+), 64 deletions(-) diff --git a/crates/cargo-cuda/src/main.rs b/crates/cargo-cuda/src/main.rs index c710811..0cd724e 100644 --- a/crates/cargo-cuda/src/main.rs +++ b/crates/cargo-cuda/src/main.rs @@ -196,7 +196,9 @@ components = ["llvm-tools-preview", "rust-src", "rustc-dev"] eprintln!(" prebuilt LLVM for Linux, so build it from source and set"); eprintln!(" LLVM_CONFIG=/path/to/llvm-7/bin/llvm-config (see the khal README)."); eprintln!(" - pkg-config + OpenSSL headers (openssl-sys), e.g. `pkg-config libssl-dev`."); - eprintln!(" - libclang with its resource headers for bindgen, e.g. `libclang-common--dev`."); + eprintln!( + " - libclang with its resource headers for bindgen, e.g. `libclang-common--dev`." + ); std::process::exit(1); } diff --git a/crates/khal-derive/src/spirv_bindgen/cuda.rs b/crates/khal-derive/src/spirv_bindgen/cuda.rs index 5941b92..7f2a045 100644 --- a/crates/khal-derive/src/spirv_bindgen/cuda.rs +++ b/crates/khal-derive/src/spirv_bindgen/cuda.rs @@ -316,7 +316,10 @@ pub(super) fn generate_cuda_oxide_entry_block( // reference in a prelude so the ABI stays aligned. let is_uniform = matches!( p.kind, - OriginalParamKind::Binding { is_uniform: true, .. } + OriginalParamKind::Binding { + is_uniform: true, + .. + } ); if is_uniform { return quote! { #name: #ty }; @@ -406,10 +409,9 @@ pub(super) fn generate_cuda_oxide_entry_block( // `SharedArray` viewed as `&mut T`. A real reference (not a // wrapper) so the body can index it directly AND pass it to // helpers that take `&mut [T; N]` / `&mut T`. - let static_name = - syn::Ident::new(&format!("__smem_{}", name), name.span()); - match &p.ty { - syn::Type::Reference(r) => match &*r.elem { + let static_name = syn::Ident::new(&format!("__smem_{}", name), name.span()); + if let syn::Type::Reference(r) = &p.ty { + match &*r.elem { // `&mut [f32; N]` shared tile -> a real `&mut [f32; N]` into // the `SharedArray`'s storage, so the body can index it, // call MaybeIndexUnchecked on it AND pass it verbatim to @@ -446,8 +448,7 @@ pub(super) fn generate_cuda_oxide_entry_block( let #name: &mut #elem_ty = &mut #sref_name[0]; }); } - }, - _ => {} + } } } _ => {} diff --git a/crates/khal-example-shaders/src/lib.rs b/crates/khal-example-shaders/src/lib.rs index 6edb7b6..9452b0d 100644 --- a/crates/khal-example-shaders/src/lib.rs +++ b/crates/khal-example-shaders/src/lib.rs @@ -17,7 +17,7 @@ pub fn add_assign( } } -/// 2D-grid probe: writes a tag encoding (batch=y, x, src[y]) per slot. +/// 2D-grid probe: writes a tag encoding (batch=y, x, `src[y]`) per slot. #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn grid_probe( @@ -31,7 +31,10 @@ pub fn grid_probe( if x >= cap || y as usize >= src.len() { return; } - out.write((y * cap + x) as usize, (y << 20) | (src.read(y as usize) << 10) | x); + out.write( + (y * cap + x) as usize, + (y << 20) | (src.read(y as usize) << 10) | x, + ); } /// Batched StepRng probe mimicking the nexus narrow-phase pattern: diff --git a/crates/khal-std/src/index.rs b/crates/khal-std/src/index.rs index a6bd075..9d09245 100644 --- a/crates/khal-std/src/index.rs +++ b/crates/khal-std/src/index.rs @@ -17,9 +17,15 @@ pub trait MaybeIndexUnchecked { impl MaybeIndexUnchecked for [T] { #[inline(always)] fn at(&self, id: usize) -> &T { - #[cfg(all(feature = "unsafe_remove_boundchecks", any(target_arch = "nvptx64", feature = "cuda-oxide")))] + #[cfg(all( + feature = "unsafe_remove_boundchecks", + any(target_arch = "nvptx64", feature = "cuda-oxide") + ))] return unsafe { self.get_unchecked(id) }; - #[cfg(all(feature = "unsafe_remove_boundchecks", not(any(target_arch = "nvptx64", feature = "cuda-oxide"))))] + #[cfg(all( + feature = "unsafe_remove_boundchecks", + not(any(target_arch = "nvptx64", feature = "cuda-oxide")) + ))] return unsafe { use spirv_std::arch::IndexUnchecked; self.index_unchecked(id) @@ -30,9 +36,15 @@ impl MaybeIndexUnchecked for [T] { #[inline(always)] fn at_mut(&mut self, id: usize) -> &mut T { - #[cfg(all(feature = "unsafe_remove_boundchecks", any(target_arch = "nvptx64", feature = "cuda-oxide")))] + #[cfg(all( + feature = "unsafe_remove_boundchecks", + any(target_arch = "nvptx64", feature = "cuda-oxide") + ))] return unsafe { self.get_unchecked_mut(id) }; - #[cfg(all(feature = "unsafe_remove_boundchecks", not(any(target_arch = "nvptx64", feature = "cuda-oxide"))))] + #[cfg(all( + feature = "unsafe_remove_boundchecks", + not(any(target_arch = "nvptx64", feature = "cuda-oxide")) + ))] return unsafe { use spirv_std::arch::IndexUnchecked; self.index_unchecked_mut(id) @@ -55,9 +67,15 @@ impl MaybeIndexUnchecked for [T] { impl MaybeIndexUnchecked for [T; N] { #[inline(always)] fn at(&self, id: usize) -> &T { - #[cfg(all(feature = "unsafe_remove_boundchecks", any(target_arch = "nvptx64", feature = "cuda-oxide")))] + #[cfg(all( + feature = "unsafe_remove_boundchecks", + any(target_arch = "nvptx64", feature = "cuda-oxide") + ))] return unsafe { self.get_unchecked(id) }; - #[cfg(all(feature = "unsafe_remove_boundchecks", not(any(target_arch = "nvptx64", feature = "cuda-oxide"))))] + #[cfg(all( + feature = "unsafe_remove_boundchecks", + not(any(target_arch = "nvptx64", feature = "cuda-oxide")) + ))] return unsafe { use spirv_std::arch::IndexUnchecked; self.index_unchecked(id) @@ -68,9 +86,15 @@ impl MaybeIndexUnchecked for [T; N] { #[inline(always)] fn at_mut(&mut self, id: usize) -> &mut T { - #[cfg(all(feature = "unsafe_remove_boundchecks", any(target_arch = "nvptx64", feature = "cuda-oxide")))] + #[cfg(all( + feature = "unsafe_remove_boundchecks", + any(target_arch = "nvptx64", feature = "cuda-oxide") + ))] return unsafe { self.get_unchecked_mut(id) }; - #[cfg(all(feature = "unsafe_remove_boundchecks", not(any(target_arch = "nvptx64", feature = "cuda-oxide"))))] + #[cfg(all( + feature = "unsafe_remove_boundchecks", + not(any(target_arch = "nvptx64", feature = "cuda-oxide")) + ))] return unsafe { use spirv_std::arch::IndexUnchecked; self.index_unchecked_mut(id) diff --git a/crates/khal-std/src/lib.rs b/crates/khal-std/src/lib.rs index 707aec3..495565e 100644 --- a/crates/khal-std/src/lib.rs +++ b/crates/khal-std/src/lib.rs @@ -59,13 +59,21 @@ pub mod cuda_oxide_glue { pub struct SmemBuf(pub &'static mut SharedArray); impl MaybeIndexUnchecked for SmemBuf { #[inline(always)] - fn read(&self, i: usize) -> T { self.0[i] } + fn read(&self, i: usize) -> T { + self.0[i] + } #[inline(always)] - fn write(&mut self, i: usize, v: T) { self.0[i] = v; } + fn write(&mut self, i: usize, v: T) { + self.0[i] = v; + } #[inline(always)] - fn at_mut(&mut self, i: usize) -> &mut T { &mut self.0[i] } + fn at_mut(&mut self, i: usize) -> &mut T { + &mut self.0[i] + } #[inline(always)] - fn at(&self, i: usize) -> &T { &self.0[i] } + fn at(&self, i: usize) -> &T { + &self.0[i] + } } impl core::ops::Index for SmemBuf { type Output = T; diff --git a/crates/khal-std/src/num_traits.rs b/crates/khal-std/src/num_traits.rs index ef5dc02..286c342 100644 --- a/crates/khal-std/src/num_traits.rs +++ b/crates/khal-std/src/num_traits.rs @@ -45,70 +45,115 @@ mod cuda_oxide_float { // while the `__nv_*` symbols are stable across toolchains. #[inline(always)] fn exp(self) -> $ty { - unsafe extern "C" { fn $nvexp(x: $ty) -> $ty; } + unsafe extern "C" { + fn $nvexp(x: $ty) -> $ty; + } unsafe { $nvexp(self) } } #[inline(always)] fn ln(self) -> $ty { - unsafe extern "C" { fn $nvlog(x: $ty) -> $ty; } + unsafe extern "C" { + fn $nvlog(x: $ty) -> $ty; + } unsafe { $nvlog(self) } } // These intrinsics still exist on every nightly cuda-oxide // supports; newer toolchains made them safe, hence the allow. #[inline(always)] #[allow(unused_unsafe)] - fn sqrt(self) -> $ty { unsafe { core::intrinsics::$sqrt(self) } } + fn sqrt(self) -> $ty { + unsafe { core::intrinsics::$sqrt(self) } + } #[inline(always)] #[allow(unused_unsafe)] - fn powf(self, n: $ty) -> $ty { unsafe { core::intrinsics::$pow(self, n) } } + fn powf(self, n: $ty) -> $ty { + unsafe { core::intrinsics::$pow(self, n) } + } #[inline(always)] #[allow(unused_unsafe)] - fn floor(self) -> $ty { unsafe { core::intrinsics::$floor(self) } } + fn floor(self) -> $ty { + unsafe { core::intrinsics::$floor(self) } + } #[inline(always)] #[allow(unused_unsafe)] - fn ceil(self) -> $ty { unsafe { core::intrinsics::$ceil(self) } } + fn ceil(self) -> $ty { + unsafe { core::intrinsics::$ceil(self) } + } // abs/max/min via plain ops (no stable intrinsic names here). #[inline(always)] - fn abs(self) -> $ty { if self < 0.0 { -self } else { self } } + fn abs(self) -> $ty { + if self < 0.0 { -self } else { self } + } #[inline(always)] - fn max(self, other: $ty) -> $ty { if self >= other { self } else { other } } + fn max(self, other: $ty) -> $ty { + if self >= other { self } else { other } + } #[inline(always)] - fn min(self, other: $ty) -> $ty { if self <= other { self } else { other } } + fn min(self, other: $ty) -> $ty { + if self <= other { self } else { other } + } #[inline(always)] fn atan(self) -> $ty { - unsafe extern "C" { fn $nvatan(x: $ty) -> $ty; } + unsafe extern "C" { + fn $nvatan(x: $ty) -> $ty; + } unsafe { $nvatan(self) } } #[inline(always)] fn sin(self) -> $ty { - unsafe extern "C" { fn $nvsin(x: $ty) -> $ty; } + unsafe extern "C" { + fn $nvsin(x: $ty) -> $ty; + } unsafe { $nvsin(self) } } #[inline(always)] fn cos(self) -> $ty { - unsafe extern "C" { fn $nvcos(x: $ty) -> $ty; } + unsafe extern "C" { + fn $nvcos(x: $ty) -> $ty; + } unsafe { $nvcos(self) } } #[inline(always)] fn asin(self) -> $ty { - unsafe extern "C" { fn $nvasin(x: $ty) -> $ty; } + unsafe extern "C" { + fn $nvasin(x: $ty) -> $ty; + } unsafe { $nvasin(self) } } #[inline(always)] fn acos(self) -> $ty { - unsafe extern "C" { fn $nvacos(x: $ty) -> $ty; } + unsafe extern "C" { + fn $nvacos(x: $ty) -> $ty; + } unsafe { $nvacos(self) } } #[inline(always)] fn atan2(self, x: $ty) -> $ty { - unsafe extern "C" { fn $nvatan2(y: $ty, x: $ty) -> $ty; } + unsafe extern "C" { + fn $nvatan2(y: $ty, x: $ty) -> $ty; + } unsafe { $nvatan2(self, x) } } } }; } - float_impl!(f32, sqrtf32, powf32, floorf32, ceilf32, __nv_expf, __nv_logf, - __nv_atanf, __nv_sinf, __nv_cosf, __nv_asinf, __nv_acosf, __nv_atan2f); - float_impl!(f64, sqrtf64, powf64, floorf64, ceilf64, __nv_exp, __nv_log, - __nv_atan, __nv_sin, __nv_cos, __nv_asin, __nv_acos, __nv_atan2); + float_impl!( + f32, + sqrtf32, + powf32, + floorf32, + ceilf32, + __nv_expf, + __nv_logf, + __nv_atanf, + __nv_sinf, + __nv_cosf, + __nv_asinf, + __nv_acosf, + __nv_atan2f + ); + float_impl!( + f64, sqrtf64, powf64, floorf64, ceilf64, __nv_exp, __nv_log, __nv_atan, __nv_sin, __nv_cos, + __nv_asin, __nv_acos, __nv_atan2 + ); } diff --git a/crates/khal-std/src/sync/barriers.rs b/crates/khal-std/src/sync/barriers.rs index a7bedc1..44da2ec 100644 --- a/crates/khal-std/src/sync/barriers.rs +++ b/crates/khal-std/src/sync/barriers.rs @@ -54,11 +54,7 @@ pub fn workgroup_memory_barrier_with_group_sync() { } } - #[cfg(not(any( - target_arch = "spirv", - target_arch = "nvptx64", - feature = "cuda-oxide" - )))] + #[cfg(not(any(target_arch = "spirv", target_arch = "nvptx64", feature = "cuda-oxide")))] #[cfg(feature = "cpu")] { crate::arch::cpu::barrier_wait(); diff --git a/crates/khal-std/src/sync/subgroup.rs b/crates/khal-std/src/sync/subgroup.rs index 8607fce..47bab84 100644 --- a/crates/khal-std/src/sync/subgroup.rs +++ b/crates/khal-std/src/sync/subgroup.rs @@ -15,11 +15,7 @@ pub fn subgroup_f_add(val: f32) -> f32 { { warp_reduce_add(val) } - #[cfg(not(any( - target_arch = "spirv", - target_arch = "nvptx64", - feature = "cuda-oxide" - )))] + #[cfg(not(any(target_arch = "spirv", target_arch = "nvptx64", feature = "cuda-oxide")))] { val } @@ -42,11 +38,7 @@ pub fn subgroup_f_max(val: f32) -> f32 { { warp_reduce_max(val) } - #[cfg(not(any( - target_arch = "spirv", - target_arch = "nvptx64", - feature = "cuda-oxide" - )))] + #[cfg(not(any(target_arch = "spirv", target_arch = "nvptx64", feature = "cuda-oxide")))] { val } diff --git a/crates/khal/src/backend/any_backend.rs b/crates/khal/src/backend/any_backend.rs index bd82521..7d1fc14 100644 --- a/crates/khal/src/backend/any_backend.rs +++ b/crates/khal/src/backend/any_backend.rs @@ -3,9 +3,9 @@ use crate::backend::WebGpu; #[cfg(feature = "cuda")] use crate::backend::cuda::{ CapturedGraph, Cuda, CudaBackendError, CudaBuffer, CudaBufferSlice, - CudaDispatch as CudaDispatchInner, - CudaEncoder as CudaEncoderInner, CudaFunction as CudaFunctionInner, - CudaModule as CudaModuleInner, CudaPass as CudaPassInner, CudaTimestamps, + CudaDispatch as CudaDispatchInner, CudaEncoder as CudaEncoderInner, + CudaFunction as CudaFunctionInner, CudaModule as CudaModuleInner, CudaPass as CudaPassInner, + CudaTimestamps, }; #[cfg(feature = "metal")] use crate::backend::metal::{ diff --git a/crates/khal/src/backend/webgpu.rs b/crates/khal/src/backend/webgpu.rs index 872ab89..995413b 100644 --- a/crates/khal/src/backend/webgpu.rs +++ b/crates/khal/src/backend/webgpu.rs @@ -80,7 +80,7 @@ pub struct WebGpuFunction { /// A WebGPU compute pass that carries its own device reference. /// /// On wasm the pass may be a MERGED pass shared by successive `begin_pass` -/// calls (see [`WebGpuEncoder::open_pass`]); dropping the handle then returns +/// calls (see `WebGpuEncoder::open_pass`, wasm only); dropping the handle then returns /// the underlying pass to the encoder's cache instead of ending it. pub struct WebGpuPass { pub(crate) pass: Option>, @@ -94,7 +94,11 @@ pub struct WebGpuPass { impl WebGpuPass { /// Begins a compute dispatch within this pass, binding the given function. pub fn begin_dispatch<'a>(&'a mut self, function: &'a WebGpuFunction) -> WebGpuDispatch<'a> { - WebGpuDispatch::new(&self.device, self.pass.as_mut().expect("pass ended"), function) + WebGpuDispatch::new( + &self.device, + self.pass.as_mut().expect("pass ended"), + function, + ) } } @@ -291,7 +295,7 @@ impl WebGpu { /// Pop a `MAP_READ | COPY_DST` staging buffer of exactly `bytes_len` bytes /// from the readback pool, allocating one if the pool is empty for that - /// size. Pairs with [`release_read_staging`]. Used by [`slow_read_buffer`] + /// size. Pairs with [`Self::release_read_staging`]. Used by [`Backend::slow_read_buffer`] /// to avoid allocating a fresh GPU buffer (and freeing it) on every call. fn acquire_read_staging(&self, bytes_len: u64) -> Buffer { if let Some(buf) = self @@ -1091,7 +1095,9 @@ async fn read_bytes(device: &Device, buffer: &Buffer) -> Result Result