From 3be864d9f952e3cd11b40462ece1ca2b92004227 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 26 May 2026 09:45:57 +0200 Subject: [PATCH] feat(rl): raw CUDA launch wrapper + cudarc cu_function() accessor raw_launch.rs: RawArgs fixed-array u64 storage + raw_launch() calling cuLaunchKernel directly. Bypasses PushKernelArg trait dispatch, Arc bookkeeping, event guards. Foundation for 500+ sps hot path. vendor/cudarc: added CudaFunction::cu_function() accessor mirroring CudaStream::cu_stream() pattern. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/src/trainer/mod.rs | 1 + crates/ml-alpha/src/trainer/raw_launch.rs | 241 ++++++++++++++++++++++ vendor/cudarc/src/driver/safe/core.rs | 4 + 3 files changed, 246 insertions(+) create mode 100644 crates/ml-alpha/src/trainer/raw_launch.rs diff --git a/crates/ml-alpha/src/trainer/mod.rs b/crates/ml-alpha/src/trainer/mod.rs index 91dfa5fb3..d8b1ee485 100644 --- a/crates/ml-alpha/src/trainer/mod.rs +++ b/crates/ml-alpha/src/trainer/mod.rs @@ -6,3 +6,4 @@ pub mod integrated; pub mod loss; pub mod optim; pub mod perception; +pub mod raw_launch; diff --git a/crates/ml-alpha/src/trainer/raw_launch.rs b/crates/ml-alpha/src/trainer/raw_launch.rs new file mode 100644 index 000000000..8fb61f3dd --- /dev/null +++ b/crates/ml-alpha/src/trainer/raw_launch.rs @@ -0,0 +1,241 @@ +//! Zero-overhead CUDA kernel launch bypassing cudarc's `PushKernelArg` +//! trait dispatch and event tracking. Used in the step hot path where +//! cudarc's ~200us-per-launch overhead dominates at ~70 launches/step. +//! +//! cudarc stays for device init, module loading, and memory allocation. +//! Only the launch path is bypassed here. + +use cudarc::driver::sys::{self, CUfunction, CUresult, CUstream}; +use std::ffi::c_void; + +/// Maximum kernel arguments supported by a single launch. +/// 32 covers every kernel in the step pipeline (largest is ~24 args). +const MAX_ARGS: usize = 32; + +/// Fixed-capacity kernel argument buffer. +/// +/// All values are stored as `u64` in a flat array. For device pointers +/// (`CUdeviceptr = u64`) this is a direct store. For smaller scalars +/// (`i32`, `f32`), the value occupies the low bytes of the `u64` slot; +/// CUDA reads `sizeof(param)` bytes from the pointer we give it, which +/// on little-endian reads exactly the right bytes. +/// +/// The `arg_ptrs` array is **not** stored inline because the struct +/// would become self-referential and invalidate on move. Instead, call +/// [`RawArgs::build_arg_ptrs`] to produce a stack-local pointer array +/// right before launch. +pub struct RawArgs { + storage: [u64; MAX_ARGS], + len: usize, +} + +impl RawArgs { + #[inline(always)] + pub fn new() -> Self { + Self { + storage: [0u64; MAX_ARGS], + len: 0, + } + } + + /// Reset the argument list for reuse on the next kernel. + #[inline(always)] + pub fn reset(&mut self) { + self.len = 0; + } + + /// Number of arguments currently stored. + #[inline(always)] + pub fn len(&self) -> usize { + self.len + } + + /// Whether the argument list is empty. + #[inline(always)] + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// Push a device pointer argument (`CUdeviceptr = u64`). + #[inline(always)] + pub fn push_ptr(&mut self, ptr: u64) { + debug_assert!(self.len < MAX_ARGS, "RawArgs overflow"); + self.storage[self.len] = ptr; + self.len += 1; + } + + /// Push a raw `u64` argument (device pointer or 64-bit scalar). + #[inline(always)] + pub fn push_u64(&mut self, val: u64) { + debug_assert!(self.len < MAX_ARGS, "RawArgs overflow"); + self.storage[self.len] = val; + self.len += 1; + } + + /// Push an `i32` scalar argument. + /// + /// Stored in the low 4 bytes of a `u64` slot. CUDA reads + /// `sizeof(int)` = 4 bytes from the pointer, which on + /// little-endian is exactly the `i32` value. + #[inline(always)] + pub fn push_i32(&mut self, val: i32) { + debug_assert!(self.len < MAX_ARGS, "RawArgs overflow"); + // Zero-extend to u64 so the upper bytes are clean. + self.storage[self.len] = val as u32 as u64; + self.len += 1; + } + + /// Push an `f32` scalar argument. + /// + /// Stored as the bit pattern in the low 4 bytes of a `u64` slot. + #[inline(always)] + pub fn push_f32(&mut self, val: f32) { + debug_assert!(self.len < MAX_ARGS, "RawArgs overflow"); + self.storage[self.len] = val.to_bits() as u64; + self.len += 1; + } + + /// Build the `void**` argument pointer array on the caller's stack. + /// + /// Each entry points to the corresponding `storage[i]` slot. + /// The returned array is valid as long as `self` is not moved or + /// dropped -- call this immediately before [`raw_launch`] and pass + /// the result directly. + #[inline(always)] + pub fn build_arg_ptrs(&mut self) -> [*mut c_void; MAX_ARGS] { + let mut ptrs = [std::ptr::null_mut(); MAX_ARGS]; + for i in 0..self.len { + ptrs[i] = (&mut self.storage[i]) as *mut u64 as *mut c_void; + } + ptrs + } +} + +impl Default for RawArgs { + fn default() -> Self { + Self::new() + } +} + +/// Launch a CUDA kernel with pre-packed argument pointers. +/// +/// Zero trait dispatch, zero Arc increment, zero event tracking. +/// The caller is responsible for stream synchronization and argument +/// validity. +/// +/// # Safety +/// +/// - `func` must be a valid, loaded `CUfunction` handle. +/// - `stream` must be a valid `CUstream` on the same device/context. +/// - `args` must contain exactly as many valid pointers as the kernel +/// expects, each pointing to a value of the correct type and size. +/// - The caller must ensure all device memory referenced by the kernel +/// remains valid until the kernel completes (async execution). +#[inline(always)] +pub unsafe fn raw_launch( + func: CUfunction, + grid: (u32, u32, u32), + block: (u32, u32, u32), + shared_mem_bytes: u32, + stream: CUstream, + args: &mut [*mut c_void], +) -> Result<(), CUresult> { + let r = sys::cuLaunchKernel( + func, + grid.0, grid.1, grid.2, + block.0, block.1, block.2, + shared_mem_bytes, + stream, + args.as_mut_ptr(), + std::ptr::null_mut(), + ); + if r != CUresult::CUDA_SUCCESS { + return Err(r); + } + Ok(()) +} + +/// Extract the raw `CUfunction` handle from a cudarc `CudaFunction`. +#[inline(always)] +pub fn extract_fn(func: &cudarc::driver::CudaFunction) -> CUfunction { + func.cu_function() +} + +/// Extract the raw `CUstream` handle from a cudarc `CudaStream`. +#[inline(always)] +pub fn extract_stream(stream: &cudarc::driver::CudaStream) -> CUstream { + stream.cu_stream() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn raw_args_push_and_build() { + let mut args = RawArgs::new(); + assert!(args.is_empty()); + + args.push_ptr(0xDEAD_BEEF_u64); + args.push_i32(42); + args.push_f32(3.14); + args.push_u64(0xCAFE); + + assert_eq!(args.len(), 4); + + let ptrs = args.build_arg_ptrs(); + // Verify the first pointer points to storage[0] which holds 0xDEAD_BEEF + unsafe { + let val = *(ptrs[0] as *const u64); + assert_eq!(val, 0xDEAD_BEEF_u64); + } + // Verify i32 stored correctly in low bytes + unsafe { + let val = *(ptrs[1] as *const u32); + assert_eq!(val, 42); + } + // Verify f32 bit pattern + unsafe { + let val = *(ptrs[2] as *const u32); + assert_eq!(f32::from_bits(val), 3.14); + } + // Verify u64 + unsafe { + let val = *(ptrs[3] as *const u64); + assert_eq!(val, 0xCAFE); + } + // Remaining pointers are null + assert!(ptrs[4].is_null()); + } + + #[test] + fn raw_args_reset_reuse() { + let mut args = RawArgs::new(); + args.push_ptr(1); + args.push_ptr(2); + assert_eq!(args.len(), 2); + + args.reset(); + assert!(args.is_empty()); + assert_eq!(args.len(), 0); + + args.push_ptr(99); + assert_eq!(args.len(), 1); + let ptrs = args.build_arg_ptrs(); + unsafe { + assert_eq!(*(ptrs[0] as *const u64), 99); + } + } + + #[test] + fn push_i32_negative() { + let mut args = RawArgs::new(); + args.push_i32(-1); + let ptrs = args.build_arg_ptrs(); + // -1i32 as u32 = 0xFFFF_FFFF, CUDA reads 4 bytes and interprets as int + unsafe { + let val = *(ptrs[0] as *const i32); + assert_eq!(val, -1); + } + } +} diff --git a/vendor/cudarc/src/driver/safe/core.rs b/vendor/cudarc/src/driver/safe/core.rs index ddee1b94c..919947233 100644 --- a/vendor/cudarc/src/driver/safe/core.rs +++ b/vendor/cudarc/src/driver/safe/core.rs @@ -1976,6 +1976,10 @@ impl CudaModule { } impl CudaFunction { + pub fn cu_function(&self) -> sys::CUfunction { + self.cu_function + } + pub fn occupancy_available_dynamic_smem_per_block( &self, num_blocks: u32,