refactor(gpu-log): generic MappedRecordBuffer<T>; MappedF32Buffer as alias

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<T: Copy> {
+    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<T: Copy + Send> Send for MappedRecordBuffer<T> {}
+unsafe impl<T: Copy + Sync> Sync for MappedRecordBuffer<T> {}
 
-impl MappedF32Buffer {
-    /// Allocate `len` f32s of mapped-pinned memory.
+impl<T: Copy> MappedRecordBuffer<T> {
+    /// 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<Self, String> {
-        let num_bytes = len * std::mem::size_of::<f32>();
+        let num_bytes = len * std::mem::size_of::<T>();
         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<f32> {
+    pub fn read_all(&self) -> Vec<T> {
         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<T: Copy> Drop for MappedRecordBuffer<T> {
+    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<f32>;
+
+impl MappedRecordBuffer<f32> {
     /// 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);
-        }
-    }
 }