Skip to main content

ariadnetor_tensor/dense/
layout.rs

1//! `DenseLayout`: interpretation half of the dense tensor split.
2//!
3//! Carries shape and memory order. Data lives on
4//! [`DenseStorage<T>`](crate::DenseStorage); the wrapper
5//! [`DenseTensorData<T>`](crate::DenseTensorData) joins the two with a
6//! length-consistency check.
7
8use ariadnetor_core::backend::MemoryOrder;
9
10use crate::TensorLayout;
11
12/// Interpretation half of the dense tensor split.
13///
14/// Holds the logical shape and the memory order the paired
15/// [`DenseStorage`](crate::DenseStorage) is laid out in. Operations
16/// consuming a [`DenseTensorData`](crate::DenseTensorData) consult
17/// `order()` to decide whether to repack at their boundary.
18#[derive(Clone, Debug)]
19pub struct DenseLayout {
20    shape: Vec<usize>,
21    order: MemoryOrder,
22}
23
24impl DenseLayout {
25    /// Construct a `DenseLayout` from shape and memory order.
26    pub fn new(shape: Vec<usize>, order: MemoryOrder) -> Self {
27        Self { shape, order }
28    }
29
30    /// Logical shape.
31    pub fn shape(&self) -> &[usize] {
32        &self.shape
33    }
34
35    /// Rank (number of dimensions).
36    pub fn rank(&self) -> usize {
37        self.shape.len()
38    }
39
40    /// Memory order the paired storage is laid out in.
41    pub fn order(&self) -> MemoryOrder {
42        self.order
43    }
44}
45
46impl TensorLayout for DenseLayout {
47    fn shape(&self) -> &[usize] {
48        &self.shape
49    }
50
51    fn storage_extent(&self) -> usize {
52        self.shape.iter().product()
53    }
54}