bring portable simd take back

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

Author: connortsui20

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

@@ -16,14 +16,15 @@ impl<T: Copy, I: UnsignedPType> Take<[I]> for &[T] {
     type Output = Buffer<T>;
 
     fn take(self, indices: &[I]) -> Buffer<T> {
-        // TODO(connor): Make the SIMD implementations bound by `Copy` instead of `NativePType`.
-        /*
-
         #[cfg(vortex_nightly)]
         {
             return portable::take_portable(self, indices);
         }
 
+        // TODO(connor): Make the SIMD implementations bound by `Copy` instead of `NativePType`.
+
+        /*
+
         #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
         {
             if is_x86_feature_detected!("avx2") {
@@ -44,6 +45,6 @@ impl<T: Copy, I: UnsignedPType> Take<[I]> for &[T] {
     reason = "Compiler may see this as unused based on enabled features"
 )]
 #[inline]
-fn take_scalar<T: Copy, I: UnsignedPType>(buffer: &[T], indices: &[I]) -> Buffer<T> {
+pub(crate) fn take_scalar<T: Copy, I: UnsignedPType>(buffer: &[T], indices: &[I]) -> Buffer<T> {
     indices.iter().map(|idx| buffer[idx.as_()]).collect()
 }

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

@@ -15,10 +15,7 @@ use multiversion::multiversion;
 use vortex_buffer::Alignment;
 use vortex_buffer::Buffer;
 use vortex_buffer::BufferMut;
-use vortex_dtype::NativePType;
-use vortex_dtype::PType;
 use vortex_dtype::UnsignedPType;
-use vortex_dtype::match_each_native_simd_ptype;
 use vortex_dtype::match_each_unsigned_integer_ptype;
 
 /// SIMD types larger than the SIMD register size are beneficial for
@@ -27,38 +24,49 @@ pub const SIMD_WIDTH: usize = 64;
 
 /// Takes the specified indices into a new [`Buffer`] using portable SIMD.
 ///
-/// This function handles the type matching required to satisfy `SimdElement` bounds.
-/// For `f16` values, it reinterprets them as `u16` since `f16` doesn't implement `SimdElement`.
+/// This function handles the type matching required to satisfy `SimdElement` bounds by casting
+/// to unsigned integers of the same size. Falls back to scalar implementation for unsupported
+/// type sizes.
 #[inline]
-pub fn take_portable<T: NativePType, I: UnsignedPType>(buffer: &[T], indices: &[I]) -> Buffer<T> {
-    if T::PTYPE == PType::F16 {
-        assert_eq!(size_of::<half::f16>(), size_of::<T>());
-
-        // Since Rust does not actually support 16-bit floats, we first reinterpret the data as
-        // `u16` integers.
-        // SAFETY: We know that f16 has the same bit pattern as u16, so this transmute is fine to
-        // make.
-        let u16_slice: &[u16] =
-            unsafe { std::slice::from_raw_parts(buffer.as_ptr() as *const u16, buffer.len()) };
-        return take_with_indices(u16_slice, indices).cast_into::<T>();
+pub fn take_portable<T: Copy, I: UnsignedPType>(buffer: &[T], indices: &[I]) -> Buffer<T> {
+    // SIMD gather operations only care about bit patterns, not semantic type. We cast to unsigned
+    // integers which implement `SimdElement` and then cast back.
+    //
+    // SAFETY: The pointer casts below are safe because:
+    // - `T` and the target type have the same size (matched by `size_of::<T>()`).
+    // - The alignment of unsigned integers is always <= their size, and `buffer` came from a valid
+    //   `&[T]` which guarantees proper alignment for types of the same size.
+    match size_of::<T>() {
+        1 => {
+            let buffer: &[u8] =
+                unsafe { std::slice::from_raw_parts(buffer.as_ptr() as *const u8, buffer.len()) };
+            take_with_indices(buffer, indices).cast_into::<T>()
+        }
+        2 => {
+            let buffer: &[u16] =
+                unsafe { std::slice::from_raw_parts(buffer.as_ptr() as *const u16, buffer.len()) };
+            take_with_indices(buffer, indices).cast_into::<T>()
+        }
+        4 => {
+            let buffer: &[u32] =
+                unsafe { std::slice::from_raw_parts(buffer.as_ptr() as *const u32, buffer.len()) };
+            take_with_indices(buffer, indices).cast_into::<T>()
+        }
+        8 => {
+            let buffer: &[u64] =
+                unsafe { std::slice::from_raw_parts(buffer.as_ptr() as *const u64, buffer.len()) };
+            take_with_indices(buffer, indices).cast_into::<T>()
+        }
+        // Fall back to scalar implementation for unsupported type sizes.
+        _ => super::take_scalar(buffer, indices),
     }
-
-    match_each_native_simd_ptype!(T::PTYPE, |TC| {
-        assert_eq!(size_of::<TC>(), size_of::<T>());
-
-        // SAFETY: This is essentially a no-op that tricks the compiler into adding the
-        // `simd::SimdElement` bound we need to call `take_with_indices`.
-        let buffer: &[TC] =
-            unsafe { std::slice::from_raw_parts(buffer.as_ptr() as *const TC, buffer.len()) };
-        take_with_indices(buffer, indices).cast_into::<T>()
-    })
 }
 
 /// Helper that matches on index type and calls `take_portable_simd`.
 ///
 /// We separate this code out from above to add the [`simd::SimdElement`] constraint.
 #[inline]
-fn take_with_indices<T: NativePType + simd::SimdElement, I: UnsignedPType>(
+fn take_with_indices<T: Copy + Default + simd::SimdElement, I: UnsignedPType>(
     buffer: &[T],
     indices: &[I],
 ) -> Buffer<T> {
@@ -78,7 +86,7 @@ fn take_with_indices<T: NativePType + simd::SimdElement, I: UnsignedPType>(
 #[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
-    T: NativePType + simd::SimdElement,
+    T: Copy + Default + simd::SimdElement,
     I: UnsignedPType + simd::SimdElement,
     simd::LaneCount<LANE_COUNT>: simd::SupportedLaneCount,
     simd::Simd<I, LANE_COUNT>: SimdUint<Cast<usize> = simd::Simd<usize, LANE_COUNT>>,
@@ -128,6 +136,7 @@ where
 }
 
 #[cfg(test)]
+#[allow(clippy::cast_possible_truncation)]
 mod tests {
     use super::take_portable_simd;
 
@@ -159,7 +168,7 @@ mod tests {
         // Strided by 4: 0, 4, 8, ..., 252.
         indices.extend((0u32..64).map(|i| i * 4));
         // Repeated: index 42 repeated 32 times.
-        indices.extend(std::iter::repeat(42u32).take(32));
+        indices.extend(std::iter::repeat_n(42u32, 32));
         // Reverse: 255, 254, ..., 216.
         indices.extend((216u32..256).rev());
 

vortex-compute/src/take/vector/tests.rs

@@ -120,7 +120,6 @@ fn test_null_vector_take() {
     assert!(result.validity().all_false());
 }
 
-#[ignore = "TODO(connor): Implement `DecimalVector::take`."]
 #[test]
 fn test_dvector_take() {
     use vortex_buffer::buffer;
@@ -154,7 +153,6 @@ fn test_dvector_take() {
     assert_eq!(validity, vec![true, false, true, false]);
 }
 
-#[ignore = "TODO(connor): Implement `DecimalVector::take`."]
 #[test]
 fn test_decimal_vector_take() {
     use vortex_buffer::buffer;