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-builder/src/lib.rs b/crates/khal-builder/src/lib.rs index edf8c82..52ffa58 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) { @@ -163,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-derive/src/spirv_bindgen/cuda.rs b/crates/khal-derive/src/spirv_bindgen/cuda.rs index 93b64e2..955153a 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,239 @@ 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_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; + }); + } + // `&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_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]; + }); + } + }, + _ => {} + } + } + _ => {} + } + } + + 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(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-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 a947d0d..20ff8dd 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,12 +25,35 @@ 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", "dep:cuda-device"] + [dependencies] glamx = { version = "0.3", default-features = false, features = ["nostd-libm", "bytemuck", "u32", "i32", "f64"] } rayon = { version = "1", optional = true } corosensei = { version = "0.3", optional = true } spirv-std-macros = "0.10.0-alpha.1" khal-derive = { version = "0.2.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 @@ -38,5 +62,7 @@ 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 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 } diff --git a/crates/khal-std/src/arch/cuda.rs b/crates/khal-std/src/arch/cuda.rs index 251b611..d56531b 100644 --- a/crates/khal-std/src/arch/cuda.rs +++ b/crates/khal-std/src/arch/cuda.rs @@ -1,8 +1,45 @@ //! 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`: `cuda_device::thread::*` accessors, recognized by name by +//! the cuda-oxide MIR importer. `cuda_std` does not build under cuda-oxide. -use cuda_std::thread; use glamx::UVec3; +#[cfg(not(feature = "cuda-oxide"))] +use cuda_std::thread; + +/// 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, $dev:ident) => { + #[inline(always)] + pub fn $name() -> u32 { + cuda_device::thread::$dev() + } + }; + } + 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`. #[inline(always)] pub fn thread_idx() -> UVec3 { @@ -55,8 +92,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/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 dd39cbd..8a404dd 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(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; @@ -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; +// 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(feature = "cuda-oxide", not(target_arch = "spirv")))] +pub mod cuda_oxide_glue { + use crate::index::MaybeIndexUnchecked; + pub use cuda_device::SharedArray; + + pub struct SmemBuf(pub &'static mut SharedArray); + impl MaybeIndexUnchecked for SmemBuf { + #[inline(always)] + fn read(&self, i: usize) -> T { self.0[i] } + #[inline(always)] + 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] } + #[inline(always)] + fn at(&self, i: usize) -> &T { &self.0[i] } + } +} #[cfg(not(any(target_arch = "spirv", target_arch = "nvptx64")))] pub use std::println; 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 d7f52c2..80ce764 100644 --- a/crates/khal-std/src/num_traits.rs +++ b/crates/khal-std/src/num_traits.rs @@ -1,9 +1,97 @@ -// 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"))] +#[cfg(not(any(target_arch = "nvptx64", feature = "cuda-oxide")))] pub use spirv_std::num_traits::Float; + +#[cfg(all(feature = "cuda-oxide", not(target_arch = "spirv")))] +pub use cuda_oxide_float::Float; + +#[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`. + 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; + 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, $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) } } + #[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) } + } + #[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, __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 8ca0887..af05a47 100644 --- a/crates/khal-std/src/sync/atomics.rs +++ b/crates/khal-std/src/sync/atomics.rs @@ -93,7 +93,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(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", feature = "cuda-oxide")))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -112,7 +120,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(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", feature = "cuda-oxide")))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *const u32 as *const AtomicU32) }; @@ -273,7 +288,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(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", feature = "cuda-oxide")))] { use core::sync::atomic::{AtomicU32, Ordering}; let atomic = unsafe { &*(ptr as *mut u32 as *const AtomicU32) }; @@ -292,7 +313,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(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", 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); diff --git a/crates/khal/src/backend/any_backend.rs b/crates/khal/src/backend/any_backend.rs index d41c1a9..68f56cf 100644 --- a/crates/khal/src/backend/any_backend.rs +++ b/crates/khal/src/backend/any_backend.rs @@ -57,10 +57,104 @@ 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. + 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] NVIDIA GPU present but this binary was built without the `cuda` feature; using WebGPU" + ); + } + } + + eprintln!("[khal] backend = WebGPU"); + Ok(Self::WebGpu(WebGpu::new(features, limits).await?)) } /// Returns `true` if this is the Metal backend. @@ -1084,6 +1178,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; @@ -1632,7 +1734,15 @@ 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")] @@ -1664,7 +1774,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); + dispatch.set_arg( + binding, + slice.offset_ptr(), + slice.byte_len / std::mem::size_of::() as u64, + ); Ok(()) } #[cfg(feature = "metal")] @@ -1701,7 +1815,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); + dispatch.set_arg( + binding, + slice.offset_ptr(), + slice.byte_len / std::mem::size_of::() as u64, + ); Ok(()) } #[cfg(feature = "metal")] diff --git a/crates/khal/src/backend/cuda.rs b/crates/khal/src/backend/cuda.rs index 2cba404..568fcb2 100644 --- a/crates/khal/src/backend/cuda.rs +++ b/crates/khal/src/backend/cuda.rs @@ -8,7 +8,41 @@ 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 ─────────────────────────────────────────────────────── @@ -26,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, @@ -39,10 +76,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 +150,8 @@ pub enum CudaBackendError { Driver(#[from] driver::DriverError), #[error("Invalid PTX module")] InvalidPtx, + #[error("CUDA stream capture produced no graph")] + CaptureFailed, } // ── Buffer ───────────────────────────────────────────────────────────── @@ -121,6 +218,7 @@ pub struct CudaModule { #[derive(Clone)] pub struct CudaFunction { pub(crate) func: driver::CudaFunction, + pub(crate) name: String, } // ── Encoder / Pass ───────────────────────────────────────────────────── @@ -318,8 +416,17 @@ impl Backend for Cuda { entry_point: &str, _push_constant_size: u32, ) -> Result { - let func = module.inner.load_function(entry_point)?; - Ok(CudaFunction { func }) + 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(), + }) } fn load_function_with_layouts( @@ -534,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)?; @@ -607,9 +717,42 @@ 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 + ); + } + 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), + 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 f4921c5..404ce11 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}; #[cfg(feature = "push_constants")] use wgpu::PushConstantRange; use wgpu::util::{BufferInitDescriptor, DeviceExt}; @@ -112,6 +113,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 { @@ -193,6 +202,7 @@ impl WebGpu { hacks: vec![], spirv_passthrough_enabled, timestamp_supported, + read_staging_pool: Arc::new(Mutex::new(HashMap::new())), }) } @@ -215,9 +225,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)); @@ -638,18 +682,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 } }