diff --git a/crates/ml/src/cuda_pipeline/gpu_her.rs b/crates/ml/src/cuda_pipeline/gpu_her.rs index 21853f9b1..6877992db 100644 --- a/crates/ml/src/cuda_pipeline/gpu_her.rs +++ b/crates/ml/src/cuda_pipeline/gpu_her.rs @@ -197,8 +197,8 @@ impl GpuHer { .collect(); let mut rng_states = stream.alloc_zeros::(her_batch) .map_err(|e| MLError::ModelError(format!("GpuHer alloc rng_states: {e}")))?; - stream.memcpy_htod(&rng_init, &mut rng_states) - .map_err(|e| MLError::ModelError(format!("GpuHer rng_states upload: {e}")))?; + super::mapped_pinned::upload_u32_via_pinned(&stream, &rng_init, &mut rng_states) + .map_err(|e| MLError::ModelError(format!("GpuHer rng_states upload via pinned: {e}")))?; let vram_bytes = (her_batch * ml_core::state_layout::STATE_DIM * 2 + her_batch * 3) * 4 + her_batch * 3 * 4; // output bufs + index bufs + rng_states @@ -273,13 +273,21 @@ impl GpuHer { ))); } - // Upload source and donor indices to GPU - self.stream - .memcpy_htod(&source_idx_host[..her_batch_size], &mut self.source_indices) - .map_err(|e| MLError::ModelError(format!("HER source_indices upload: {e}")))?; - self.stream - .memcpy_htod(&donor_idx_host[..her_batch_size], &mut self.donor_indices) - .map_err(|e| MLError::ModelError(format!("HER donor_indices upload: {e}")))?; + // Upload source and donor indices to GPU via mapped-pinned staging + // (per-relabel call; pinned staging is allocated+freed each call — + // acceptable on this cold path which runs once per epoch). + super::mapped_pinned::upload_i32_via_pinned( + &self.stream, + &source_idx_host[..her_batch_size], + &mut self.source_indices, + ) + .map_err(|e| MLError::ModelError(format!("HER source_indices upload via pinned: {e}")))?; + super::mapped_pinned::upload_i32_via_pinned( + &self.stream, + &donor_idx_host[..her_batch_size], + &mut self.donor_indices, + ) + .map_err(|e| MLError::ModelError(format!("HER donor_indices upload via pinned: {e}")))?; // Launch relabel kernel: one warp (32 threads) per sample let launch_config = LaunchConfig { diff --git a/crates/ml/src/cuda_pipeline/mapped_pinned.rs b/crates/ml/src/cuda_pipeline/mapped_pinned.rs index 1484478b1..11bd4f463 100644 --- a/crates/ml/src/cuda_pipeline/mapped_pinned.rs +++ b/crates/ml/src/cuda_pipeline/mapped_pinned.rs @@ -107,6 +107,83 @@ impl Drop for MappedI32Buffer { } } +// ── MappedU32Buffer ─────────────────────────────────────────────────────────── + +/// CPU+GPU visible buffer of `u32`s allocated via +/// `cuMemHostAlloc(DEVICEMAP|PORTABLE)`. Kernels read/write via `dev_ptr`; +/// the CPU reads via `read_volatile` on `host_ptr`. +pub struct MappedU32Buffer { + pub host_ptr: *mut u32, + pub dev_ptr: cudarc::driver::sys::CUdeviceptr, + pub len: usize, +} + +impl std::fmt::Debug for MappedU32Buffer { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("MappedU32Buffer") + .field("dev_ptr", &self.dev_ptr) + .field("len", &self.len) + .finish() + } +} + +unsafe impl Send for MappedU32Buffer {} +unsafe impl Sync for MappedU32Buffer {} + +impl MappedU32Buffer { + /// Allocate `len` u32s of mapped pinned memory. + /// + /// # Safety + /// Caller must ensure a CUDA context is active on the current thread. + pub unsafe fn new(len: usize) -> Result { + let num_bytes = len * std::mem::size_of::(); + let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP + | cudarc::driver::sys::CU_MEMHOSTALLOC_PORTABLE; + + let host_ptr = cudarc::driver::result::malloc_host(num_bytes, flags) + .map_err(|e| format!("MappedU32Buffer alloc ({len} u32): {e}"))? + as *mut u32; + + std::ptr::write_bytes(host_ptr, 0, len); + + let mut dev_ptr_raw = MaybeUninit::uninit(); + cudarc::driver::sys::cuMemHostGetDevicePointer_v2( + dev_ptr_raw.as_mut_ptr(), + host_ptr as *mut c_void, + 0, + ) + .result() + .map_err(|e| format!("cuMemHostGetDevicePointer (u32 buf): {e}"))?; + + Ok(Self { + host_ptr, + dev_ptr: dev_ptr_raw.assume_init(), + len, + }) + } + + /// Write CPU-side data into the buffer via host_ptr (no memcpy). + /// Caller must ensure `slice.len() <= self.len`. + pub fn write_from_slice(&self, slice: &[u32]) { + assert!(slice.len() <= self.len, "MappedU32Buffer write overflow"); + // Safety: host_ptr is valid for self.len u32s by construction. + unsafe { + for (i, &v) in slice.iter().enumerate() { + std::ptr::write_volatile(self.host_ptr.add(i), v); + } + } + } +} + +impl Drop for MappedU32Buffer { + fn drop(&mut self) { + unsafe { + #[allow(clippy::let_underscore_must_use)] + let _ = cudarc::driver::result::free_host(self.host_ptr as *mut c_void); + } + } +} + // ── MappedF32Buffer ─────────────────────────────────────────────────────────── /// CPU+GPU visible buffer of `f32`s allocated via @@ -359,3 +436,39 @@ pub fn upload_i32_via_pinned( .map_err(|e| format!("upload_i32_via_pinned sync: {e}"))?; Ok(()) } + +/// Seed an existing `CudaSlice` with `host` via mapped-pinned staging. +/// +/// # Safety +/// Caller must hold an active CUDA context. +pub fn upload_u32_via_pinned( + stream: &Arc, + host: &[u32], + dst: &mut CudaSlice, +) -> Result<(), String> { + let n = host.len(); + if n == 0 { + return Ok(()); + } + let staging = unsafe { MappedU32Buffer::new(n) }?; + staging.write_from_slice(host); + + let nbytes = n * std::mem::size_of::(); + unsafe { + let (dst_ptr, _g) = { + use cudarc::driver::DevicePtrMut; + dst.device_ptr_mut(stream) + }; + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, + staging.dev_ptr, + nbytes, + stream.cu_stream(), + ) + .map_err(|e| format!("upload_u32_via_pinned DtoD: {e}"))?; + } + stream + .synchronize() + .map_err(|e| format!("upload_u32_via_pinned sync: {e}"))?; + Ok(()) +} diff --git a/docs/dqn-gpu-hot-path-audit.md b/docs/dqn-gpu-hot-path-audit.md index 1d7f7d770..fd16edf29 100644 --- a/docs/dqn-gpu-hot-path-audit.md +++ b/docs/dqn-gpu-hot-path-audit.md @@ -112,3 +112,6 @@ ### Fix 9 — Migrate `gpu_iqn_head.rs` constructors + DtoD-only target clone (5 sites) `gpu_iqn_head.rs`: 4 `stream.clone_htod` calls (cos_features precompute, online_taus + target_taus tile broadcast, init_iqn_xavier_weights upload) rewritten to `mapped_pinned::clone_to_device_f32_via_pinned`. Additionally the helper `clone_cuda_slice` was replaced — it previously did device→host (`clone_dtoh`) → host→device (`clone_htod`) round-trip, double-violating both the no-DtoH and no-HtoD rules. Now uses `clone_cuda_slice_f32` (single async DtoD), the only callsite being the target-network seed at constructor time (`target_params = clone_cuda_slice(&online_params)`). + +### Fix 10 — Migrate `gpu_her.rs` to mapped-pinned + add MappedU32Buffer (3 sites) +`gpu_her.rs`: 1 cold-path constructor `stream.memcpy_htod` (rng_states u32 init) and 2 warm-path `memcpy_htod` (source_indices, donor_indices in `relabel_batch`) rewritten to `mapped_pinned::upload_{u32,i32}_via_pinned`. The warm sites re-allocate pinned staging per relabel call — acceptable on this cold path (relabel runs once per epoch). New `MappedU32Buffer` type added to `mapped_pinned.rs` mirroring `MappedI32Buffer`; carries a manual `Debug` impl so the workspace warning count stays at 13.