From b1eaef5db2c75c79a8e9af59c6479470f2c26e13 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 21 May 2026 01:11:22 +0200 Subject: [PATCH] refactor(gpu-log): generic MappedRecordBuffer; MappedF32Buffer as alias --- crates/ml-alpha/src/pinned_mem.rs | 107 +++++++++++++++++++++--------- 1 file changed, 76 insertions(+), 31 deletions(-) diff --git a/crates/ml-alpha/src/pinned_mem.rs b/crates/ml-alpha/src/pinned_mem.rs index c91c70719..0a63fe640 100644 --- a/crates/ml-alpha/src/pinned_mem.rs +++ b/crates/ml-alpha/src/pinned_mem.rs @@ -1,4 +1,4 @@ -//! Mapped-pinned f32 buffer (CPU+GPU visible via `cuMemHostAlloc(DEVICEMAP)`). +//! Mapped-pinned record buffers (CPU+GPU visible via `cuMemHostAlloc(DEVICEMAP)`). //! //! This is the **only** permitted CPU↔GPU path per //! `feedback_no_htod_htoh_only_mapped_pinned.md`. Mirror of @@ -12,37 +12,47 @@ use std::ffi::c_void; use std::mem::MaybeUninit; -/// CPU+GPU visible buffer of `f32`s allocated via -/// `cuMemHostAlloc(DEVICEMAP|PORTABLE)`. The kernel writes via `dev_ptr` +/// Mapped-pinned host buffer of `T` records. Allocated via +/// `cuMemHostAlloc(DEVICEMAP|PORTABLE)`; both host and device see the +/// same physical memory. Device pointer obtained via +/// `cuMemHostGetDevicePointer_v2`. The kernel writes via `dev_ptr` /// (with `__threadfence_system()` in the kernel); the CPU reads via /// `host_ptr` after stream synchronisation. +/// +/// Per `feedback_no_htod_htoh_only_mapped_pinned`: the ONLY legitimate +/// channel for cross-direction transfer in the captured-graph hot path. +/// +/// SAFETY: callers must ensure `T` is plain-old-data (no Drop, no +/// references, `repr(C)` recommended). The `Copy` bound enforces a +/// subset of this at compile time. #[allow(missing_debug_implementations)] -pub struct MappedF32Buffer { - pub host_ptr: *mut f32, +pub struct MappedRecordBuffer { + pub host_ptr: *mut T, pub dev_ptr: cudarc::driver::sys::CUdeviceptr, + /// Number of `T`-records (NOT bytes). pub len: usize, } // Safety: only accessed from the thread that constructed it (which holds // the active CUDA context). The host pointer is valid for the life of the // allocation. -unsafe impl Send for MappedF32Buffer {} -unsafe impl Sync for MappedF32Buffer {} +unsafe impl Send for MappedRecordBuffer {} +unsafe impl Sync for MappedRecordBuffer {} -impl MappedF32Buffer { - /// Allocate `len` f32s of mapped-pinned memory. +impl MappedRecordBuffer { + /// Allocate `len` records of `T` as 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 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!("MappedF32Buffer alloc ({len} f32): {e}"))? as *mut f32; + .map_err(|e| format!("MappedF32Buffer alloc ({len} f32): {e}"))? as *mut T; - // Zero-init so a stale read returns a safe default (0.0). + // Zero-init so a stale read returns a safe default (all-zero bit pattern). std::ptr::write_bytes(host_ptr, 0, len); let mut dev_ptr_raw = MaybeUninit::uninit(); @@ -61,9 +71,35 @@ impl MappedF32Buffer { }) } + /// Write a single record at index `idx`. Volatile write — the + /// compiler cannot reorder the caller's store past this point. + pub fn write_record(&self, idx: usize, val: T) { + assert!( + idx < self.len, + "MappedRecordBuffer: idx {} >= len {}", + idx, + self.len + ); + unsafe { + std::ptr::write_volatile(self.host_ptr.add(idx), val); + } + } + + /// Read a single record at index `idx`. Volatile read — caller must + /// have synchronised the producing stream first. + pub fn read_record(&self, idx: usize) -> T { + assert!( + idx < self.len, + "MappedRecordBuffer: idx {} >= len {}", + idx, + self.len + ); + unsafe { std::ptr::read_volatile(self.host_ptr.add(idx)) } + } + /// Read all `len` entries via `read_volatile`. Caller must have /// synchronised the producing stream first. - pub fn read_all(&self) -> Vec { + pub fn read_all(&self) -> Vec { let mut out = Vec::with_capacity(self.len); unsafe { for i in 0..self.len { @@ -73,9 +109,36 @@ impl MappedF32Buffer { out } + /// Mutable host slice over the mapped pages. + /// + /// # Safety + /// `host_ptr` is valid for `self.len` `T`s by construction; the + /// caller must not invalidate the allocation while the slice is + /// live and must respect unique-mutable-borrow semantics. + pub fn host_slice_mut(&mut self) -> &mut [T] { + unsafe { std::slice::from_raw_parts_mut(self.host_ptr, self.len) } + } +} + +impl Drop for MappedRecordBuffer { + fn drop(&mut self) { + unsafe { + let _ = cudarc::driver::result::free_host(self.host_ptr as *mut c_void); + } + } +} + +/// Backward-compatible alias for the original `f32`-only buffer name. +/// All pre-existing call sites continue to work unchanged. +pub type MappedF32Buffer = MappedRecordBuffer; + +impl MappedRecordBuffer { /// Write CPU-side data into the buffer via `host_ptr`. The kernel /// reading via `dev_ptr` sees the data after a stream sync barrier /// (mapped-pinned coherence). + /// + /// `f32`-specific bulk-write convenience for the legacy + /// `MappedF32Buffer` call sites that pass `&[f32]`. pub fn write_from_slice(&self, slice: &[f32]) { assert!( slice.len() <= self.len, @@ -89,24 +152,6 @@ impl MappedF32Buffer { } } } - - /// Mutable host slice over the mapped pages. - /// - /// # Safety - /// `host_ptr` is valid for `self.len` f32s by construction; the caller - /// must not invalidate the allocation while the slice is live and - /// must respect unique-mutable-borrow semantics. - pub fn host_slice_mut(&mut self) -> &mut [f32] { - unsafe { std::slice::from_raw_parts_mut(self.host_ptr, self.len) } - } -} - -impl Drop for MappedF32Buffer { - fn drop(&mut self) { - unsafe { - let _ = cudarc::driver::result::free_host(self.host_ptr as *mut c_void); - } - } }