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 {