metrale_core/
tensor.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `TensorRef`, a non-owning description of a device tensor.
4//!
5//! Owner: metrale-core.
6//! Invariants: none beyond the types.
7
8use crate::dtype::DType;
9
10/// 2026-09-25: A non-owning reference to a device tensor: a raw pointer plus
11/// shape, strides and element type. It neither allocates nor frees.
12#[derive(Debug, Clone)]
13pub struct TensorRef {
14    pub ptr: u64,
15
16    pub shape: Vec<usize>,
17
18    /// 2026-09-25: Strides in elements, not bytes.
19    pub strides: Vec<usize>,
20
21    pub dtype: DType,
22}
23
24impl TensorRef {
25    /// 2026-09-25: Create a tensor reference with contiguous row-major strides.
26    pub fn new(ptr: u64, shape: Vec<usize>, dtype: DType) -> Self {
27        let strides = Self::contiguous_strides(&shape);
28        Self {
29            ptr,
30            shape,
31            strides,
32            dtype,
33        }
34    }
35
36    pub fn numel(&self) -> usize {
37        self.shape.iter().product()
38    }
39
40    /// 2026-09-25: Total size in bytes, rounded up for sub-byte element types.
41    pub fn size_bytes(&self) -> usize {
42        let bits = self.numel() * self.dtype.element_size_bits();
43        bits.div_ceil(8)
44    }
45
46    pub fn ndim(&self) -> usize {
47        self.shape.len()
48    }
49
50    fn contiguous_strides(shape: &[usize]) -> Vec<usize> {
51        let mut strides = vec![1usize; shape.len()];
52        for i in (0..shape.len().saturating_sub(1)).rev() {
53            strides[i] = strides[i + 1] * shape[i + 1];
54        }
55        strides
56    }
57
58    pub fn as_device_ptr<T>(&self) -> *const T {
59        self.ptr as *const T
60    }
61
62    pub fn as_device_ptr_mut<T>(&self) -> *mut T {
63        self.ptr as *mut T
64    }
65}
metrale_core/tensor.rs

metrale_core/
tensor.rs
