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/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/CUDA_OXIDE_AUTOPORT.md b/CUDA_OXIDE_AUTOPORT.md new file mode 100644 index 0000000..779663f --- /dev/null +++ b/CUDA_OXIDE_AUTOPORT.md @@ -0,0 +1,241 @@ +# 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. +> +> **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/README.md b/README.md index fc3a270..b8bbdd7 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,92 @@ 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. + +### 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. + +### 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 +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 +152,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..0cd724e 100644 --- a/crates/cargo-cuda/src/main.rs +++ b/crates/cargo-cuda/src/main.rs @@ -190,6 +190,15 @@ 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); } @@ -278,6 +287,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-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 edf8c82..1f288da 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, @@ -93,7 +94,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 +141,8 @@ 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"); } fn build_spirv(&self, output_dir: impl AsRef) { @@ -163,9 +177,55 @@ 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 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(); + + 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) { + // 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}") + }); + 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![ @@ -194,4 +254,329 @@ 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")] +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. +/// +/// 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() + .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}")) +} + +/// `__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")] +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 93b64e2..7f2a045 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, @@ -94,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() } @@ -125,7 +122,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 +146,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 +163,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, ) }; }, @@ -214,17 +211,385 @@ 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(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),* + #(#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 { 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], + 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| { + 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, 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, 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 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_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! { + (__khal_tid.z * #wg_x * #wg_y + __khal_tid.y * #wg_x + __khal_tid.x) + } + } + _ => quote! { Default::default() }, + }; + builtin_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()); + 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 + // 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; + 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 #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` -> + // `&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 + // (`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_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) + // 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 + // 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,)* + __khal_indirect_args: &[u32], + ) { + #(#preludes)* + #(#array_reconstructions)* + #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 + } + } + } +} + +/// 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_"; + +/// 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-derive/src/spirv_bindgen/mod.rs b/crates/khal-derive/src/spirv_bindgen/mod.rs index 48a0854..abc22ef 100644 --- a/crates/khal-derive/src/spirv_bindgen/mod.rs +++ b/crates/khal-derive/src/spirv_bindgen/mod.rs @@ -660,6 +660,16 @@ 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, + workgroup_size, + &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 +830,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-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-shaders/src/lib.rs b/crates/khal-example-shaders/src/lib.rs index 167a116..9452b0d 100644 --- a/crates/khal-example-shaders/src/lib.rs +++ b/crates/khal-example-shaders/src/lib.rs @@ -16,3 +16,47 @@ 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..6075b1a 100644 --- a/crates/khal-example/Cargo.toml +++ b/crates/khal-example/Cargo.toml @@ -7,9 +7,17 @@ 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"] +# Native CUDA with kernels compiled by cuda-oxide (`cargo oxide`) rather than +# 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-example/shaders-spirv/batch_probe.spv b/crates/khal-example/shaders-spirv/batch_probe.spv new file mode 100644 index 0000000..083223f Binary files /dev/null and b/crates/khal-example/shaders-spirv/batch_probe.spv differ 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 0000000..0529959 Binary files /dev/null and b/crates/khal-example/shaders-spirv/grid_probe.spv differ 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-std/Cargo.toml b/crates/khal-std/Cargo.toml index eae18d9..726ac25 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,44 @@ 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 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"] +# 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. +# `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"] } 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. +# +# 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/NVlabs/cuda-oxide.git", rev = "62472763", optional = true } [lints] workspace = true @@ -38,5 +71,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..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", 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 +36,15 @@ 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 +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", 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 +86,15 @@ 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..495565e 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::{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))] /// Architecture-specific runtime support (CPU coroutines, CUDA intrinsics). pub mod arch; @@ -32,8 +39,56 @@ 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] + } + } + 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")))] 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..286c342 100644 --- a/crates/khal-std/src/num_traits.rs +++ b/crates/khal-std/src/num_traits.rs @@ -1,9 +1,159 @@ -// 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, $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 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 { $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)] + 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, + 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/atomics.rs b/crates/khal-std/src/sync/atomics.rs index a1629a2..078ce9a 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(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) }; @@ -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(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) }; @@ -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(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) }; @@ -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(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..44da2ec 100644 --- a/crates/khal-std/src/sync/barriers.rs +++ b/crates/khal-std/src/sync/barriers.rs @@ -8,15 +8,43 @@ 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. + // 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(); + } + // 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")))] { - // 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(); @@ -24,10 +52,9 @@ pub fn workgroup_memory_barrier_with_group_sync() { unsafe { nvvm_barrier0(); } - // 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..47bab84 100644 --- a/crates/khal-std/src/sync/subgroup.rs +++ b/crates/khal-std/src/sync/subgroup.rs @@ -11,11 +11,11 @@ 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 +34,26 @@ 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 +68,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 +79,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/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 bf720f6..7d1fc14 100644 --- a/crates/khal/src/backend/any_backend.rs +++ b/crates/khal/src/backend/any_backend.rs @@ -2,9 +2,10 @@ use crate::backend::WebGpu; #[cfg(feature = "cuda")] use crate::backend::cuda::{ - Cuda, CudaBackendError, CudaBuffer, CudaBufferSlice, CudaDispatch as CudaDispatchInner, - CudaEncoder as CudaEncoderInner, CudaFunction as CudaFunctionInner, - CudaModule as CudaModuleInner, CudaPass as CudaPassInner, CudaTimestamps, + CapturedGraph, Cuda, CudaBackendError, CudaBuffer, CudaBufferSlice, + CudaDispatch as CudaDispatchInner, CudaEncoder as CudaEncoderInner, + CudaFunction as CudaFunctionInner, CudaModule as CudaModuleInner, CudaPass as CudaPassInner, + CudaTimestamps, }; #[cfg(feature = "metal")] use crate::backend::metal::{ @@ -57,10 +58,160 @@ impl GpuBackend { } } - /// Returns `true` if this is the CUDA backend. - #[cfg(feature = "cuda")] + /// 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 { - 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 { + // 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() { + 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. @@ -555,6 +706,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)) => { @@ -674,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. @@ -879,6 +1053,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. @@ -895,6 +1074,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 @@ -908,10 +1093,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, }) } @@ -981,8 +1180,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); @@ -1037,6 +1261,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]); } } @@ -1061,6 +1315,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(..); @@ -1085,6 +1352,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; @@ -1633,7 +1908,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")] @@ -1665,7 +1948,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")] @@ -1702,7 +1989,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..8454be6 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 ─────────────────────────────────────────────────────── @@ -20,17 +54,38 @@ 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 { /// 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()?; + // 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, }) } @@ -39,10 +94,80 @@ 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)?; + // 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 }) + } +} + +/// 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 +180,8 @@ pub enum CudaBackendError { Driver(#[from] driver::DriverError), #[error("Invalid PTX module")] InvalidPtx, + #[error("CUDA stream capture produced no graph")] + CaptureFailed, } // ── Buffer ───────────────────────────────────────────────────────────── @@ -121,6 +248,7 @@ pub struct CudaModule { #[derive(Clone)] pub struct CudaFunction { pub(crate) func: driver::CudaFunction, + pub(crate) name: String, } // ── Encoder / Pass ───────────────────────────────────────────────────── @@ -129,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. @@ -136,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 { @@ -174,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> { @@ -318,8 +452,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( @@ -339,6 +482,7 @@ impl Backend for Cuda { fn begin_encoding(&self) -> Self::Encoder { CudaEncoder { stream: self.stream.clone(), + persistent_blocks: self.persistent_blocks, } } @@ -353,6 +497,7 @@ impl Backend for Cuda { args: Vec::new(), #[cfg(feature = "push_constants")] push_constants: Vec::new(), + persistent_blocks: self.persistent_blocks, } } @@ -384,6 +529,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 { @@ -410,6 +567,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 { @@ -486,6 +658,7 @@ impl Encoder for CudaEncoder { CudaPass { stream: self.stream.clone(), timing, + persistent_blocks: self.persistent_blocks, } } @@ -524,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) => [ @@ -532,22 +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. - self.stream.synchronize()?; - if let Some(ref inner) = buffer.inner { + if std::env::var("KHAL_TRACE_INDIRECT").is_ok() { + eprintln!("[khal] indirect dispatch: {}", self.function.name); + } + 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] } } }; @@ -594,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 { @@ -607,9 +804,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/metal.rs b/crates/khal/src/backend/metal.rs index 1efeb67..4906c19 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) @@ -735,12 +815,29 @@ 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 .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 +876,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 +885,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 +932,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 +961,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); 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 1fae606..995413b 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::{ @@ -77,15 +78,36 @@ 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`, 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: 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(); + } } } @@ -93,6 +115,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. @@ -110,6 +177,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 +266,7 @@ impl WebGpu { hacks: vec![], spirv_passthrough_enabled, timestamp_supported, + read_staging_pool: Arc::new(Mutex::new(HashMap::new())), }) } @@ -213,7 +289,41 @@ 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 [`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 + .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. @@ -282,6 +392,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 +404,7 @@ impl WebGpu { self.device .create_shader_module_passthrough(wgpu::ShaderModuleDescriptorPassthrough { spirv: Some(spirv), + entry_points: Cow::Owned(entry_points), ..Default::default() }) }; @@ -305,6 +420,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 { @@ -509,6 +685,8 @@ impl Backend for WebGpu { .device .create_command_encoder(&CommandEncoderDescriptor::default()), device: self.device.clone(), + #[cfg(target_arch = "wasm32")] + open_pass: Default::default(), } } @@ -521,6 +699,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(()) } @@ -576,6 +760,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)); @@ -628,18 +814,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 } } @@ -650,6 +843,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()) { @@ -658,11 +852,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, } } @@ -674,6 +892,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, @@ -768,10 +991,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]); } @@ -783,6 +1017,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]); } @@ -859,6 +1094,10 @@ async fn read_bytes(device: &Device, buffer: &Buffer) -> Result Result