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

cleanup(borrow): fix CudaSlice<i32>→u32 aliasing UB in episode_ids — [BORROW-001]

Browse Source 
gpu_replay_buffer.rs had three sites that cast `&CudaSlice<i32>` to
`&CudaSlice<u32>` via raw-pointer transmutes to satisfy kernel
signatures (gather_u32, scatter_insert_u32). This changes the element
type of a `&mut` reference through `as *mut`, which is UB under Rust's
strict aliasing model even though i32 and u32 share a memory layout.

Episode-id values are always non-negative (counters of the form
`(write_cursor + j) % capacity`), so u32 is the correct storage type.
Changed:
- 3 field types: episode_ids, sample_episode_ids, insert_ep_buf
  (CudaSlice<i32> → CudaSlice<u32>)
- Allocators: a32i → a32u at 4 sites
- Vec<i32> → Vec<u32> at ep_ids_host, with `as u32` cast
- Deleted 3 raw-ptr transmute blocks (lines 491-492, 689-690)
- Public sample_episode_ids_ref() return type i32 → u32
- Deleted now-unused a32i helper (fn never called after the change)

Per BORROW learned pattern option (3): "change the declared type".
This commit is contained in:
jgrusewski
2026-04-21 00:12:08 +02:00
parent 3e8003b7f2
commit 8a31d3b99b
1 changed files with 19 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
crates/ml-dqn/src/gpu_replay_buffer.rs
View File

