feat(taxis): derive storage dtype from the element type [recovered; design proposal]

crates/hipcore/src/memory.rs

@@ -205,6 +205,28 @@ impl<T: BytePod> DeviceBuffer<T> {
             "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::<c_void>(), 0, byte_len) },
+            "hipMemset",
+        )
+    }
 }
 
 impl<T: BytePod> Drop for DeviceBuffer<T> {

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::<Self>())`. 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::<f32>()`.
+unsafe impl DTyped for f32 {
+    const DTYPE: DType = DType::F32;
+}
+// SAFETY: `DType::F16.size_in_bytes_exact() == Some(2) == size_of::<half::f16>()`.
+unsafe impl DTyped for half::f16 {
+    const DTYPE: DType = DType::F16;
+}
+// SAFETY: `DType::BF16.size_in_bytes_exact() == Some(2) == size_of::<half::bf16>()`.
+unsafe impl DTyped for half::bf16 {
+    const DTYPE: DType = DType::BF16;
+}
+// SAFETY: `DType::I32.size_in_bytes_exact() == Some(4) == size_of::<i32>()`.
+unsafe impl DTyped for i32 {
+    const DTYPE: DType = DType::I32;
+}
+// SAFETY: `DType::I8.size_in_bytes_exact() == Some(1) == size_of::<i8>()`.
+unsafe impl DTyped for i8 {
+    const DTYPE: DType = DType::I8;
+}
+// SAFETY: `DType::U8.size_in_bytes_exact() == Some(1) == size_of::<u8>()`.
+unsafe impl DTyped for u8 {
+    const DTYPE: DType = DType::U8;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn assert_dtyped_size_matches<T: DTyped>() {
+        assert_eq!(
+            T::DTYPE.size_in_bytes_exact(),
+            Some(core::mem::size_of::<T>()),
+            "DTyped mapping for {:?} disagrees with size_of::<T>()",
+            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::<f32>();
+        assert_dtyped_size_matches::<half::f16>();
+        assert_dtyped_size_matches::<half::bf16>();
+        assert_dtyped_size_matches::<i32>();
+        assert_dtyped_size_matches::<i8>();
+        assert_dtyped_size_matches::<u8>();
+    }
+
+    #[test]
+    fn dtyped_mapping_assigns_the_expected_tag() {
+        assert_eq!(<f32 as DTyped>::DTYPE, DType::F32);
+        assert_eq!(<half::f16 as DTyped>::DTYPE, DType::F16);
+        assert_eq!(<half::bf16 as DTyped>::DTYPE, DType::BF16);
+        assert_eq!(<i32 as DTyped>::DTYPE, DType::I32);
+        assert_eq!(<i8 as DTyped>::DTYPE, DType::I8);
+        assert_eq!(<u8 as DTyped>::DTYPE, DType::U8);
+    }
+}

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;

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<Self> {
         let bytes = dtype.byte_count(elem_count);
-        let buffer = DeviceBuffer::<u8>::alloc(device, bytes)?;
+        let mut buffer = DeviceBuffer::<u8>::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<T: BytePod>(device: &Device, dtype: DType, data: &[T]) -> Result<Self> {
-        // 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<T: DTyped>(device: &Device, data: &[T]) -> Result<Self> {
+        // 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::<u8>(), core::mem::size_of_val(data))
         };
         let buffer = DeviceBuffer::<u8>::from_host(device, bytes)?;
         Ok(Self {
             buffer: Arc::new(buffer),
-            dtype,
+            dtype: T::DTYPE,
             elem_count: data.len(),
             device: device.clone(),
         })

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> {
-        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> {
-        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> {
-        Self::from_host_typed(device, data, shape, DType::BF16)
+        Self::from_host_typed(device, data, shape)
     }
 
-    fn from_host_typed<T: BytePod>(
-        device: &Device,
-        data: &[T],
-        shape: Shape,
-        dtype: DType,
-    ) -> Result<Self> {
+    /// `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<T: DTyped>(device: &Device, data: &[T], shape: Shape) -> Result<Self> {
         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<Self> {
         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 {