From 5f299d6ebd224f521e233bd8f1b13373bf95fba3 Mon Sep 17 00:00:00 2001 From: forkwright Date: Mon, 24 Aug 2026 09:05:09 -0500 Subject: [PATCH] feat(taxis): bind host element types to their runtime dtype tag Adds BytePod, a trait pairing a host element type with the DType it represents, and derives HipStorage::from_host's stored dtype from T::DTYPE rather than taking a free DType parameter. A caller can no longer construct storage whose declared dtype disagrees with the byte layout of the data actually copied: the mismatch stops being a checked error and becomes unrepresentable. RECOVERED WORK. This was found uncommitted in the shared logismos clone, against 8b82822 rather than the current main, with no branch and no author recorded. It is committed here at its true base so it stops depending on one machine's working tree surviving. It is preserved, not verified. Nothing here has been compiled or tested, and the tree it came from may have been mid-edit. Treat the safety contract in BytePod's doc comment as an assertion needing proof rather than one already discharged, and rebase onto main before believing any build result. --- crates/hipcore/src/memory.rs | 22 +++++++++ crates/taxis/src/dtype.rs | 87 ++++++++++++++++++++++++++++++++++++ crates/taxis/src/lib.rs | 2 +- crates/taxis/src/storage.rs | 26 +++++++---- crates/taxis/src/tensor.rs | 31 ++++++------- 5 files changed, 141 insertions(+), 27 deletions(-) diff --git a/crates/hipcore/src/memory.rs b/crates/hipcore/src/memory.rs index 5426f5c..ca62f12 100644 --- a/crates/hipcore/src/memory.rs +++ b/crates/hipcore/src/memory.rs @@ -205,6 +205,28 @@ impl DeviceBuffer { "hipMemcpyAsync(HtoD)", ) } + + /// Zero every byte of the allocation via `hipMemset`. + /// + /// `hipMalloc` never zeroes device memory on its own — a fresh + /// allocation can carry residual bytes from whatever this VRAM + /// region held before. Call this whenever a caller-facing "zeroed" + /// contract depends on it. + /// + /// # Errors + /// + /// [`Error::Runtime`] on HIP failure. + pub fn zero_fill(&mut self) -> Result<()> { + self.device.make_current()?; + let byte_len = self.byte_len(); + // SAFETY: `self.ptr` is a live device allocation owned by this + // `DeviceBuffer`, sized `byte_len` bytes; `hipMemset` writes + // exactly that many bytes starting at the pointer. + check( + unsafe { ffi::hipMemset(self.ptr.as_ptr().cast::(), 0, byte_len) }, + "hipMemset", + ) + } } impl Drop for DeviceBuffer { diff --git a/crates/taxis/src/dtype.rs b/crates/taxis/src/dtype.rs index 64faeff..7ecc01f 100644 --- a/crates/taxis/src/dtype.rs +++ b/crates/taxis/src/dtype.rs @@ -1,5 +1,7 @@ //! Runtime dtype enumeration. +use hipcore::BytePod; + /// Runtime dtype tag. /// /// `#[non_exhaustive]` so the public surface can grow without breaking @@ -60,3 +62,88 @@ impl DType { matches!(self, Self::F32 | Self::F16 | Self::BF16) } } + +/// Binds a host element type to the runtime [`DType`] tag it represents. +/// +/// [`HipStorage::from_host`](crate::storage::HipStorage::from_host) +/// derives its stored `dtype` from `T::DTYPE` instead of accepting a +/// free [`DType`] parameter, so a caller cannot construct storage whose +/// declared dtype disagrees with the byte layout of the data actually +/// copied — the mismatch is unrepresentable rather than checked. +/// +/// # Safety +/// +/// Implementors must guarantee `DTYPE.size_in_bytes_exact()` is +/// `Some(size_of::())`. Violating this lets `from_host` construct +/// a [`crate::storage::HipStorage`] whose `dtype` / `elem_count` / +/// byte length are mutually inconsistent — the exact corruption this +/// trait exists to foreclose. +pub unsafe trait DTyped: BytePod { + /// Runtime dtype tag this Rust type represents. + const DTYPE: DType; +} + +// SAFETY: `DType::F32.size_in_bytes_exact() == Some(4) == size_of::()`. +unsafe impl DTyped for f32 { + const DTYPE: DType = DType::F32; +} +// SAFETY: `DType::F16.size_in_bytes_exact() == Some(2) == size_of::()`. +unsafe impl DTyped for half::f16 { + const DTYPE: DType = DType::F16; +} +// SAFETY: `DType::BF16.size_in_bytes_exact() == Some(2) == size_of::()`. +unsafe impl DTyped for half::bf16 { + const DTYPE: DType = DType::BF16; +} +// SAFETY: `DType::I32.size_in_bytes_exact() == Some(4) == size_of::()`. +unsafe impl DTyped for i32 { + const DTYPE: DType = DType::I32; +} +// SAFETY: `DType::I8.size_in_bytes_exact() == Some(1) == size_of::()`. +unsafe impl DTyped for i8 { + const DTYPE: DType = DType::I8; +} +// SAFETY: `DType::U8.size_in_bytes_exact() == Some(1) == size_of::()`. +unsafe impl DTyped for u8 { + const DTYPE: DType = DType::U8; +} + +#[cfg(test)] +mod tests { + use super::*; + + fn assert_dtyped_size_matches() { + assert_eq!( + T::DTYPE.size_in_bytes_exact(), + Some(core::mem::size_of::()), + "DTyped mapping for {:?} disagrees with size_of::()", + T::DTYPE + ); + } + + /// Pins the exact invariant `HipStorage::from_host` relies on as + /// its ONLY source of `dtype`: if any `DTyped` impl declared a tag + /// whose byte size does not match its Rust type, `from_host` would + /// silently reproduce a dtype/byte-layout mismatch. There is no + /// second, independent `dtype` input left to cross-check against — + /// this mapping table IS the contract. + #[test] + fn dtyped_mapping_matches_declared_byte_size() { + assert_dtyped_size_matches::(); + assert_dtyped_size_matches::(); + assert_dtyped_size_matches::(); + assert_dtyped_size_matches::(); + assert_dtyped_size_matches::(); + assert_dtyped_size_matches::(); + } + + #[test] + fn dtyped_mapping_assigns_the_expected_tag() { + assert_eq!(::DTYPE, DType::F32); + assert_eq!(::DTYPE, DType::F16); + assert_eq!(::DTYPE, DType::BF16); + assert_eq!(::DTYPE, DType::I32); + assert_eq!(::DTYPE, DType::I8); + assert_eq!(::DTYPE, DType::U8); + } +} diff --git a/crates/taxis/src/lib.rs b/crates/taxis/src/lib.rs index c33f641..9560abc 100644 --- a/crates/taxis/src/lib.rs +++ b/crates/taxis/src/lib.rs @@ -37,7 +37,7 @@ pub mod shape; pub mod storage; pub mod tensor; -pub use crate::dtype::DType; +pub use crate::dtype::{DType, DTyped}; pub use crate::error::{Error, Result}; pub use crate::layout::Layout; pub use crate::shape::Shape; diff --git a/crates/taxis/src/storage.rs b/crates/taxis/src/storage.rs index bee6581..ba934d5 100644 --- a/crates/taxis/src/storage.rs +++ b/crates/taxis/src/storage.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use hipcore::{BytePod, Device, DeviceBuffer}; -use crate::dtype::DType; +use crate::dtype::{DType, DTyped}; use crate::error::{Error, Result}; /// Type-erased CPU storage. One variant per supported dtype that has @@ -138,10 +138,15 @@ impl HipStorage { /// /// # Errors /// - /// [`Error::Hip`] on allocation failure. + /// [`Error::Hip`] on allocation or zero-fill failure. pub fn alloc(device: &Device, dtype: DType, elem_count: usize) -> Result { let bytes = dtype.byte_count(elem_count); - let buffer = DeviceBuffer::::alloc(device, bytes)?; + let mut buffer = DeviceBuffer::::alloc(device, bytes)?; + // WHY: `hipMalloc` never zeroes memory; without this, any + // caller trusting this fn's "zeroed" contract (or a Phase-1 + // kernel that does not overwrite every output element) would + // observe residual bytes from a prior allocation. + buffer.zero_fill()?; Ok(Self { buffer: Arc::new(buffer), dtype, @@ -152,20 +157,25 @@ impl HipStorage { /// Copy a typed host slice to a freshly allocated HIP storage. /// + /// The stored `dtype` is derived from `T::DTYPE` + /// ([`DTyped`]) rather than accepted as a separate + /// parameter, so the declared dtype cannot disagree with the byte + /// layout of `data`. + /// /// # Errors /// /// [`Error::Hip`] on allocation or memcpy failure. - pub fn from_host(device: &Device, dtype: DType, data: &[T]) -> Result { - // SAFETY: `T: BytePod` guarantees every bit pattern is valid - // and the type is `Copy`. Transmuting the slice to a byte - // view is defined. + pub fn from_host(device: &Device, data: &[T]) -> Result { + // SAFETY: `T: BytePod` (via `DTyped`'s supertrait) guarantees + // every bit pattern is valid and the type is `Copy`. + // Transmuting the slice to a byte view is defined. let bytes: &[u8] = unsafe { core::slice::from_raw_parts(data.as_ptr().cast::(), core::mem::size_of_val(data)) }; let buffer = DeviceBuffer::::from_host(device, bytes)?; Ok(Self { buffer: Arc::new(buffer), - dtype, + dtype: T::DTYPE, elem_count: data.len(), device: device.clone(), }) diff --git a/crates/taxis/src/tensor.rs b/crates/taxis/src/tensor.rs index b5d67a0..420370f 100644 --- a/crates/taxis/src/tensor.rs +++ b/crates/taxis/src/tensor.rs @@ -2,9 +2,9 @@ use std::sync::Arc; -use hipcore::{BytePod, Device}; +use hipcore::Device; -use crate::dtype::DType; +use crate::dtype::{DType, DTyped}; use crate::error::{Error, Result}; use crate::layout::Layout; use crate::shape::Shape; @@ -44,7 +44,7 @@ impl Tensor { /// [`Error::ShapeMismatch`] when `data.len() != shape.elem_count()`. /// [`Error::Hip`] on device allocation or copy failure. pub fn from_host_f32(device: &Device, data: &[f32], shape: Shape) -> Result { - Self::from_host_typed(device, data, shape, DType::F32) + Self::from_host_typed(device, data, shape) } /// Construct a HIP tensor from a host slice of `half::f16`. @@ -53,7 +53,7 @@ impl Tensor { /// /// See [`Self::from_host_f32`]. pub fn from_host_f16(device: &Device, data: &[half::f16], shape: Shape) -> Result { - Self::from_host_typed(device, data, shape, DType::F16) + Self::from_host_typed(device, data, shape) } /// Construct a HIP tensor from a host slice of `half::bf16`. @@ -62,15 +62,13 @@ impl Tensor { /// /// See [`Self::from_host_f32`]. pub fn from_host_bf16(device: &Device, data: &[half::bf16], shape: Shape) -> Result { - Self::from_host_typed(device, data, shape, DType::BF16) + Self::from_host_typed(device, data, shape) } - fn from_host_typed( - device: &Device, - data: &[T], - shape: Shape, - dtype: DType, - ) -> Result { + /// `dtype` is derived from `T::DTYPE` ([`DTyped`]), never accepted + /// as a free parameter — a caller cannot request a `T`/dtype pair + /// that disagrees, because there is no second value to disagree. + fn from_host_typed(device: &Device, data: &[T], shape: Shape) -> Result { if data.len() != shape.elem_count() { return Err(Error::ShapeMismatch { op: "from_host_typed", @@ -81,27 +79,24 @@ impl Tensor { ), }); } - let storage = HipStorage::from_host(device, dtype, data)?; + let storage = HipStorage::from_host(device, data)?; let layout = Layout::contiguous(shape); Ok(Self { inner: Arc::new(TensorInner { - dtype, + dtype: T::DTYPE, storage: Arc::new(Storage::Hip(storage)), layout, }), }) } - /// Allocate a HIP tensor of the given shape + dtype, uninitialised. + /// Allocate a HIP tensor of the given shape + dtype, zero-filled. /// /// # Errors /// - /// [`Error::Hip`] on allocation failure. + /// [`Error::Hip`] on allocation or zero-fill failure. pub fn zeros_hip(device: &Device, dtype: DType, shape: Shape) -> Result { let elem = shape.elem_count(); - // Caller may want genuinely zeroed memory; `hipMalloc` does - // not zero. Phase-1 kernels always overwrite their outputs, - // so this allocation stays uninitialised by design. let storage = HipStorage::alloc(device, dtype, elem)?; let layout = Layout::contiguous(shape); Ok(Self {