vortex_compute/lane_kernels/
source.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Read-only lane source — the [`IndexedSource`] trait and the [`LaneZip`] adapter.
5
6/// A length-known source supporting unchecked indexed reads.
7///
8/// Implemented for `&[T]` (with `T: Copy`) and for [`LaneZip`] over two `IndexedSource`s.
9/// The kernels in this crate require this trait instead of `Iterator` so that lane
10/// reads carry no inter-iteration data dependency — the autovectorizer treats each
11/// lane independently.
12pub trait IndexedSource {
13    /// The per-lane item type. Must be `Copy` so the kernels can pass it through
14    /// the closure by value without extra moves.
15    type Item: Copy;
16    /// Logical lane count.
17    fn len(&self) -> usize;
18    /// Returns true when there are no lanes.
19    fn is_empty(&self) -> bool {
20        self.len() == 0
21    }
22    /// Read the lane at `i` without bounds checking.
23    ///
24    /// # Safety
25    ///
26    /// `i` must be strictly less than `self.len()`.
27    unsafe fn get_unchecked(&self, i: usize) -> Self::Item;
28}
29
30impl<T: Copy> IndexedSource for &[T] {
31    type Item = T;
32    #[inline]
33    fn len(&self) -> usize {
34        <[T]>::len(self)
35    }
36    #[inline]
37    unsafe fn get_unchecked(&self, i: usize) -> T {
38        // SAFETY: caller guarantees i < self.len().
39        unsafe { *<[T]>::get_unchecked(self, i) }
40    }
41}
42
43impl<T: Copy> IndexedSource for &mut [T] {
44    type Item = T;
45    #[inline]
46    fn len(&self) -> usize {
47        <[T]>::len(self)
48    }
49    #[inline]
50    unsafe fn get_unchecked(&self, i: usize) -> T {
51        // SAFETY: caller guarantees i < self.len().
52        unsafe { *<[T]>::get_unchecked(self, i) }
53    }
54}
55
56/// Pair of two [`IndexedSource`]s of equal length. Yields `(A::Item, B::Item)` per lane.
57///
58/// Use this to drive a binary kernel from two columns. Length equality is enforced
59/// at construction.
60pub struct LaneZip<A, B>(pub A, pub B);
61
62impl<A: IndexedSource, B: IndexedSource> LaneZip<A, B> {
63    /// Build a `LaneZip` from two equal-length sources.
64    ///
65    /// # Panics
66    ///
67    /// Panics if the two operands have different lengths.
68    pub fn new(a: A, b: B) -> Self {
69        assert_eq!(
70            a.len(),
71            b.len(),
72            "LaneZip operands must have the same length"
73        );
74        Self(a, b)
75    }
76}
77
78impl<A: IndexedSource, B: IndexedSource> IndexedSource for LaneZip<A, B> {
79    type Item = (A::Item, B::Item);
80    #[inline]
81    fn len(&self) -> usize {
82        debug_assert_eq!(self.0.len(), self.1.len());
83        self.0.len()
84    }
85    #[inline]
86    unsafe fn get_unchecked(&self, i: usize) -> (A::Item, B::Item) {
87        // SAFETY: caller guarantees i < self.len(); `new` enforces matching lengths.
88        unsafe { (self.0.get_unchecked(i), self.1.get_unchecked(i)) }
89    }
90}
vortex_compute/lane_kernels/source.rs

vortex_compute/lane_kernels/
source.rs
