hrx/
tensor.rs

1//! Checked tensor metadata and owned views in one coordinated runtime.
2
3use crate::{Completion, Error, Result, Runtime, execution::BufferView};
4use std::sync::Arc;
5mod ops;
6pub use ops::TensorOps;
7
8/// Tensor element representation. Strides are measured in elements.
9#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, serde::Serialize)]
10pub enum DType {
11    /// Unsigned byte.
12    U8,
13    /// Signed 32-bit integer.
14    I32,
15    /// Unsigned 32-bit integer.
16    U32,
17    /// Signed 64-bit integer.
18    I64,
19    /// IEEE half precision.
20    F16,
21    /// Brain floating point.
22    BF16,
23    /// IEEE single precision.
24    F32,
25}
26impl DType {
27    /// Bytes per element.
28    pub const fn bytes(self) -> usize {
29        match self {
30            Self::U8 => 1,
31            Self::F16 | Self::BF16 => 2,
32            Self::I32 | Self::U32 | Self::F32 => 4,
33            Self::I64 => 8,
34        }
35    }
36}
37
38/// Logical axis convention, independent of element strides.
39#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, serde::Serialize)]
40pub enum Layout {
41    /// Arbitrary-rank tensor with axes in caller-defined order.
42    #[default]
43    General,
44    /// Batch, channels, height, width.
45    Nchw,
46    /// Batch, height, width, channels.
47    Nhwc,
48    /// Rows and columns.
49    Rows,
50}
51
52/// Validated tensor shape and non-overlapping element strides.
53#[derive(Clone, Debug, PartialEq, Eq, Hash, serde::Serialize)]
54pub struct TensorDesc {
55    dtype: DType,
56    shape: Vec<usize>,
57    strides: Vec<usize>,
58    layout: Layout,
59    elements: usize,
60    bytes: usize,
61}
62impl TensorDesc {
63    /// Create a contiguous row-major tensor. An empty shape is one scalar;
64    /// any zero dimension makes the tensor empty.
65    pub fn new(dtype: DType, shape: impl Into<Vec<usize>>) -> Result<Self> {
66        let shape = shape.into();
67        if shape.contains(&0) {
68            let strides = vec![1; shape.len()];
69            return Self::strided(dtype, shape, strides);
70        }
71        let mut stride = 1usize;
72        let mut strides = vec![0; shape.len()];
73        for (index, &dimension) in shape.iter().enumerate().rev() {
74            strides[index] = stride;
75            stride = stride
76                .checked_mul(dimension.max(1))
77                .ok_or_else(|| Error::Message("tensor stride overflow".into()))?;
78        }
79        Self::strided(dtype, shape, strides)
80    }
81
82    /// Create a tensor with positive element strides. Overlapping writable
83    /// layouts are rejected; slicing and permutations may leave gaps.
84    pub fn strided(dtype: DType, shape: Vec<usize>, strides: Vec<usize>) -> Result<Self> {
85        if shape.len() != strides.len() || strides.contains(&0) {
86            return Err(Error::Message("invalid tensor strides".into()));
87        }
88        let empty = shape.contains(&0);
89        let elements = if empty {
90            0
91        } else {
92            shape
93                .iter()
94                .try_fold(1usize, |n, &d| n.checked_mul(d))
95                .ok_or_else(|| Error::Message("tensor element count overflow".into()))?
96        };
97        let mut extent = usize::from(!empty);
98        if !empty {
99            let mut axes = shape
100                .iter()
101                .copied()
102                .zip(strides.iter().copied())
103                .filter(|&(d, _)| d > 1)
104                .collect::<Vec<_>>();
105            axes.sort_unstable_by_key(|&(_, stride)| stride);
106            for (dimension, stride) in axes {
107                if stride < extent {
108                    return Err(Error::Message("overlapping tensor strides".into()));
109                }
110                extent = (dimension - 1)
111                    .checked_mul(stride)
112                    .and_then(|n| n.checked_add(extent))
113                    .ok_or_else(|| Error::Message("tensor extent overflow".into()))?;
114            }
115        }
116        let bytes = extent
117            .checked_mul(dtype.bytes())
118            .filter(|&n| n <= isize::MAX as usize)
119            .ok_or_else(|| Error::Message("tensor byte extent overflow".into()))?;
120        Ok(Self {
121            dtype,
122            shape,
123            strides,
124            layout: Layout::General,
125            elements,
126            bytes,
127        })
128    }
129
130    /// Set an axis convention after validating its rank.
131    pub fn with_layout(mut self, layout: Layout) -> Result<Self> {
132        let valid = match layout {
133            Layout::General => true,
134            Layout::Rows => self.shape.len() == 2,
135            Layout::Nchw | Layout::Nhwc => self.shape.len() == 4,
136        };
137        if !valid {
138            return Err(Error::Message("tensor layout rank mismatch".into()));
139        }
140        self.layout = layout;
141        Ok(self)
142    }
143    /// Element representation.
144    pub fn dtype(&self) -> DType {
145        self.dtype
146    }
147    /// Logical dimensions.
148    pub fn shape(&self) -> &[usize] {
149        &self.shape
150    }
151    /// Element strides.
152    pub fn strides(&self) -> &[usize] {
153        &self.strides
154    }
155    /// Axis convention.
156    pub fn layout(&self) -> Layout {
157        self.layout
158    }
159    /// Number of logical elements.
160    pub fn elements(&self) -> usize {
161        self.elements
162    }
163    /// Accessible byte extent including stride gaps.
164    pub fn bytes(&self) -> usize {
165        self.bytes
166    }
167    /// Whether the tensor contains no elements.
168    pub fn is_empty(&self) -> bool {
169        self.elements == 0
170    }
171    /// Whether elements occupy contiguous row-major storage.
172    pub fn is_contiguous(&self) -> bool {
173        if self.is_empty() {
174            return true;
175        }
176        let mut stride = 1usize;
177        for (&dimension, &actual) in self.shape.iter().zip(&self.strides).rev() {
178            if dimension > 1 && actual != stride {
179                return false;
180            }
181            let Some(next) = stride.checked_mul(dimension.max(1)) else {
182                return false;
183            };
184            stride = next;
185        }
186        true
187    }
188}
189
190/// An owned tensor view. Clones retain storage, producer completion, and any
191/// inference slot lease. Empty tensors carry metadata without allocating.
192#[derive(Clone)]
193pub struct DeviceTensor {
194    pub(crate) runtime: Runtime,
195    pub(crate) desc: TensorDesc,
196    pub(crate) view: Option<BufferView>,
197    pub(crate) producer: Completion,
198    pub(crate) lease: Option<Arc<dyn Send + Sync>>,
199}
200impl DeviceTensor {
201    /// Validated shape, dtype, layout, and extent.
202    pub fn desc(&self) -> &TensorDesc {
203        &self.desc
204    }
205    /// Completion of the operation that produced these values.
206    pub fn completion(&self) -> &Completion {
207        &self.producer
208    }
209    /// Checked owned binding. Its lifetime also retains the inference slot.
210    /// Empty tensors have no native binding.
211    pub fn binding(&self) -> Option<BufferView> {
212        self.view.clone().map(|view| match &self.lease {
213            Some(lease) => view.retain(lease.clone()),
214            None => view,
215        })
216    }
217    /// Reinterpret a checked subregion, retaining its producer and slot lease.
218    pub fn view(&self, byte_offset: usize, desc: TensorDesc) -> Result<Self> {
219        let end = byte_offset
220            .checked_add(desc.bytes())
221            .filter(|&end| end <= self.desc.bytes())
222            .ok_or_else(|| Error::Message("tensor view exceeds storage".into()))?;
223        let origin = self.view.as_ref().map_or(0, BufferView::offset);
224        if !(origin + byte_offset).is_multiple_of(desc.dtype().bytes()) {
225            return Err(Error::Message("unaligned tensor view".into()));
226        }
227        let view = if desc.is_empty() {
228            None
229        } else {
230            Some(
231                self.view
232                    .as_ref()
233                    .ok_or_else(|| Error::Message("empty tensor has no storage".into()))?
234                    .slice(byte_offset..end)?,
235            )
236        };
237        Ok(Self {
238            runtime: self.runtime.clone(),
239            desc,
240            view,
241            producer: self.producer.clone(),
242            lease: self.lease.clone(),
243        })
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250    #[test]
251    fn shapes_distinguish_empty_scalar_and_strided_storage() {
252        let scalar = TensorDesc::new(DType::F32, vec![]).unwrap();
253        assert_eq!((scalar.elements(), scalar.bytes()), (1, 4));
254        let empty = TensorDesc::new(DType::F32, vec![2, 0, 3]).unwrap();
255        assert_eq!((empty.elements(), empty.bytes()), (0, 0));
256        let huge_empty = TensorDesc::new(DType::F32, vec![usize::MAX, 0, usize::MAX]).unwrap();
257        assert!(huge_empty.is_empty() && huge_empty.is_contiguous());
258        let padded = TensorDesc::strided(DType::F32, vec![2, 3], vec![4, 1]).unwrap();
259        assert_eq!((padded.elements(), padded.bytes()), (6, 28));
260        assert!(!padded.is_contiguous());
261        assert!(TensorDesc::strided(DType::F32, vec![2, 3], vec![1, 1]).is_err());
262        assert!(TensorDesc::new(DType::F32, vec![usize::MAX]).is_err());
263        assert!(
264            TensorDesc::new(DType::U8, vec![2, 3])
265                .unwrap()
266                .with_layout(Layout::Nchw)
267                .is_err()
268        );
269        let transpose = TensorDesc::strided(DType::U8, vec![3, 2], vec![1, 3]).unwrap();
270        assert_eq!(transpose.bytes(), 6);
271    }
272}
hrx/tensor.rs

hrx/
tensor.rs
