atlas_core/
tensor.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3use crate::dtype::DType;
4
5/// A zero-copy reference to a GPU tensor, typically from PyTorch via data_ptr().
6///
7/// Atlas never allocates or copies tensor data — it receives raw CUDA device
8/// pointers from Python and wraps them for kernel launches.
9#[derive(Debug, Clone)]
10pub struct TensorRef {
11    /// Raw CUDA device pointer (from torch.Tensor.data_ptr())
12    pub ptr: u64,
13
14    /// Shape dimensions (e.g., [batch, seq_len, hidden_size])
15    pub shape: Vec<usize>,
16
17    /// Strides in elements (not bytes)
18    pub strides: Vec<usize>,
19
20    /// Element data type
21    pub dtype: DType,
22}
23
24impl TensorRef {
25    /// Create a new tensor reference from a raw pointer and shape.
26    /// Assumes contiguous (row-major) layout.
27    pub fn new(ptr: u64, shape: Vec<usize>, dtype: DType) -> Self {
28        let strides = Self::contiguous_strides(&shape);
29        Self {
30            ptr,
31            shape,
32            strides,
33            dtype,
34        }
35    }
36
37    /// Total number of elements.
38    pub fn numel(&self) -> usize {
39        self.shape.iter().product()
40    }
41
42    /// Total size in bytes.
43    pub fn size_bytes(&self) -> usize {
44        let bits = self.numel() * self.dtype.element_size_bits();
45        bits.div_ceil(8) // round up for sub-byte types
46    }
47
48    /// Number of dimensions.
49    pub fn ndim(&self) -> usize {
50        self.shape.len()
51    }
52
53    /// Compute contiguous (C-order / row-major) strides from shape.
54    fn contiguous_strides(shape: &[usize]) -> Vec<usize> {
55        let mut strides = vec![1usize; shape.len()];
56        for i in (0..shape.len().saturating_sub(1)).rev() {
57            strides[i] = strides[i + 1] * shape[i + 1];
58        }
59        strides
60    }
61
62    /// Raw pointer cast to a typed device pointer (for kernel launches).
63    pub fn as_device_ptr<T>(&self) -> *const T {
64        self.ptr as *const T
65    }
66
67    /// Mutable raw pointer cast.
68    pub fn as_device_ptr_mut<T>(&self) -> *mut T {
69        self.ptr as *mut T
70    }
71}
atlas_core/tensor.rs

atlas_core/
tensor.rs
