Struct TensorRef 

pub struct TensorRef {
    pub ptr: u64,
    pub shape: Vec<usize>,
    pub strides: Vec<usize>,
    pub dtype: DType,
}

A zero-copy reference to a GPU tensor, typically from PyTorch via data_ptr().

Atlas never allocates or copies tensor data — it receives raw CUDA device pointers from Python and wraps them for kernel launches.

Fields

ptr: u64

Raw CUDA device pointer (from torch.Tensor.data_ptr())

shape: Vec<usize>

Shape dimensions (e.g., [batch, seq_len, hidden_size])

strides: Vec<usize>

Strides in elements (not bytes)

dtype: DType

Element data type

Implementations

impl TensorRef

pub fn new(ptr: u64, shape: Vec<usize>, dtype: DType) -> Self

Create a new tensor reference from a raw pointer and shape. Assumes contiguous (row-major) layout.

pub fn numel(&self) -> usize

Total number of elements.

pub fn size_bytes(&self) -> usize

Total size in bytes.

pub fn ndim(&self) -> usize

Number of dimensions.

pub fn as_device_ptr<T>(&self) -> *const T

Raw pointer cast to a typed device pointer (for kernel launches).

pub fn as_device_ptr_mut<T>(&self) -> *mut T

Mutable raw pointer cast.

Trait Implementations

impl Clone for TensorRef

fn clone(&self) -> TensorRef

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for TensorRef

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.