[Draft/RFC] Native-CUDA (cuda-oxide) device backend + CUDA-graph capture#9
Draft
haixuanTao wants to merge 10 commits into
Draft
[Draft/RFC] Native-CUDA (cuda-oxide) device backend + CUDA-graph capture#9haixuanTao wants to merge 10 commits into
haixuanTao wants to merge 10 commits into
Conversation
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. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…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). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…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). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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 #[kernel] proc-macro, so pure shader crates need no cuda-host dependency; #[unsafe(no_mangle)] forces local codegen (rustc's cross-crate-inlinable 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. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…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. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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
1 participant
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.
Draft for discussion — the khal half of the native-CUDA (cuda-oxide) stack that now runs nexus end-to-end from Python (see the companion dimforge/nexus and dimforge/vortx drafts).
What's in here
#[spirv_bindgen]grows a cuda-oxide entry-point flavor (typed params,SmemBufshared memory,llvm.nvvm.*intrinsics lowered by cuda-oxide) selected by the shader crate'scuda-oxidefeature.khal-stdgains the panic-free device shims (trigFloatmethods, genericSmemBuf, atomics).compute_capability, graph capture (begin_capture/end_capture→CapturedGraph::{upload,launch}),KHAL_CUDA_PROFILE, and theCUDA_OXIDE_SHADERS_PTX_<CRATE>prebuilt-cubin escape hatch in khal-builder (skipscargo cuda buildas a build prerequisite).default_stream()— the legacy NULL stream, which CUDA refuses to capture. It now owns a non-blocking stream. AlsoGpuBackend(the enum) never delegatedBackend::as_cuda, so the capture API was unreachable through it. PlusKHAL_TRACE_INDIRECT=1for hunting indirect dispatches (each one costs a full stream drain on CUDA and invalidates a capture).Measured (nexus cube-drop capture demo, RTX 5080 laptop)
WebGPU 92 fps → CUDA with indirect dispatch 5.6 fps → fixed-grid dispatch 16.4 fps → CUDA graph 117 fps. The 13-link LeRobot multibody: 6.9 → 28.3 fps.
Blocked on upstream cuda-oxide
Building device code with this backend needs the still-open NVlabs/cuda-oxide PRs: #350 (unsize address-space coercion), #358 (collector kernel rooting), #360 (canonical
sync_threadsdispatch) — until those land, the backend.socomes from thehaixuanTao/cuda-oxidefeat/nexus3d-vortx-native-cudabranch.Note: includes the commit already open standalone as #6 (slice arg element count).
🤖 Generated with Claude Code