luma_tensor/tensor/
layout.rs

1use crate::{Dim, Error, Result, Shape};
2
3#[derive(Debug, PartialEq, Eq, Clone)]
4pub struct Layout {
5    pub(crate) shape: Shape,
6    pub(crate) stride: Vec<usize>,
7    pub(crate) start_offset: usize,
8}
9
10impl Layout {
11    pub fn new<S: Into<Shape>>(shape: S, stride: Vec<usize>, start_offset: usize) -> Self {
12        Self { shape: shape.into(), stride, start_offset }
13    }
14
15    pub fn contiguous<S: Into<Shape>>(shape: S) -> Self {
16        let shape = shape.into();
17        let stride = shape.stride_contiguous();
18        Self { shape, stride, start_offset: 0 }
19    }
20
21    pub fn contiguous_with_offset<S: Into<Shape>>(shape: S, start_offset: usize) -> Self {
22        let shape = shape.into();
23        let stride = shape.stride_contiguous();
24        Self { shape, stride, start_offset }
25    }
26
27    pub fn dims(&self) -> &[usize] {
28        self.shape.dims()
29    }
30
31    pub fn dim<D: Dim>(&self, dim: D) -> Result<usize> {
32        let dim = dim.to_index(&self.shape, "dim")?;
33        Ok(self.dims()[dim])
34    }
35
36    pub fn shape(&self) -> &Shape {
37        &self.shape
38    }
39
40    pub fn stride(&self) -> &[usize] {
41        &self.stride
42    }
43
44    pub fn start_offset(&self) -> usize {
45        self.start_offset
46    }
47
48    pub fn element_count(&self) -> usize {
49        self.shape().element_count()
50    }
51
52    pub fn is_contiguous(&self) -> bool {
53        self.shape.is_contiguous(&self.stride)
54    }
55
56    pub fn slice(&self, dim: usize, start: usize, end: usize, step: usize) -> Result<Self> {
57        let dims = self.shape().dims();
58        if dim >= dims.len() {
59            Err(Error::DimOutOfRange { shape: self.shape().clone(), dim: dim as i32, op: "slice" })?;
60        }
61        if step == 0 {
62            return Err(Error::NarrowInvalidArgs { shape: self.shape.clone(), dim, start, len: 0, msg: "step cannot be 0" }.into());
63        }
64
65        if start > end || end > dims[dim] {
66            return Err(Error::NarrowInvalidArgs {
67                shape: self.shape.clone(),
68                dim,
69                start,
70                len: end.saturating_sub(start),
71                msg: "index out of range",
72            }
73            .into());
74        }
75
76        let new_len = if start == end { 0 } else { (start..end).step_by(step).len() };
77
78        let mut new_dims = dims.to_vec();
79        new_dims[dim] = new_len;
80
81        let mut new_stride = self.stride.clone();
82        new_stride[dim] *= step;
83
84        Ok(Self::new(new_dims, new_stride, self.start_offset + self.stride[dim] * start))
85    }
86
87    pub fn narrow(&self, dim: usize, start: usize, len: usize) -> Result<Self> {
88        self.slice(dim, start, start + len, 1)
89    }
90
91    pub fn transpose(&self, dim1: usize, dim2: usize) -> Result<Self> {
92        let rank = self.shape.rank();
93        if rank <= dim1 || rank <= dim2 {
94            Err(Error::UnexpectedNumberOfDims { expected: usize::max(dim1, dim2), got: rank, shape: self.shape().clone() })?
95        }
96
97        let mut stride = self.stride().to_vec();
98        let mut dims = self.shape().dims().to_vec();
99        dims.swap(dim1, dim2);
100        stride.swap(dim1, dim2);
101
102        Ok(Self::new(dims, stride, self.start_offset))
103    }
104
105    pub fn broadcast_as<S: Into<Shape>>(&self, shape: S) -> Result<Self> {
106        let shape = shape.into();
107        if shape.rank() < self.shape().rank() {
108            return Err(Error::BroadcastIncompatibleShapes { src_shape: self.shape().clone(), dst_shape: shape })?;
109        }
110
111        let added_dims = shape.rank() - self.shape().rank();
112        let mut stride = vec![0; added_dims];
113        for (&dst_dim, (&src_dim, &src_stride)) in shape.dims()[added_dims..].iter().zip(self.dims().iter().zip(self.stride())) {
114            let s = if dst_dim == src_dim {
115                src_stride
116            } else if src_dim != 1 {
117                return Err(Error::BroadcastIncompatibleShapes { src_shape: self.shape().clone(), dst_shape: shape })?;
118            } else {
119                0
120            };
121            stride.push(s)
122        }
123        Ok(Self { shape, stride, start_offset: self.start_offset })
124    }
125
126    pub fn permute(&self, idxs: &[usize]) -> Result<Self> {
127        let is_permutation = idxs.len() == self.shape.rank() && (0..idxs.len()).all(|i| idxs.contains(&i));
128        if !is_permutation {
129            crate::bail!("dimension mismatch in permute, tensor {:?}, dims: {:?}", self.dims(), idxs)
130        }
131        let stride = self.stride();
132        let dims = self.shape().dims();
133        let mut perm_stride = stride.to_vec();
134        let mut perm_dims = dims.to_vec();
135        for (i, &idx) in idxs.iter().enumerate() {
136            perm_stride[i] = stride[idx];
137            perm_dims[i] = dims[idx];
138        }
139        Ok(Self { shape: Shape::from(perm_dims), stride: perm_stride, start_offset: self.start_offset })
140    }
141
142    /// Returns an iterator over **storage indices**.
143    ///
144    /// This iterator yields the linear (flat) indices as they are laid out
145    /// in the underlying storage buffer. The order depends on the memory
146    /// layout (e.g., row-major / column-major / with strides).
147    ///
148    /// Example for shape = (2, 2) in row-major layout:
149    /// yields: `0, 1, 2, 3`
150    pub fn storage_indices<'a>(&'a self) -> StorageIndices<'a> {
151        StorageIndices::from_layout(self)
152    }
153}
154
155//////////////////////////////////////////////////////////////////////////////////////
156///                  StorageIndices
157//////////////////////////////////////////////////////////////////////////////////////
158
159#[derive(Debug, Clone)]
160pub enum StorageIndices<'a> {
161    Uncontiguous(UncontiguousStorageIndices<'a>),
162    Contiguous(ContiguousStorageIndices),
163}
164
165impl<'a> StorageIndices<'a> {
166    pub fn from_layout(l: &'a Layout) -> Self {
167        if l.is_contiguous() {
168            Self::Contiguous(ContiguousStorageIndices::from_layout(l))
169        } else {
170            Self::Uncontiguous(UncontiguousStorageIndices::from_layout(l))
171        }
172    }
173
174    pub fn reset(&mut self) {
175        match self {
176            Self::Uncontiguous(index) => index.reset(),
177            Self::Contiguous(index) => index.reset(),
178        }
179    }
180
181    pub fn len(&self) -> usize {
182        match self {
183            Self::Uncontiguous(index) => index.len(),
184            Self::Contiguous(index) => index.len(),
185        }
186    }
187
188    pub fn is_empty(&self) -> bool {
189        self.len() == 0
190    }
191}
192
193impl<'a> Iterator for StorageIndices<'a> {
194    type Item = usize;
195
196    fn next(&mut self) -> Option<Self::Item> {
197        match self {
198            Self::Contiguous(i) => i.next(),
199            Self::Uncontiguous(i) => i.next(),
200        }
201    }
202}
203
204#[derive(Debug, Clone)]
205pub struct ContiguousStorageIndices {
206    pub storage_index: usize,
207
208    pub begin_index: usize,
209    pub end_index: usize,
210}
211
212impl ContiguousStorageIndices {
213    fn from_layout(l: &Layout) -> Self {
214        Self { begin_index: l.start_offset(), storage_index: l.start_offset(), end_index: l.start_offset() + l.element_count() }
215    }
216
217    fn reset(&mut self) {
218        self.storage_index = self.begin_index;
219    }
220
221    fn len(&self) -> usize {
222        self.end_index - self.begin_index
223    }
224}
225
226impl Iterator for ContiguousStorageIndices {
227    type Item = usize;
228
229    fn next(&mut self) -> Option<Self::Item> {
230        if self.storage_index >= self.end_index {
231            None
232        } else {
233            let index = self.storage_index;
234            self.storage_index += 1;
235            Some(index)
236        }
237    }
238}
239
240impl<S: Into<Shape>> From<S> for Layout {
241    fn from(value: S) -> Self {
242        Layout::contiguous(value.into())
243    }
244}
245
246#[derive(Debug, Clone)]
247pub struct UncontiguousStorageIndices<'a> {
248    begin_index: Option<usize>,
249    next_storage_index: Option<usize>,
250    multi_index: Vec<usize>,
251    dims: &'a [usize],
252    stride: &'a [usize],
253    len: usize,
254}
255
256impl<'a> UncontiguousStorageIndices<'a> {
257    fn new(dims: &'a [usize], stride: &'a [usize], start_offset: usize) -> Self {
258        let elem_count: usize = dims.iter().product();
259        let next_storage_index = if elem_count == 0 {
260            None
261        } else {
262            // This applies to the scalar case.
263            Some(start_offset)
264        };
265        UncontiguousStorageIndices {
266            begin_index: next_storage_index,
267            next_storage_index,
268            multi_index: vec![0; dims.len()],
269            dims,
270            stride,
271            len: elem_count,
272        }
273    }
274
275    fn from_layout(l: &'a Layout) -> Self {
276        Self::new(l.dims(), l.stride(), l.start_offset())
277    }
278
279    pub fn reset(&mut self) {
280        self.next_storage_index = self.begin_index;
281    }
282
283    pub fn len(&self) -> usize {
284        self.len
285    }
286
287    pub fn is_empty(&self) -> bool {
288        self.len == 0
289    }
290}
291
292impl Iterator for UncontiguousStorageIndices<'_> {
293    type Item = usize;
294
295    fn next(&mut self) -> Option<Self::Item> {
296        let storage_index = self.next_storage_index?;
297        let mut updated = false;
298        let mut next_storage_index = storage_index;
299        for ((multi_i, max_i), stride_i) in self.multi_index.iter_mut().zip(self.dims.iter()).zip(self.stride.iter()).rev() {
300            let next_i = *multi_i + 1;
301            if next_i < *max_i {
302                *multi_i = next_i;
303                updated = true;
304                next_storage_index += stride_i;
305                break;
306            } else {
307                next_storage_index -= *multi_i * stride_i;
308                *multi_i = 0
309            }
310        }
311        self.next_storage_index = if updated { Some(next_storage_index) } else { None };
312        Some(storage_index)
313    }
314}
luma_tensor/tensor/layout.rs

luma_tensor/tensor/
layout.rs
