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

cleanup(pinmem): rename dtod_copy → dtod_copy_async — [PINMEM-007,PINMEM-013]

Browse Source 
Both dtod_copy wrappers (noisy_layers.rs:372, gpu_dqn_trainer.rs:13220)
already call memcpy_dtod_async internally — the name didn't advertise
the async semantics, which caused the PINMEM scan to false-positive on
them.

Per user question "why use wrappers?" — they provide per-call-site
error context (label / op / idx) with ~5 LOC of setup each. The value
is real but small; renaming to dtod_copy_async makes the async
semantics visible at every call site and lets the scan regex
(`\bdtod_copy\b`) stop matching them.

27 sites across 4 files: noisy_layers.rs (7), gpu_dqn_trainer.rs
(def + many callers), gpu_iqn_head.rs (3), fused_training.rs (imports).
This commit is contained in:
jgrusewski
2026-04-21 00:07:57 +02:00
parent b0fdf9b3e4
commit 3e8003b7f2
4 changed files with 27 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
crates/ml-dqn/src/noisy_layers.rs
View File

@@ -359,17 +359,17 @@ impl NoisyLinear {
/// Copies `weight_mu`, `bias_mu`, `weight_sigma`, `bias_sigma`, and epsilon buffers.
/// Zero CPU involvement — all copies are async device-to-device.
pub fn copy_params_from(&mut self, src: &NoisyLinear) -> Result<(), MLError> {
Self::dtod_copy(&src.weight_mu, &mut self.weight_mu, "weight_mu", &self.stream)?;
Self::dtod_copy(&src.bias_mu, &mut self.bias_mu, "bias_mu", &self.stream)?;
Self::dtod_copy(&src.weight_sigma, &mut self.weight_sigma, "weight_sigma", &self.stream)?;
Self::dtod_copy(&src.bias_sigma, &mut self.bias_sigma, "bias_sigma", &self.stream)?;
Self::dtod_copy(&src.weight_epsilon, &mut self.weight_epsilon, "weight_epsilon", &self.stream)?;
Self::dtod_copy(&src.bias_epsilon, &mut self.bias_epsilon, "bias_epsilon", &self.stream)?;
Self::dtod_copy_async(&src.weight_mu, &mut self.weight_mu, "weight_mu", &self.stream)?;
Self::dtod_copy_async(&src.bias_mu, &mut self.bias_mu, "bias_mu", &self.stream)?;
Self::dtod_copy_async(&src.weight_sigma, &mut self.weight_sigma, "weight_sigma", &self.stream)?;
Self::dtod_copy_async(&src.bias_sigma, &mut self.bias_sigma, "bias_sigma", &self.stream)?;
Self::dtod_copy_async(&src.weight_epsilon, &mut self.weight_epsilon, "weight_epsilon", &self.stream)?;
Self::dtod_copy_async(&src.bias_epsilon, &mut self.bias_epsilon, "bias_epsilon", &self.stream)?;
Ok(())
}
/// Async device-to-device memcpy for a single `CudaSlice` buffer.
fn dtod_copy(
fn dtod_copy_async(
src: &CudaSlice<f32>,
dst: &mut CudaSlice<f32>,
label: &str,

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

@@ -8397,7 +8397,7 @@ impl GpuDqnTrainer {
// td_errors → readback[3..3+B] (only td_errors needs DtoD → DtoH)
let td_src = self.td_errors_buf.raw_ptr();
let td_bytes = b * f32_bytes;
dtod_copy(readback_base + (3 * f32_bytes) as u64, td_src, td_bytes, &self.stream, 3, "readback_gather")?;
dtod_copy_async(readback_base + (3 * f32_bytes) as u64, td_src, td_bytes, &self.stream, 3, "readback_gather")?;
// ── Single DtoH transfer (only td_errors portion matters) ─────
super::dtoh_f32(&self.stream, &self.readback_buf, &mut self.readback_host)?;
@@ -10851,26 +10851,26 @@ impl GpuDqnTrainer {
)?;
// Rewards, dones: f32 DtoD
dtod_copy(
dtod_copy_async(
self.rewards_buf.raw_ptr(),
gpu_batch.rewards_ptr,
b * f32_size, &self.stream, 2, "rewards",
)?;
dtod_copy(
dtod_copy_async(
self.dones_buf.raw_ptr(),
gpu_batch.dones_ptr,
b * f32_size, &self.stream, 3, "dones",
)?;
// IS-weights: f32 DtoD
dtod_copy(
dtod_copy_async(
self.is_weights_buf.raw_ptr(),
gpu_batch.weights_ptr,
b * f32_size, &self.stream, 4, "is_weights",
)?;
// Actions: CudaSlice<i32> → i32 buf (async DtoD, zero CPU sync)
dtod_copy(
dtod_copy_async(
self.actions_buf.raw_ptr(), gpu_batch.actions_ptr,
b * std::mem::size_of::<i32>(), &self.stream, 4, "actions",
)?;
@@ -12038,7 +12038,7 @@ impl GpuDqnTrainer {
)));
}
let num_bytes = sizes[i] * std::mem::size_of::<f32>();
dtod_copy(dst_base + byte_offset, src, num_bytes, &self.stream, i, "flatten")?;
dtod_copy_async(dst_base + byte_offset, src, num_bytes, &self.stream, i, "flatten")?;
byte_offset += num_bytes as u64;
}
@@ -12085,7 +12085,7 @@ impl GpuDqnTrainer {
let mut byte_offset: u64 = 0;
for (i, &dst) in ptrs.iter().enumerate() {
let num_bytes = sizes[i] * std::mem::size_of::<f32>();
dtod_copy(dst, src_base + byte_offset, num_bytes, &self.stream, i, "unflatten")?;
dtod_copy_async(dst, src_base + byte_offset, num_bytes, &self.stream, i, "unflatten")?;
byte_offset += num_bytes as u64;
}
@@ -12371,7 +12371,7 @@ impl GpuDqnTrainer {
for (i, &src) in ptrs.iter().enumerate() {
let num_bytes = sizes[i] * std::mem::size_of::<f32>();
let dst = dst_base + byte_offset;
dtod_copy(dst, src, num_bytes, &self.stream, i, "flatten_target")?;
dtod_copy_async(dst, src, num_bytes, &self.stream, i, "flatten_target")?;
byte_offset += num_bytes as u64;
}
@@ -12412,7 +12412,7 @@ impl GpuDqnTrainer {
let mut byte_offset: u64 = 0;
for (i, &dst) in ptrs.iter().enumerate() {
let num_bytes = sizes[i] * std::mem::size_of::<f32>();
dtod_copy(dst, src_base + byte_offset, num_bytes, &self.stream, i, "unflatten_target")?;
dtod_copy_async(dst, src_base + byte_offset, num_bytes, &self.stream, i, "unflatten_target")?;
byte_offset += num_bytes as u64;
}
@@ -13217,7 +13217,7 @@ pub(crate) fn query_max_shmem_bytes() -> usize {
// raw_device_ptr wrappers removed — use CudaSlice::raw_ptr() directly (no event tracking).
/// Async device-to-device memcpy with error context.
pub(crate) fn dtod_copy(
pub(crate) fn dtod_copy_async(
dst: u64,
src: u64,
num_bytes: usize,
@@ -13247,7 +13247,7 @@ fn dtod_from_slice_f32(
) -> Result<(), MLError> {
let src_ptr = src.raw_ptr();
let dst_ptr = dst.raw_ptr();
dtod_copy(dst_ptr, src_ptr, num_elements * std::mem::size_of::<f32>(), stream, 0, ctx)
dtod_copy_async(dst_ptr, src_ptr, num_elements * std::mem::size_of::<f32>(), stream, 0, ctx)
}
/// DtoD copy from a CudaSlice<u32> to a CudaSlice<i32>.
@@ -13262,7 +13262,7 @@ fn dtod_from_u32_to_i32(
) -> Result<(), MLError> {
let src_ptr = src.raw_ptr();
let dst_ptr = dst.raw_ptr();
dtod_copy(dst_ptr, src_ptr, num_elements * std::mem::size_of::<u32>(), stream, 0, ctx)
dtod_copy_async(dst_ptr, src_ptr, num_elements * std::mem::size_of::<u32>(), stream, 0, ctx)
}
// ── Allocation helpers ──────────────────────────────────────────────────────
@@ -13331,7 +13331,7 @@ fn dtod_from_staging(
) -> Result<(), MLError> {
let src_ptr = src.raw_ptr();
let dst_ptr = dst.raw_ptr();
dtod_copy(dst_ptr, src_ptr, num_elements * std::mem::size_of::<f32>(), stream, 0, ctx)
dtod_copy_async(dst_ptr, src_ptr, num_elements * std::mem::size_of::<f32>(), stream, 0, ctx)
}
/// Load ensemble aggregate + diversity kernels from precompiled cubin.

6
crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
View File

@@ -697,15 +697,15 @@ impl GpuIqnHead {
.map_err(|e| MLError::ModelError(format!("IQN decode_actions kernel: {e}")))?;
}
use super::gpu_dqn_trainer::dtod_copy;
use super::gpu_dqn_trainer::dtod_copy_async;
let f32_bytes = std::mem::size_of::<f32>();
let r_src = dqn_rewards_buf.raw_ptr();
let r_dst = self.rewards_buf.raw_ptr();
dtod_copy(r_dst, r_src, b * f32_bytes, &self.stream, 0, "iqn_rewards_dtod")?;
dtod_copy_async(r_dst, r_src, b * f32_bytes, &self.stream, 0, "iqn_rewards_dtod")?;
let d_src = dqn_dones_buf.raw_ptr();
let d_dst = self.dones_buf.raw_ptr();
dtod_copy(d_dst, d_src, b * f32_bytes, &self.stream, 0, "iqn_dones_dtod")?;
dtod_copy_async(d_dst, d_src, b * f32_bytes, &self.stream, 0, "iqn_dones_dtod")?;
Ok(())
}

8
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -2991,13 +2991,13 @@ fn clone_dueling_weights(
stream: &Arc<cudarc::driver::CudaStream>,
head_idx: usize,
) -> Result<(gpu_weights::DuelingWeightBacking, DuelingWeightSet)> {
use crate::cuda_pipeline::gpu_dqn_trainer::dtod_copy;
use crate::cuda_pipeline::gpu_dqn_trainer::dtod_copy_async;
let clone_from_ptr = |ptr: u64, len: usize| -> Result<cudarc::driver::CudaSlice<f32>> {
let dst = stream.alloc_zeros::<f32>(len)
.map_err(|e| anyhow::anyhow!("Clone alloc {len}xf32: {e}"))?;
let n_bytes = len * std::mem::size_of::<f32>();
dtod_copy(dst.raw_ptr(), ptr, n_bytes, stream, 0, "clone_dueling")
dtod_copy_async(dst.raw_ptr(), ptr, n_bytes, stream, 0, "clone_dueling")
.map_err(|e| anyhow::anyhow!("DtoD clone: {e}"))?;
Ok(dst)
};
@@ -3032,13 +3032,13 @@ fn clone_branching_weights(
stream: &Arc<cudarc::driver::CudaStream>,
head_idx: usize,
) -> Result<(gpu_weights::BranchingWeightBacking, BranchingWeightSet)> {
use crate::cuda_pipeline::gpu_dqn_trainer::dtod_copy;
use crate::cuda_pipeline::gpu_dqn_trainer::dtod_copy_async;
let clone_from_ptr = |ptr: u64, len: usize| -> Result<cudarc::driver::CudaSlice<f32>> {
let dst = stream.alloc_zeros::<f32>(len)
.map_err(|e| anyhow::anyhow!("Clone alloc {len}xf32: {e}"))?;
let n_bytes = len * std::mem::size_of::<f32>();
dtod_copy(dst.raw_ptr(), ptr, n_bytes, stream, 0, "clone_branching")
dtod_copy_async(dst.raw_ptr(), ptr, n_bytes, stream, 0, "clone_branching")
.map_err(|e| anyhow::anyhow!("DtoD clone: {e}"))?;
Ok(dst)
};