hermes_simd_core/view/
gather.rs

1use crate::align::Alignment;
2use crate::arch::SimdArch;
3use crate::execution::ExecutionMode;
4use crate::kernel::SimdKernel;
5use crate::scalar::Scalar;
6use crate::view::{SimdError, SimdView};
7use core::mem::MaybeUninit;
8
9impl<'a, T: 'a, Arch: SimdArch + SimdKernel<T>, Align: Alignment, Mode: ExecutionMode, Ref: 'a>
10    SimdView<'a, T, Arch, Align, Mode, Ref>
11where
12    T: Scalar,
13{
14    /// Indirectly load (gather) elements from this view using indices, writing them contiguous to `out`.
15    ///
16    /// # Errors
17    /// Returns `SimdError::InsufficientOutputLength` if `out.len() < indices.len()`.
18    /// Returns `SimdError::IndexOutOfBounds` if any index in `indices` is out of bounds (negative or >= view len).
19    #[inline(always)]
20    pub fn gather(&self, indices: &[i32], out: &mut [T]) -> Result<(), SimdError> {
21        // SAFETY: an initialized `[T]` is a valid `[MaybeUninit<T>]` — the cast
22        // only widens the permitted state, never narrows it — and `T: Scalar` is
23        // `Copy`, so overwriting slots with `MaybeUninit::write` drops nothing.
24        let out_uninit = unsafe {
25            core::slice::from_raw_parts_mut(out.as_mut_ptr() as *mut MaybeUninit<T>, out.len())
26        };
27        self.gather_into_uninit(indices, out_uninit)?;
28        Ok(())
29    }
30
31    /// Gather into a possibly-uninitialized buffer, returning the initialized prefix.
32    ///
33    /// This is the single gather implementation; [`gather`](Self::gather) is the
34    /// initialized-slice wrapper over it. On `Ok` exactly the first
35    /// `indices.len()` elements of `out` are initialized (and returned as an
36    /// initialized slice); on `Err` nothing is written, since indices are fully
37    /// validated before any store. Filling an `AlignedVec`'s
38    /// [`spare_capacity_mut`](crate::vec::AlignedVec::spare_capacity_mut) through
39    /// this method and then advancing its length avoids the zero-fill that an
40    /// initialized-slice API would otherwise force.
41    ///
42    /// # Errors
43    /// Returns `SimdError::InsufficientOutputLength` if `out.len() < indices.len()`.
44    /// Returns `SimdError::IndexOutOfBounds` if any index is negative or `>=` the view length.
45    #[inline]
46    pub fn gather_into_uninit<'o>(
47        &self,
48        indices: &[i32],
49        out: &'o mut [MaybeUninit<T>],
50    ) -> Result<&'o mut [T], SimdError> {
51        let len = self.len();
52        if out.len() < indices.len() {
53            return Err(SimdError::InsufficientOutputLength);
54        }
55        let max_idx = len as i32;
56        // Validate all indices first: no element is written unless every index
57        // is in range, so the `Err` path leaves `out` untouched.
58        for &idx in indices {
59            if idx < 0 || idx >= max_idx {
60                return Err(SimdError::IndexOutOfBounds);
61            }
62        }
63
64        let indices_len = indices.len();
65        let lane_count = Arch::LANE_COUNT;
66        let simd_len = (indices_len / lane_count) * lane_count;
67        let base_ptr = self.as_slice().as_ptr();
68        let slice = self.as_slice();
69        // Derive the output pointer once and write exclusively through it: mixing
70        // it with `out[i]` slice reborrows would invalidate its provenance under
71        // Stacked Borrows.
72        let out_ptr = out.as_mut_ptr().cast::<T>();
73
74        // SAFETY: every gathered index was validated in range, so each load
75        // reads a live element of the view. `out` holds at least `indices_len`
76        // slots (checked above); `MaybeUninit<T>` shares `T`'s layout, so the
77        // vector and scalar stores below fill `[0, indices_len)` through
78        // `out_ptr` without reading any slot first. No slot is written twice.
79        unsafe {
80            for i in (0..simd_len).step_by(lane_count) {
81                let idx_slice = &indices[i..i + lane_count];
82                let idx_vec = crate::sparse::spmv::build_index_vector::<T, Arch>(idx_slice);
83                let v = Arch::gather(base_ptr, idx_vec);
84                Arch::store_unaligned(out_ptr.add(i), v);
85            }
86            for i in simd_len..indices_len {
87                core::ptr::write(out_ptr.add(i), slice[indices[i] as usize]);
88            }
89            // Every element of `[0, indices_len)` is now initialized.
90            Ok(core::slice::from_raw_parts_mut(out_ptr, indices_len))
91        }
92    }
93}
hermes_simd_core/view/gather.rs

hermes_simd_core/view/
gather.rs
