gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf(htod): migrate gpu_her.rs to mapped-pinned + add MappedU32Buffer (3 sites)

Browse Source 
- 1 constructor site (rng_states u32 init) rewritten to
  mapped_pinned::upload_u32_via_pinned.
- 2 warm-path sites (source_indices, donor_indices in relabel_batch)
  rewritten to mapped_pinned::upload_i32_via_pinned. The warm sites
  re-allocate pinned staging per relabel call — acceptable on this
  cold path (HER relabel runs once per epoch).

Adds MappedU32Buffer type to mapped_pinned.rs mirroring MappedI32Buffer,
plus upload_u32_via_pinned helper. Manual Debug impl so warning count
stays at 13.

Audit row appended (Fix 10) in docs/dqn-gpu-hot-path-audit.md.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-28 21:04:29 +02:00
parent f77ebeb214
commit efeb950751
3 changed files with 133 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
crates/ml/src/cuda_pipeline/gpu_her.rs
View File

@@ -197,8 +197,8 @@ impl GpuHer {
.collect();
let mut rng_states = stream.alloc_zeros::<u32>(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 {

113
crates/ml/src/cuda_pipeline/mapped_pinned.rs
View File

@@ -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<Self, String> {
let num_bytes = len * std::mem::size_of::<u32>();
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<u32>` with `host` via mapped-pinned staging.
///
/// # Safety
/// Caller must hold an active CUDA context.
pub fn upload_u32_via_pinned(
stream: &Arc<CudaStream>,
host: &[u32],
dst: &mut CudaSlice<u32>,
) -> 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::<u32>();
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(())
}

3
docs/dqn-gpu-hot-path-audit.md
View File

@@ -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.