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 passed through the kernel by value.
14    type Item;
15    /// Logical lane count.
16    fn len(&self) -> usize;
17    /// Returns true when there are no lanes.
18    fn is_empty(&self) -> bool {
19        self.len() == 0
20    }
21    /// Read the lane at `i` without bounds checking.
22    ///
23    /// # Safety
24    ///
25    /// `i` must be strictly less than `self.len()`.
26    unsafe fn get_unchecked(&self, i: usize) -> Self::Item;
27}
28
29impl<T: Copy> IndexedSource for &[T] {
30    type Item = T;
31    #[inline]
32    fn len(&self) -> usize {
33        <[T]>::len(self)
34    }
35    #[inline]
36    unsafe fn get_unchecked(&self, i: usize) -> T {
37        // SAFETY: caller guarantees i < self.len().
38        unsafe { *<[T]>::get_unchecked(self, i) }
39    }
40}
41
42impl<T: Copy> IndexedSource for &mut [T] {
43    type Item = T;
44    #[inline]
45    fn len(&self) -> usize {
46        <[T]>::len(self)
47    }
48    #[inline]
49    unsafe fn get_unchecked(&self, i: usize) -> T {
50        // SAFETY: caller guarantees i < self.len().
51        unsafe { *<[T]>::get_unchecked(self, i) }
52    }
53}
54
55/// Pair of two [`IndexedSource`]s of equal length. Yields `(A::Item, B::Item)` per lane.
56///
57/// Use this to drive a binary kernel from two columns. Length equality is enforced at
58/// construction, and the private fields prevent callers from bypassing that check.
59#[derive(Clone, Copy)]
60pub struct LaneZip<A, B>(A, 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        self.0.len()
83    }
84    #[inline]
85    unsafe fn get_unchecked(&self, i: usize) -> (A::Item, B::Item) {
86        // SAFETY: caller guarantees i < self.len(); `new` enforces matching lengths.
87        unsafe { (self.0.get_unchecked(i), self.1.get_unchecked(i)) }
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::LaneZip;
94
95    #[test]
96    #[should_panic(expected = "LaneZip operands must have the same length")]
97    fn rejects_mismatched_lengths() {
98        _ = LaneZip::new(&[1_u8][..], &[2_u8, 3][..]);
99    }
100}
vortex_compute/lane_kernels/source.rs

vortex_compute/lane_kernels/
source.rs
