Repository navigation
feat: add support for cuda-oxide (without breaking rust-cuda) - #13
Merged
Merged
Conversation
This was referenced Sep 28, 2026
Adds a cuda-oxide code path so khal shader crates compile to PTX/cubin via the cuda-oxide Rust→PTX backend (alternative to Rust-CUDA's rustc_codegen_nvvm, which is wall-walled on LLVM 7), enabling `vortx --features cuda` and the nexus3d physics shaders to run on native CUDA. - khal-derive/spirv_bindgen: generate a `<fn>_cuda_entry_<hash>` cuda-oxide kernel (typed params, inlined body, builtins → khal_std::arch::cuda, `#[spirv(workgroup)]` → SharedArray/SmemBuf). CRITICAL ABI fix: the host pushes a (ptr, byte_len) pair for every storage binding, so sized-array (`&[T; N]`) and scalar (`&T`) storage params are received as slices and the original reference rebuilt in a prelude — otherwise they'd consume one fewer kernel arg than the host pushes and shift every later binding (illegal-address at runtime). Uniforms stay single-pointer. - khal-std: `cuda-oxide` feature — own `llvm.nvvm.*` sreg/barrier intrinsics, a libdevice-backed `Float`, `SmemBuf<T,N>` shared-memory glue, generic `MaybeIndexUnchecked`; `cuda_std` made optional (it doesn't build under cuda-oxide's `-Zbuild-std`). - khal/backend/cuda: load a pre-linked CUBIN (ELF magic) in addition to PTX text — libdevice kernels (exp/sqrt/…) must arrive as a cubin since their PTX carries unresolved `__nv_*` externs. Optional `KHAL_CUDA_TRACE` launch trace.
…ubin/host hash mismatches)
glam 0.33 dropped UVec3/UVec4 under default-features=false, which breaks spirv-std 0.10.0-alpha.1's compile (112 errors). The `cargo gpu build` sub-invocation re-resolves without the workspace lock and would otherwise pick the newest 0.33.x via spirv-std's open `glam >=0.30.8`. Pin to 0.32.1, the newest pre-0.33 version that still works (mirrors the pin in vortx-shaders).
cuda-oxide lowers a `&[T]` kernel parameter to a `(ptr, len)` pair where `len` is an ELEMENT count, but the khal CUDA backend was pushing the buffer's BYTE length. So a kernel calling `slice.len()` got `byte_len` (4x too large for `&[u32]`) and read out of bounds — fatal in the batched broad-phase radix sort (`gpu_init_sort_dispatch` derives the workgroup grid from `num_keys_arr.len()`, producing a garbage indirect dispatch and an illegal memory access) and in the LBVH traversal (`*_len.len()`). Fix: in the three CUDA `write_arg` arms (`GpuBuffer`/`GpuBufferSlice`/ `GpuBufferSliceMut`, all generic over `T`), push `byte_len / size_of::<T>()`. Sized arrays (`&[T; N]`), scalars (`&T`) and uniforms reconstruct via the pointer and ignore this value, so they are unaffected; only true `&[T]` slices change — to correct. Latent in vortx (its kernels bound loops by `Shape` uniforms, never `slice.len()`) and in single-env physics; surfaced once N>1 batched physics ran on native CUDA.
…stream); khal-builder: KHAL_SKIP_SPIRV skip for cuda-only builds
…capture (begin/end_capture, CapturedGraph), KHAL_CUDA_PROFILE kernel profiling, prebuilt-PTX escape hatch (CUDA_OXIDE_SHADERS_PTX_*)
… cuda-oxide - num_traits: Float gains sin/cos/asin/acos/atan2 via __nv_* libdevice externs (nexus shaders route trig through Float on the device). - SmemBuf is generic over the element type (SmemBuf<T, N>) — nexus/vortx workgroup tiles are u32/Vec3/i32, not just f32. - atomics: on nvptx64+cuda-oxide, load/store go through cuda-device atomics; core's AtomicU32 load/store keep a reachable AcqRel panic arm that the cuda-oxide backend rejects. - Cargo.toml: wire the cuda-device path dep for the cuda-oxide feature (machine-local path, as documented in the manifest comment).
Rework the cuda-oxide backend path to compile through upstream cuda-oxide's supported model — unified interception on the HOST target — instead of a custom --target nvptx64 device build. Gates move from cfg(target_arch = "nvptx64") to cfg(feature = "cuda-oxide"), and every raw LLVM construct is replaced by the cuda-device API the importer recognizes by name on any compilation target: - sregs: `llvm.nvvm.read.ptx.sreg.*` link-name externs -> cuda_device::thread - barrier: `llvm.nvvm.barrier0` extern -> cuda_device::sync_threads - warp shuffles: nvptx inline asm -> cuda_device::warp::shuffle_xor_f32 - cuda-device/num-traits become plain (not target-gated) optional deps khal-derive generates the CUDA entries with cuda-oxide's reserved kernel prefix directly (cuda_oxide_kernel_246e25db_<entry>) instead of the heuristic otherwise skips trivial kernels that nothing in-crate calls). khal-std deliberately does NOT re-export cuda-device at the crate root: the khal_std::cuda_device::* visible path breaks the importer's by-name dispatch. Shader crates now build with plain cargo build -p <shaders> --features "cuda-oxide ... glamx/scalar-math" (host target, glam scalar math instead of host SIMD). Verified: nexus + vortx shader crates compile against upstream cuda-oxide main (+ the two open fix PRs), cubins assemble, and pendulum PPO / biped training on an RTX 5080 match the device-build reward trajectories exactly. Bonus: the vortx obs/reward kernels (previously excluded from device builds) now compile too.
Three hardenings in KhalGpuBackend::auto: - An explicit KHAL_BACKEND=cuda on a binary built without the `cuda` feature is now a hard error instead of a warning + WebGPU fallback: an explicitly requested backend must not be silently substituted (results and performance would differ from what the user asked to measure). - Auto-detect CUDA init failure now prints the error before falling back to WebGPU; on a machine with an NVIDIA GPU a quiet fallback usually hides real breakage behind a slower working run. - A cuda-less binary auto-detecting on a machine with an NVIDIA GPU (/dev/nvidiactl present) now says so — a stale or misbuilt binary otherwise degrades silently to WebGPU.
…on the backend enum - The CUDA backend ran on cudarc's default_stream(), i.e. the legacy NULL stream, which CUDA refuses to capture (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). Use an owned non-blocking stream instead; all khal work is ordered on this one stream anyway. - GpuBackend (the enum) never overrode Backend::as_cuda, so the graph-capture API was unreachable through it (always None). Delegate to the Cuda variant. - KHAL_TRACE_INDIRECT=1 prints each indirect dispatch's kernel name: on CUDA every indirect dispatch costs a full stream drain + host count read, and invalidates a graph capture — this is the tool for hunting them down.
…path (#5) khal-std pinned cuda-device at path = "/home/xavier/cuda-oxide-src/crates/cuda-device" an absolute path into one developer's home directory. Cargo reads path dependencies eagerly, even optional ones that no active feature selects, so this did not merely break `--features cuda-oxide` -- it broke `cargo metadata` for the whole graph on every other machine: error: failed to load source for dependency `cuda-device` Caused by: Unable to update /home/xavier/cuda-oxide-src/crates/cuda-device Caused by: No such file or directory (os error 2) Found while building nexus (which path-patches khal) on a fresh box: the build is unreachable off that one machine. This also defeats exactly what the feature comment two lines up promises -- "This manifest keeps the feature host-buildable everywhere else." The intent was clearly that the dependency be added locally on the build box and not committed; the local edit escaped into the branch, and the comment describing the intent stayed behind. cuda-device is not on crates.io, so a git rev is the portable way to say this. Pinned rather than floating on a branch, so the resolve is reproducible. Hacking on both trees still works -- add a [patch] with a local checkout in the consuming workspace. Verified with the local cuda-oxide checkout deleted, so no stale path could satisfy it: cargo metadata resolves clean, and with `-p khal-std --features cuda-oxide` cargo fetches the rev and records source = "git+https://github.com/haixuanTao/cuda-oxide.git?rev=8c3b5e2..." cuda-device then still needs nightly (#![feature(f16)]; cuda-oxide pins nightly-2026-04-03 in rust-toolchain.toml) -- that is pre-existing and identical under the old path dep, not a regression from this change. Co-authored-by: haixuanTao <shavtao@gmail.com>
The fused solver kernels index workgroup buffers directly (vels_smem[i]); the SPIR-V path gets that via the array type, the cuda-oxide SmemBuf wrapper only had the MaybeIndexUnchecked methods. Claude-Session: https://claude.ai/code/session_01U2n9RqmxTJb8UG5d1Sjw4W
…nd/swap) The load/store fns already routed through cuda_device atomics on cuda-oxide; the RMW fns still fell back to core::sync::atomic, whose fetch_* lower to AtomicRmwOp — rejected by the legacy-NVVM legalizer. cuda_device's fetch_* lower to atom.* intrinsics. Claude-Session: https://claude.ai/code/session_01U2n9RqmxTJb8UG5d1Sjw4W
…min/or/and/swap)" This reverts commit 9f15a5d.
The env short-circuit declared rerun-if-env-changed only, so replacing the cubin file at the same path silently kept the stale embedded shaders. Claude-Session: https://claude.ai/code/session_01B7NC7U2wDx2tSmT9mkeF9h
KHAL_CUDA_GRAPH_DOT=<path> dumps the captured graph via cuGraphDebugDotPrint (node-type census: MEM_ALLOC nodes explain relaunch INVALID_VALUE). KHAL_CUDA_ALLOC_TRACE=1 backtraces any buffer allocation recorded while the stream is capturing. Both no-ops unless the env vars are set. Claude-Session: https://claude.ai/code/session_01B7NC7U2wDx2tSmT9mkeF9h
The fork pin predates the unified conversion; the API surface consumed here (kernel macro, SharedArray, sync_threads) is upstream. Consuming workspaces still [patch] to a local checkout while two fixes (NVlabs #518, #519) are pending upstream.
The early-return prebuilt branch copied into OUT_DIR/shaders-spirv without creating it — masked whenever the SPIR-V build ran first, breaks cuda-feature builds that hit the prebuilt path before any spirv step.
Counts submits, real pass opens (merged-cache hits excluded), encoder-level copies, write_buffer uploads and map requests, behind take-and-zero accessors. This is how the zealot demo's "is there a hidden readback" question got answered with numbers (steady state: copies=0 writes=0 maps=0, ~450 submits/s) instead of a guess. Rides on the existing (uncommitted until now) wasm pass-merging in this file.
The native Metal backend produced wrong physics on multi-env scenes: the zealot biped ran every env with the LAST env's robot state. Root cause (proven by bisection down to a 4-env repro): `write_buffer` memcpys into shared memory immediately, racing in-flight kernel reads — wgpu's `queue.write_buffer` is ordered after submitted work. Staged per-env uploads corrupted accordingly. - write_buffer: while work may be in flight (`gpu_busy`, set on submit), route bytes through a staging buffer + ONE lazily-committed blit command buffer per write burst (`pending_writes`, flushed by the next submit/synchronize). Queue FIFO order = wgpu's semantics, no CPU stall. Direct-memcpy fast path only before any submit. - MSL translation now matches wgpu's WebGPU semantics: bounds checks ReadZeroSkipWrite (kernels rely on zero-read/skip-write for by-design OOB accesses), zero-initialized workgroup memory, force_loop_bounding, fast math OFF, MSL 3.0. - uninit_buffer zero-fills (WebGPU guarantees zeroed buffers). - KHAL_BACKEND=metal|msl now selects the backend in GpuBackend::auto (never auto-picked; errors when built without the feature). - Warn when a pipeline's maxTotalThreadsPerThreadgroup is below the shader's declared workgroup size (dispatching more is Metal UB). - KHAL_METAL_WGSL_ROUNDTRIP=1 debug lever: normalize IR through a WGSL round-trip before MSL generation (entry-point names restored). - khal-example: grid_probe/batch_probe 2D-dispatch validation kernels. Verified against WebGPU on the zealot biped (M4 Pro): bare-physics and multi-env terrain contact sets now match; zero-action survival curves identical (61% @1s). Perf at 1024 envs: 8.7 s/iter vs WebGPU 8.2 — the previously observed 6.5 was the race illegally overlapping uploads with compute. gpuwait dominates (270 vs 76 ms/step); closing that (concurrent dispatch + explicit hazards) is future work. Claude-Session: https://claude.ai/code/session_01EvbTzx6ZJNtcAKwVQ1B68i
All kernels use 64-thread (2x SIMD-width) groups dispatched at exactly that size; the hint lets Metal drop per-dispatch edge masking. Measured ~3% on zealot biped training @1024 envs (8.7-8.8 -> 8.4-8.6 s/iter). Claude-Session: https://claude.ai/code/session_01EvbTzx6ZJNtcAKwVQ1B68i
wgpu's validation error doesn't say which dispatch overflowed; panic at the khal layer with the grid so the kernel is identifiable from the backtrace. Claude-Session: https://claude.ai/code/session_01EvbTzx6ZJNtcAKwVQ1B68i
… nvidia-smi fails
sebcrozet
force-pushed
the
cuda-oxide-build-path
branch
from
September 30, 2026 17:01
720af9e to
78724f3
Compare
sebcrozet
changed the title
Merged
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This adds support for cuda-oxide. This is mainly based on #9 cleaned up and modified such that:
rust-cudaversion still works.Close #9