From 8a31d3b99b060382e2418960b9757cd67f14821f Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 21 Apr 2026 00:12:08 +0200 Subject: [PATCH] =?UTF-8?q?cleanup(borrow):=20fix=20CudaSlice?= =?UTF-8?q?=E2=86=92u32=20aliasing=20UB=20in=20episode=5Fids=20=E2=80=94?= =?UTF-8?q?=20[BORROW-001]?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit gpu_replay_buffer.rs had three sites that cast `&CudaSlice` to `&CudaSlice` 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 → CudaSlice) - Allocators: a32i → a32u at 4 sites - Vec → Vec 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". --- crates/ml-dqn/src/gpu_replay_buffer.rs | 44 +++++++++++--------------- 1 file changed, 19 insertions(+), 25 deletions(-) diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index d237dd7ce..e7736f436 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -141,9 +141,11 @@ pub struct GpuReplayBuffer { states: CudaSlice, next_states: CudaSlice, actions: CudaSlice, rewards: CudaSlice, dones: CudaSlice, priorities: CudaSlice, - /// 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, + /// u32 chosen to match the gather_u32 / scatter_insert_u32 kernel signatures — + /// episode-id values are always non-negative counters. + episode_ids: CudaSlice, write_cursor: usize, size: usize, max_priority: CudaSlice, pending_max_priority: Option>, @@ -166,7 +168,7 @@ pub struct GpuReplayBuffer { sample_priorities: CudaSlice, sample_weights: CudaSlice, sample_max_weight: CudaSlice, - sample_episode_ids: CudaSlice, + sample_episode_ids: CudaSlice, total_sum_buf: CudaSlice, // Pre-allocated scratch for update_priorities_gpu (zero cuMemAlloc per step) update_batch_max: CudaSlice, // [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, // [insert_scratch_cap] priority broadcast insert_idx_buf: CudaSlice, // [insert_scratch_cap] per_insert_pa indices - insert_ep_buf: CudaSlice, // [insert_scratch_cap] episode IDs + insert_ep_buf: CudaSlice, // [insert_scratch_cap] episode IDs (non-negative counters) insert_scratch_cap: usize, // current allocation size scratch_f32: CudaSlice, // [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 = (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 = (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 as *const CudaSlice); - let ep_src = &*(&self.insert_ep_buf as *const CudaSlice as *const CudaSlice); 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, 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 + // natively — no element-type aliasing needed. unsafe { - let ep_src = &*(&self.episode_ids as *const CudaSlice as *const CudaSlice); - let ep_dst = &mut *(&mut self.sample_episode_ids as *mut CudaSlice as *mut CudaSlice); 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 { &self.sample_episode_ids } + pub fn sample_episode_ids_ref(&self) -> &CudaSlice { &self.sample_episode_ids } /// Reference to pre-allocated sample weights buffer (valid after sample_proportional). /// @@ -1017,9 +1014,6 @@ fn a32f(s: &Arc, n: usize, nm: &str) -> Result, MLErr fn a32u(s: &Arc, n: usize, nm: &str) -> Result, MLError> { s.alloc_zeros::(n).map_err(|e| MLError::ModelError(format!("alloc {nm}: {e}"))) } -fn a32i(s: &Arc, n: usize, nm: &str) -> Result, MLError> { - s.alloc_zeros::(n).map_err(|e| MLError::ModelError(format!("alloc {nm}: {e}"))) -} #[cfg(test)] mod tests {