@@ -141,9 +141,11 @@ pub struct GpuReplayBuffer {
states: CudaSlice<f32>, next_states: CudaSlice<f32>,
actions: CudaSlice<u32>, rewards: CudaSlice<f32>,
dones: CudaSlice<f32>, priorities: CudaSlice<f32>,
/// Episode IDs per buffer slot `[capacity]` i32 on GPU.
/// Episode IDs per buffer slot `[capacity]` u32 on GPU.
/// `episode_ids[i] = i / episode_length`. Written during `insert_batch_with_episode_ids`.
episode_ids: CudaSlice<i32>,
/// u32 chosen to match the gather_u32 / scatter_insert_u32 kernel signatures —
/// episode-id values are always non-negative counters.
episode_ids: CudaSlice<u32>,
write_cursor: usize, size: usize,
max_priority: CudaSlice<f32>,
pending_max_priority: Option<CudaSlice<f32>>,
@@ -166,7 +168,7 @@ pub struct GpuReplayBuffer {
sample_priorities: CudaSlice<f32>,
sample_weights: CudaSlice<f32>,
sample_max_weight: CudaSlice<f32>,
sample_episode_ids: CudaSlice<i32>,
sample_episode_ids: CudaSlice<u32>,
total_sum_buf: CudaSlice<f32>,
// Pre-allocated scratch for update_priorities_gpu (zero cuMemAlloc per step)
update_batch_max: CudaSlice<f32>, // [1] per-batch atomicMax accumulator
@@ -175,7 +177,7 @@ pub struct GpuReplayBuffer {
// Pre-allocated insert scratch buffers (zero cuMemAlloc per insert)
insert_prio_buf: CudaSlice<f32>, // [insert_scratch_cap] priority broadcast
insert_idx_buf: CudaSlice<u32>, // [insert_scratch_cap] per_insert_pa indices
insert_ep_buf: CudaSlice<i32>, // [insert_scratch_cap] episode IDs
insert_ep_buf: CudaSlice<u32>, // [insert_scratch_cap] episode IDs (non-negative counters)
insert_scratch_cap: usize, // current allocation size
scratch_f32: CudaSlice<f32>, // [1] reusable temp for flush/apply_max
_rng_step: u32,
@@ -262,8 +264,8 @@ impl GpuReplayBuffer {
let sw = a32f(stream, mbs, "s_wt")?;
let smw = a32f(stream, 1, "s_mw")?;
let tsb = a32f(stream, 1, "ts_buf")?;
let ep_ids = a32i(stream, cap, "episode_ids")?;
let s_ep = a32i(stream, mbs, "s_episode_ids")?;
let ep_ids = a32u(stream, cap, "episode_ids")?;
let s_ep = a32u(stream, mbs, "s_episode_ids")?;
// Pre-allocate PER update scratch buffers (zero cuMemAlloc per step)
let update_batch_max = a32f(stream, 1, "update_batch_max")?;
@@ -275,7 +277,7 @@ impl GpuReplayBuffer {
let isc = mbs;
let insert_prio_buf = a32f(stream, isc, "ins_prio")?;
let insert_idx_buf = a32u(stream, isc, "ins_idx")?;
let insert_ep_buf = a32i(stream, isc, "ins_ep")?;
let insert_ep_buf = a32u(stream, isc, "ins_ep")?;
let scratch_f32 = a32f(stream, 1, "scratch")?;
// Pinned device-mapped buffer size for graph-safe PER — GPU reads current size at replay
@@ -393,7 +395,7 @@ impl GpuReplayBuffer {
if eff <= self.insert_scratch_cap { return Ok(()); }
self.insert_prio_buf = a32f(&self.stream, eff, "ins_prio")?;
self.insert_idx_buf = a32u(&self.stream, eff, "ins_idx")?;
self.insert_ep_buf = a32i(&self.stream, eff, "ins_ep")?;
self.insert_ep_buf = a32u(&self.stream, eff, "ins_ep")?;
self.insert_scratch_cap = eff;
Ok(())
}
@@ -481,17 +483,15 @@ impl GpuReplayBuffer {
.launch(lcfg(eff))
.map_err(|e| MLError::ModelError(format!("per_insert_pa: {e}")))?;
}
// Episode IDs using pre-allocated buffer
let ep_ids_host: Vec<i32> = (0..eff)
.map(|j| ((self.write_cursor + j) % cap) as i32)
// Episode IDs using pre-allocated buffer (u32 — always non-negative counters).
let ep_ids_host: Vec<u32> = (0..eff)
.map(|j| ((self.write_cursor + j) % cap) as u32)
.collect();
self.stream.memcpy_htod(&ep_ids_host, &mut self.insert_ep_buf)
.map_err(|e| MLError::ModelError(format!("ep htod: {e}")))?;
unsafe {
let ep_dst = &*(&self.episode_ids as *const CudaSlice<i32> as *const CudaSlice<u32>);
let ep_src = &*(&self.insert_ep_buf as *const CudaSlice<i32> as *const CudaSlice<u32>);
self.stream.launch_builder(&self.kernels.scatter_insert_u32)
.arg(ep_dst).arg(ep_src).arg(&ci).arg(&cpi).arg(&bsi)
.arg(&self.episode_ids).arg(&self.insert_ep_buf).arg(&ci).arg(&cpi).arg(&bsi)
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc ep: {e}")))?;
}
@@ -681,16 +681,13 @@ impl GpuReplayBuffer {
}
}
// Step 3b: gather episode_ids for HER strategies.
// Reinterpret i32 buffers as u32 (same bit width) for the gather_u32 kernel.
// SAFETY: episode_ids and sample_episode_ids are CudaSlice<i32>, same layout as u32.
// The gather kernel copies raw bytes, so reinterpret is safe for same-size types.
// Step 3b: gather episode_ids for HER strategies. Buffers are CudaSlice<u32>
// natively — no element-type aliasing needed.
unsafe {
let ep_src = &*(&self.episode_ids as *const CudaSlice<i32> as *const CudaSlice<u32>);
let ep_dst = &mut *(&mut self.sample_episode_ids as *mut CudaSlice<i32> as *mut CudaSlice<u32>);
let cap_ep = self.capacity() as i32;
self.stream.launch_builder(&self.kernels.gather_u32)
.arg(ep_dst).arg(ep_src).arg(&self.sample_indices_i64).arg(&bsi).arg(&cap_ep)
.arg(&mut self.sample_episode_ids).arg(&self.episode_ids)
.arg(&self.sample_indices_i64).arg(&bsi).arg(&cap_ep)
.launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g ep: {e}")))?;
}
@@ -791,7 +788,7 @@ impl GpuReplayBuffer {
}
/// Reference to pre-allocated sample episode IDs buffer (valid after sample_proportional).
pub fn sample_episode_ids_ref(&self) -> &CudaSlice<i32> { &self.sample_episode_ids }
pub fn sample_episode_ids_ref(&self) -> &CudaSlice<u32> { &self.sample_episode_ids }
/// Reference to pre-allocated sample weights buffer (valid after sample_proportional).
///
@@ -1017,9 +1014,6 @@ fn a32f(s: &Arc<CudaStream>, n: usize, nm: &str) -> Result<CudaSlice<f32>, MLErr
fn a32u(s: &Arc<CudaStream>, n: usize, nm: &str) -> Result<CudaSlice<u32>, MLError> {
s.alloc_zeros::<u32>(n).map_err(|e| MLError::ModelError(format!("alloc {nm}: {e}")))
}
fn a32i(s: &Arc<CudaStream>, n: usize, nm: &str) -> Result<CudaSlice<i32>, MLError> {
s.alloc_zeros::<i32>(n).map_err(|e| MLError::ModelError(format!("alloc {nm}: {e}")))
}
#[cfg(test)]
mod tests {