add OOB check + safety comments

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

vortex-compute/src/take/slice/portable.rs

@@ -7,8 +7,8 @@
 
 use std::mem::MaybeUninit;
 use std::mem::size_of;
-use std::mem::transmute;
 use std::simd;
+use std::simd::cmp::SimdPartialOrd;
 use std::simd::num::SimdUint;
 
 use multiversion::multiversion;
@@ -83,6 +83,10 @@ fn take_with_indices<T: Copy + Default + simd::SimdElement, I: UnsignedPType>(
 /// buffer. Uses SIMD instructions to process `LANE_COUNT` indices in parallel.
 ///
 /// Returns a `Buffer<T>` where each element corresponds to `values[indices[i]]`.
+///
+/// # Panics
+///
+/// Panics if any index is out of bounds for `values`.
 #[multiversion(targets("x86_64+avx2", "x86_64+avx", "aarch64+neon"))]
 pub fn take_portable_simd<T, I, const LANE_COUNT: usize>(values: &[T], indices: &[I]) -> Buffer<T>
 where
@@ -100,37 +104,67 @@ where
 
     let buf_slice = buffer.spare_capacity_mut();
 
+    // Set up a vector that we can SIMD compare against for out-of-bounds indices.
+    let len_vec = simd::Simd::<usize, LANE_COUNT>::splat(values.len());
+    let mut all_valid = simd::Mask::<isize, LANE_COUNT>::splat(true);
+
     for chunk_idx in 0..(indices_len / LANE_COUNT) {
         let offset = chunk_idx * LANE_COUNT;
-        let mask = simd::Mask::from_bitmask(u64::MAX);
         let codes_chunk = simd::Simd::<I, LANE_COUNT>::from_slice(&indices[offset..]);
-
-        let selection = simd::Simd::gather_select(
-            values,
-            mask,
-            codes_chunk.cast::<usize>(),
-            simd::Simd::<T, LANE_COUNT>::default(),
-        );
-
+        let codes_usize = codes_chunk.cast::<usize>();
+
+        // Accumulate validity and use as gather mask. An out-of-bounds index will turn a bit off.
+        all_valid &= codes_usize.simd_lt(len_vec);
+
+        // SAFETY: We use `all_valid` to mask the gather, preventing OOB memory access. If any
+        // index is OOB, `all_valid` will have those bits turned off, masking out the invalid
+        // indices.
+        // Note that this may also mask out valid indices in subsequent iterations. This is fine
+        // because we will panic after the loop if **any** index was OOB, so we do not care if the
+        // resulting gathered data is correct or not.
+        let selection = unsafe {
+            simd::Simd::gather_select_unchecked(
+                values,
+                all_valid,
+                codes_usize,
+                simd::Simd::<T, LANE_COUNT>::default(),
+            )
+        };
+
+        // SAFETY: `MaybeUninit<T>` has the same layout as `T`, and we are about to initialize these
+        // elements with the store.
+        let uninit = unsafe {
+            std::mem::transmute::<&mut [MaybeUninit<T>], &mut [T]>(
+                &mut buf_slice[offset..][..LANE_COUNT],
+            )
+        };
+
+        // SAFETY: The slice `buf_slice[offset..][..LANE_COUNT]` is guaranteed to have exactly
+        // `LANE_COUNT` elements since `offset` is a multiple of `LANE_COUNT` and we only iterate
+        // while `offset + LANE_COUNT <= indices_len`.
         unsafe {
-            selection.store_select_unchecked(
-                transmute::<&mut [MaybeUninit<T>], &mut [T]>(&mut buf_slice[offset..][..64]),
-                mask.cast(),
-            );
+            selection.store_select_unchecked(uninit, simd::Mask::splat(true));
         }
     }
 
+    // Check accumulated validity after hot loop. If there are any 0's, then there was an
+    // out-of-bounds index.
+    assert!(all_valid.all(), "index out of bounds in SIMD take");
+
+    // Fall back to scalar iteration for the remainder.
     for idx in ((indices_len / LANE_COUNT) * LANE_COUNT)..indices_len {
+        // SAFETY: `idx` is in bounds for `buf_slice` since `idx < indices_len == buf_slice.len()`.
+        // Note that the `values[...]` access is already bounds-checked and will panic if OOB.
         unsafe {
             buf_slice
                 .get_unchecked_mut(idx)
                 .write(values[indices[idx].as_()]);
         }
     }
 
-    unsafe {
-        buffer.set_len(indices_len);
-    }
+    // SAFETY: All elements have been initialized: the SIMD loop handles `0..chunks * LANE_COUNT`
+    // and the scalar loop handles the remainder up to `indices_len`.
+    unsafe { buffer.set_len(indices_len) };
 
     buffer.freeze()
 }
@@ -141,12 +175,12 @@ mod tests {
     use super::take_portable_simd;
 
     #[test]
+    #[should_panic(expected = "index out of bounds")]
     fn test_take_out_of_bounds() {
         let indices = vec![2_000_000u32; 64];
         let values = vec![1i32];
 
-        let result = take_portable_simd::<i32, u32, 64>(&values, &indices);
-        assert_eq!(result.as_slice(), [0i32; 64]);
+        drop(take_portable_simd::<i32, u32, 64>(&values, &indices));
     }
 
     /// Tests SIMD gather with a mix of sequential, strided, and repeated indices. This exercises