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

hrx/
tensor.rs
