tenferro_tensor_core/
layout.rs

1use crate::{
2    checked_logical_element_count, checked_product, col_major_strides, validate_permutation,
3    DynRank, Error, Result, ShapeVec, SliceSpec, StrideVec, TensorRank,
4};
5use smallvec::SmallVec;
6use std::collections::HashSet;
7
8/// Maximum logical elements for exact mutable-overlap validation.
9///
10/// Larger layouts must pass the sufficient stride-span proof. This keeps the
11/// fallback bounded because it enumerates logical elements and stores visited
12/// physical offsets.
13const MUTABLE_NO_OVERLAP_EXACT_ELEMENT_LIMIT: usize = 4096;
14
15pub(crate) fn reachable_offset_range(
16    shape: &[usize],
17    strides: &[isize],
18    offset: isize,
19) -> Result<Option<(isize, isize)>> {
20    if shape.contains(&0) {
21        return Ok(None);
22    }
23
24    let mut min = offset;
25    let mut max = offset;
26    for (&extent, &stride) in shape.iter().zip(strides) {
27        let last = isize::try_from(extent.saturating_sub(1)).map_err(|_| Error::IntegerOverflow)?;
28        let delta = last.checked_mul(stride).ok_or(Error::IntegerOverflow)?;
29        if delta < 0 {
30            min = min.checked_add(delta).ok_or(Error::IntegerOverflow)?;
31        } else {
32            max = max.checked_add(delta).ok_or(Error::IntegerOverflow)?;
33        }
34    }
35    Ok(Some((min, max)))
36}
37
38pub(crate) fn validate_reachable_bounds(
39    shape: &[usize],
40    strides: &[isize],
41    offset: isize,
42    buffer_len: usize,
43) -> Result<()> {
44    if shape.len() != strides.len() {
45        return Err(Error::RankMismatch {
46            expected: shape.len(),
47            actual: strides.len(),
48        });
49    }
50
51    match reachable_offset_range(shape, strides, offset)? {
52        Some((min, max)) => {
53            if min < 0 {
54                return Err(Error::ViewOutOfBounds);
55            }
56            let max = usize::try_from(max).map_err(|_| Error::IntegerOverflow)?;
57            if max < buffer_len {
58                Ok(())
59            } else {
60                Err(Error::ViewOutOfBounds)
61            }
62        }
63        None => {
64            if offset < 0 {
65                return Err(Error::ViewOutOfBounds);
66            }
67            let offset = usize::try_from(offset).map_err(|_| Error::IntegerOverflow)?;
68            if offset <= buffer_len {
69                Ok(())
70            } else {
71                Err(Error::ViewOutOfBounds)
72            }
73        }
74    }
75}
76
77fn layout_from_vecs<R: TensorRank>(
78    shape: ShapeVec,
79    strides: StrideVec,
80    offset: isize,
81    buffer_len: usize,
82) -> Result<TensorLayout<R>> {
83    TensorLayout::from_parts(
84        R::shape_from_vec(shape)?,
85        R::strides_from_vec(strides)?,
86        offset,
87        buffer_len,
88    )
89}
90
91fn positive_ceil_div(numerator: isize, denominator: isize) -> Result<usize> {
92    if numerator < 0 || denominator <= 0 {
93        return Err(Error::IntegerOverflow);
94    }
95    let extent = if numerator == 0 {
96        0
97    } else {
98        1 + (numerator - 1) / denominator
99    };
100    usize::try_from(extent).map_err(|_| Error::IntegerOverflow)
101}
102
103fn normalize_slice(slice: SliceSpec, axis_len: usize) -> Result<(isize, usize)> {
104    if slice.step == 0 {
105        return Err(Error::InvalidSliceStep { step: slice.step });
106    }
107    if axis_len == 0 {
108        return Ok((0, 0));
109    }
110
111    let axis_len = isize::try_from(axis_len).map_err(|_| Error::IntegerOverflow)?;
112    if slice.step > 0 {
113        let start = if slice.start < 0 {
114            slice
115                .start
116                .checked_add(axis_len)
117                .ok_or(Error::IntegerOverflow)?
118        } else {
119            slice.start
120        };
121        let end = if slice.end < 0 {
122            slice
123                .end
124                .checked_add(axis_len)
125                .ok_or(Error::IntegerOverflow)?
126        } else {
127            slice.end
128        };
129        if start < 0 || start > axis_len || end < 0 || end > axis_len {
130            return Err(Error::InvalidSliceBounds {
131                start: slice.start,
132                end: slice.end,
133                axis_len: usize::try_from(axis_len).map_err(|_| Error::IntegerOverflow)?,
134            });
135        }
136        if start >= end {
137            return Ok((start, 0));
138        }
139        return Ok((start, positive_ceil_div(end - start, slice.step)?));
140    }
141
142    let start = if slice.start < 0 {
143        slice
144            .start
145            .checked_add(axis_len)
146            .ok_or(Error::IntegerOverflow)?
147    } else {
148        slice.start
149    };
150    let end = if slice.end < -1 {
151        slice
152            .end
153            .checked_add(axis_len)
154            .ok_or(Error::IntegerOverflow)?
155    } else {
156        slice.end
157    };
158    if start < 0 || start >= axis_len || end < -1 || end >= axis_len {
159        return Err(Error::InvalidSliceBounds {
160            start: slice.start,
161            end: slice.end,
162            axis_len: usize::try_from(axis_len).map_err(|_| Error::IntegerOverflow)?,
163        });
164    }
165    if start <= end {
166        return Ok((start, 0));
167    }
168    let step = slice.step.checked_neg().ok_or(Error::IntegerOverflow)?;
169    Ok((start, positive_ceil_div(start - end, step)?))
170}
171
172/// Storage-neutral tensor layout metadata.
173///
174/// # Examples
175///
176/// ```rust
177/// use tenferro_tensor_core::{Rank, TensorLayout};
178///
179/// let layout = TensorLayout::<Rank<2>>::compact([2, 3])?;
180/// assert_eq!(layout.shape(), &[2, 3]);
181/// assert_eq!(layout.strides(), &[1, 2]);
182/// # Ok::<(), tenferro_tensor_core::Error>(())
183/// ```
184#[derive(Clone, Debug, PartialEq, Eq)]
185pub struct TensorLayout<R: TensorRank = DynRank> {
186    shape: R::Shape,
187    strides: R::Strides,
188    offset: isize,
189}
190
191impl<R: TensorRank> TensorLayout<R> {
192    /// Create a compact column-major layout with zero offset.
193    ///
194    /// # Examples
195    ///
196    /// ```rust
197    /// use tenferro_tensor_core::{Rank, TensorLayout};
198    ///
199    /// let layout = TensorLayout::<Rank<2>>::compact([2, 3])?;
200    /// assert_eq!(layout.strides(), &[1, 2]);
201    /// # Ok::<(), tenferro_tensor_core::Error>(())
202    /// ```
203    pub fn compact(shape: R::Shape) -> Result<Self> {
204        let strides = R::strides_from_vec(col_major_strides(shape.as_ref())?)?;
205        Ok(Self {
206            shape,
207            strides,
208            offset: 0,
209        })
210    }
211
212    /// Create a layout from shape, strides, element offset, and backing buffer length.
213    ///
214    /// # Examples
215    ///
216    /// ```rust
217    /// use tenferro_tensor_core::{DynRank, TensorLayout};
218    ///
219    /// let layout = TensorLayout::<DynRank>::from_parts(
220    ///     vec![2, 3].into(),
221    ///     vec![1, 2].into(),
222    ///     0,
223    ///     6,
224    /// )?;
225    /// assert!(layout.is_compact_col_major()?);
226    /// # Ok::<(), tenferro_tensor_core::Error>(())
227    /// ```
228    pub fn from_parts(
229        shape: R::Shape,
230        strides: R::Strides,
231        offset: isize,
232        buffer_len: usize,
233    ) -> Result<Self> {
234        checked_logical_element_count(shape.as_ref())?;
235        validate_reachable_bounds(shape.as_ref(), strides.as_ref(), offset, buffer_len)?;
236        Ok(Self {
237            shape,
238            strides,
239            offset,
240        })
241    }
242
243    /// Return the layout shape.
244    ///
245    /// # Examples
246    ///
247    /// ```rust
248    /// use tenferro_tensor_core::{Rank, TensorLayout};
249    ///
250    /// let layout = TensorLayout::<Rank<1>>::compact([4])?;
251    /// assert_eq!(layout.shape(), &[4]);
252    /// # Ok::<(), tenferro_tensor_core::Error>(())
253    /// ```
254    pub fn shape(&self) -> &[usize] {
255        self.shape.as_ref()
256    }
257
258    /// Return the layout strides in element units.
259    ///
260    /// # Examples
261    ///
262    /// ```rust
263    /// use tenferro_tensor_core::{Rank, TensorLayout};
264    ///
265    /// let layout = TensorLayout::<Rank<2>>::compact([2, 3])?;
266    /// assert_eq!(layout.strides(), &[1, 2]);
267    /// # Ok::<(), tenferro_tensor_core::Error>(())
268    /// ```
269    pub fn strides(&self) -> &[isize] {
270        self.strides.as_ref()
271    }
272
273    /// Return the layout element offset.
274    ///
275    /// # Examples
276    ///
277    /// ```rust
278    /// use tenferro_tensor_core::{DynRank, TensorLayout};
279    ///
280    /// let layout = TensorLayout::<DynRank>::from_parts(vec![3].into(), vec![1].into(), 2, 5)?;
281    /// assert_eq!(layout.offset(), 2);
282    /// # Ok::<(), tenferro_tensor_core::Error>(())
283    /// ```
284    pub fn offset(&self) -> isize {
285        self.offset
286    }
287
288    /// Return whether the layout has compact column-major strides.
289    ///
290    /// # Examples
291    ///
292    /// ```rust
293    /// use tenferro_tensor_core::{Rank, TensorLayout};
294    ///
295    /// let layout = TensorLayout::<Rank<2>>::compact([2, 3])?;
296    /// assert!(layout.is_compact_col_major()?);
297    /// # Ok::<(), tenferro_tensor_core::Error>(())
298    /// ```
299    pub fn is_compact_col_major(&self) -> Result<bool> {
300        if self.shape().contains(&0) {
301            return Ok(true);
302        }
303
304        col_major_strides(self.shape()).map(|strides| strides.as_slice() == self.strides())
305    }
306
307    /// Validate that the layout can be used for mutable access without aliasing.
308    ///
309    /// Empty logical views are accepted. Non-empty layouts are accepted when a
310    /// conservative stride-span proof succeeds, or when exact enumeration of a
311    /// small bounded view proves that all logical elements map to distinct
312    /// physical offsets.
313    ///
314    /// # Examples
315    ///
316    /// ```rust
317    /// use tenferro_tensor_core::{DynRank, TensorLayout};
318    ///
319    /// let layout = TensorLayout::<DynRank>::from_parts(vec![3].into(), vec![-1].into(), 2, 3)?;
320    /// layout.validate_mutable_no_overlap()?;
321    /// # Ok::<(), tenferro_tensor_core::Error>(())
322    /// ```
323    pub fn validate_mutable_no_overlap(&self) -> Result<()> {
324        if self.shape().contains(&0) {
325            return Ok(());
326        }
327
328        for (&extent, &stride) in self.shape().iter().zip(self.strides()) {
329            if extent > 1 && stride == 0 {
330                return Err(Error::OverlappingMutableLayout);
331            }
332        }
333
334        let element_count = checked_product(self.shape())?;
335
336        let mut axes = self
337            .shape()
338            .iter()
339            .zip(self.strides())
340            .filter(|&(&extent, _)| extent > 1)
341            .map(|(&extent, &stride)| (extent, stride.unsigned_abs()))
342            .collect::<SmallVec<[(usize, usize); 8]>>();
343        axes.sort_by_key(|&(_, stride)| stride);
344
345        let mut span = 0usize;
346        for (extent, stride) in axes {
347            if stride <= span {
348                return self.validate_mutable_no_overlap_exact_or_reject(element_count);
349            }
350            span = span
351                .checked_add(
352                    (extent - 1)
353                        .checked_mul(stride)
354                        .ok_or(Error::IntegerOverflow)?,
355                )
356                .ok_or(Error::IntegerOverflow)?;
357        }
358
359        Ok(())
360    }
361
362    fn validate_mutable_no_overlap_exact_or_reject(&self, element_count: usize) -> Result<()> {
363        if element_count > MUTABLE_NO_OVERLAP_EXACT_ELEMENT_LIMIT {
364            return Err(Error::OverlappingMutableLayout);
365        }
366
367        let mut seen = HashSet::with_capacity(element_count);
368        let rank = self.shape().len();
369        let mut indices = vec![0usize; rank];
370
371        loop {
372            let mut physical_offset = self.offset;
373            for (&index, &stride) in indices.iter().zip(self.strides()) {
374                let index = isize::try_from(index).map_err(|_| Error::IntegerOverflow)?;
375                let delta = index.checked_mul(stride).ok_or(Error::IntegerOverflow)?;
376                physical_offset = physical_offset
377                    .checked_add(delta)
378                    .ok_or(Error::IntegerOverflow)?;
379            }
380
381            if !seen.insert(physical_offset) {
382                return Err(Error::OverlappingMutableLayout);
383            }
384
385            let mut axis = 0;
386            while axis < rank {
387                indices[axis] += 1;
388                if indices[axis] < self.shape()[axis] {
389                    break;
390                }
391                indices[axis] = 0;
392                axis += 1;
393            }
394            if axis == rank {
395                return Ok(());
396            }
397        }
398    }
399
400    /// Return a metadata-only axis permutation of this layout.
401    ///
402    /// # Examples
403    ///
404    /// ```rust
405    /// use tenferro_tensor_core::{Rank, TensorLayout};
406    ///
407    /// let layout = TensorLayout::<Rank<2>>::compact([2, 3])?;
408    /// let transposed = layout.transpose_view([1, 0])?;
409    /// assert_eq!(transposed.shape(), &[3, 2]);
410    /// assert_eq!(transposed.strides(), &[2, 1]);
411    /// # Ok::<(), tenferro_tensor_core::Error>(())
412    /// ```
413    pub fn transpose_view(&self, axes: impl AsRef<[usize]>) -> Result<Self> {
414        let axes = axes.as_ref();
415        validate_permutation(self.shape().len(), axes)?;
416        let shape = axes
417            .iter()
418            .map(|&axis| self.shape()[axis])
419            .collect::<ShapeVec>();
420        let strides = axes
421            .iter()
422            .map(|&axis| self.strides()[axis])
423            .collect::<StrideVec>();
424        Ok(Self {
425            shape: R::shape_from_vec(shape)?,
426            strides: R::strides_from_vec(strides)?,
427            offset: self.offset,
428        })
429    }
430
431    /// Return a metadata-only slice of this layout.
432    ///
433    /// # Examples
434    ///
435    /// ```rust
436    /// use tenferro_tensor_core::{Rank, SliceSpec, TensorLayout};
437    ///
438    /// let layout = TensorLayout::<Rank<1>>::compact([4])?;
439    /// let view = layout.slice_view([SliceSpec { start: 3, end: -1, step: -2 }], 4)?;
440    /// assert_eq!(view.shape(), &[2]);
441    /// assert_eq!(view.strides(), &[-2]);
442    /// # Ok::<(), tenferro_tensor_core::Error>(())
443    /// ```
444    pub fn slice_view(&self, spec: impl AsRef<[SliceSpec]>, buffer_len: usize) -> Result<Self> {
445        let spec = spec.as_ref();
446        if spec.len() != self.shape().len() {
447            return Err(Error::RankMismatch {
448                expected: self.shape().len(),
449                actual: spec.len(),
450            });
451        }
452
453        let mut shape = ShapeVec::new();
454        let mut strides = StrideVec::new();
455        let mut offset = self.offset;
456        for ((&axis_len, &stride), &slice) in self
457            .shape()
458            .iter()
459            .zip(self.strides().iter())
460            .zip(spec.iter())
461        {
462            let (start, extent) = normalize_slice(slice, axis_len)?;
463            let start_offset = start.checked_mul(stride).ok_or(Error::IntegerOverflow)?;
464            offset = offset
465                .checked_add(start_offset)
466                .ok_or(Error::IntegerOverflow)?;
467            shape.push(extent);
468            strides.push(
469                stride
470                    .checked_mul(slice.step)
471                    .ok_or(Error::IntegerOverflow)?,
472            );
473        }
474        layout_from_vecs(shape, strides, offset, buffer_len)
475    }
476
477    /// Return a metadata-only reshape of this compact column-major layout.
478    ///
479    /// # Examples
480    ///
481    /// ```rust
482    /// use tenferro_tensor_core::{Rank, TensorLayout};
483    ///
484    /// let layout = TensorLayout::<Rank<2>>::compact([2, 3])?;
485    /// let reshaped = layout.reshape_view_as::<Rank<1>>([6], 6)?;
486    /// assert_eq!(reshaped.shape(), &[6]);
487    /// assert_eq!(reshaped.strides(), &[1]);
488    /// # Ok::<(), tenferro_tensor_core::Error>(())
489    /// ```
490    pub fn reshape_view_as<R2: TensorRank>(
491        &self,
492        shape: R2::Shape,
493        buffer_len: usize,
494    ) -> Result<TensorLayout<R2>> {
495        if !self.is_compact_col_major()? {
496            return Err(Error::NonContiguousViewAsSlice);
497        }
498        let from = checked_product(self.shape())?;
499        let to = checked_product(shape.as_ref())?;
500        if from != to {
501            return Err(Error::ReshapeElementCountMismatch { from, to });
502        }
503        let strides = R2::strides_from_vec(col_major_strides(shape.as_ref())?)?;
504        TensorLayout::from_parts(shape, strides, self.offset, buffer_len)
505    }
506
507    /// Return a metadata-only explicit broadcast of this layout into a target rank.
508    ///
509    /// # Examples
510    ///
511    /// ```rust
512    /// use tenferro_tensor_core::{Rank, TensorLayout};
513    ///
514    /// let layout = TensorLayout::<Rank<1>>::compact([3])?;
515    /// let broadcast = layout.broadcast_in_dim_view::<Rank<2>>([2, 3], [1], 3)?;
516    /// assert_eq!(broadcast.shape(), &[2, 3]);
517    /// assert_eq!(broadcast.strides(), &[0, 1]);
518    /// # Ok::<(), tenferro_tensor_core::Error>(())
519    /// ```
520    pub fn broadcast_in_dim_view<R2: TensorRank>(
521        &self,
522        shape: R2::Shape,
523        broadcast_dims: impl AsRef<[usize]>,
524        buffer_len: usize,
525    ) -> Result<TensorLayout<R2>> {
526        let broadcast_dims = broadcast_dims.as_ref();
527        if broadcast_dims.len() != self.shape().len() {
528            return Err(Error::RankMismatch {
529                expected: self.shape().len(),
530                actual: broadcast_dims.len(),
531            });
532        }
533
534        let output_rank = shape.as_ref().len();
535        let mut seen = vec![false; output_rank];
536        let mut strides = StrideVec::new();
537        strides.resize(output_rank, 0);
538        for (input_axis, &output_axis) in broadcast_dims.iter().enumerate() {
539            if output_axis >= output_rank {
540                return Err(Error::AxisOutOfBounds {
541                    axis: output_axis,
542                    rank: output_rank,
543                });
544            }
545            if seen[output_axis] {
546                return Err(Error::DuplicateAxis { axis: output_axis });
547            }
548            seen[output_axis] = true;
549
550            let input_extent = self.shape()[input_axis];
551            let output_extent = shape.as_ref()[output_axis];
552            if input_extent != output_extent && input_extent != 1 {
553                return Err(Error::ShapeDataLengthMismatch {
554                    expected: input_extent,
555                    actual: output_extent,
556                });
557            }
558            if input_extent == output_extent {
559                strides[output_axis] = self.strides()[input_axis];
560            }
561        }
562
563        TensorLayout::from_parts(
564            shape,
565            R2::strides_from_vec(strides)?,
566            self.offset,
567            buffer_len,
568        )
569    }
570}
571
572#[cfg(test)]
573mod tests {
574    use super::positive_ceil_div;
575    use crate::Error;
576    use std::panic::{catch_unwind, AssertUnwindSafe};
577
578    #[test]
579    fn positive_ceil_div_rejects_invalid_preconditions_without_panicking() {
580        for (numerator, denominator) in [(-1, 1), (1, 0), (1, -1)] {
581            let result = catch_unwind(AssertUnwindSafe(|| {
582                positive_ceil_div(numerator, denominator)
583            }));
584
585            assert!(
586                result.is_ok(),
587                "invalid positive_ceil_div inputs should return Err"
588            );
589            assert!(matches!(result.unwrap(), Err(Error::IntegerOverflow)));
590        }
591    }
592}
tenferro_tensor_core/layout.rs

tenferro_tensor_core/
layout.rs
