call into vortex-compute instead of duplication

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

Author: connortsui20

vortex-array/src/arrays/primitive/compute/take/portable.rs

@@ -18,6 +18,7 @@ use num_traits::AsPrimitive;
 use vortex_buffer::Alignment;
 use vortex_buffer::Buffer;
 use vortex_buffer::BufferMut;
+use vortex_compute::take::slice::portable;
 use vortex_dtype::NativePType;
 use vortex_dtype::PType;
 use vortex_dtype::match_each_native_simd_ptype;
@@ -46,9 +47,9 @@ impl TakeImpl for TakeKernelPortableSimd {
         if array.ptype() == PType::F16 {
             // Special handling for f16 to treat as opaque u16
             let decoded = match_each_unsigned_integer_ptype!(unsigned_indices.ptype(), |C| {
-                take_portable_simd::<C, u16, SIMD_WIDTH>(
-                    unsigned_indices.as_slice(),
+                portable::take_portable_simd::<u16, C, SIMD_WIDTH>(
                     array.reinterpret_cast(PType::U16).as_slice(),
+                    unsigned_indices.as_slice(),
                 )
             });
             Ok(PrimitiveArray::new(decoded, validity)
@@ -57,9 +58,9 @@ impl TakeImpl for TakeKernelPortableSimd {
         } else {
             match_each_unsigned_integer_ptype!(unsigned_indices.ptype(), |C| {
                 match_each_native_simd_ptype!(array.ptype(), |V| {
-                    let decoded = take_portable_simd::<C, V, SIMD_WIDTH>(
-                        unsigned_indices.as_slice(),
+                    let decoded = portable::take_portable_simd::<V, C, SIMD_WIDTH>(
                         array.as_slice(),
+                        unsigned_indices.as_slice(),
                     );
                     Ok(PrimitiveArray::new(decoded, validity).into_array())
                 })
@@ -68,72 +69,6 @@ impl TakeImpl for TakeKernelPortableSimd {
     }
 }
 
-/// Takes elements from an array using SIMD indexing.
-///
-/// # Type Parameters
-/// * `C` - Index type
-/// * `V` - Value type
-/// * `LANE_COUNT` - Number of SIMD lanes to process in parallel
-///
-/// # Parameters
-/// * `indices` - Indices to gather values from
-/// * `values` - Source values to index
-///
-/// # Returns
-/// A `PrimitiveArray` containing the gathered values where each index has been replaced with
-/// the corresponding value from the source array.
-#[multiversion(targets("x86_64+avx2", "x86_64+avx", "aarch64+neon"))]
-fn take_portable_simd<I, V, const LANE_COUNT: usize>(indices: &[I], values: &[V]) -> Buffer<V>
-where
-    I: simd::SimdElement + AsPrimitive<usize>,
-    V: simd::SimdElement + NativePType,
-    simd::LaneCount<LANE_COUNT>: simd::SupportedLaneCount,
-    simd::Simd<I, LANE_COUNT>: SimdUint<Cast<usize> = simd::Simd<usize, LANE_COUNT>>,
-{
-    let indices_len = indices.len();
-
-    let mut buffer = BufferMut::<V>::with_capacity_aligned(
-        indices_len,
-        Alignment::of::<simd::Simd<V, LANE_COUNT>>(),
-    );
-
-    let buf_slice = buffer.spare_capacity_mut();
-
-    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::<V, LANE_COUNT>::default(),
-        );
-
-        unsafe {
-            selection.store_select_unchecked(
-                transmute::<&mut [MaybeUninit<V>], &mut [V]>(&mut buf_slice[offset..][..64]),
-                mask.cast(),
-            );
-        }
-    }
-
-    for idx in ((indices_len / LANE_COUNT) * LANE_COUNT)..indices_len {
-        unsafe {
-            buf_slice
-                .get_unchecked_mut(idx)
-                .write(values[indices[idx].as_()]);
-        }
-    }
-
-    unsafe {
-        buffer.set_len(indices_len);
-    }
-
-    buffer.freeze()
-}
-
 #[cfg(test)]
 mod tests {
     use super::take_portable_simd;

vortex-compute/src/take/mod.rs

@@ -4,7 +4,7 @@
 //! Take function.
 
 mod buffer;
-mod slice;
+pub mod slice;
 
 /// Function for taking based on indices (which can have different representations).
 pub trait Take<Indices: ?Sized> {

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

@@ -1,12 +1,15 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+//! Take function implementations on slices using AVX2 SIMD.
+
 #![cfg(any(target_arch = "x86_64", target_arch = "x86"))]
 
 use vortex_buffer::Buffer;
 use vortex_dtype::NativePType;
 use vortex_dtype::UnsignedPType;
 
+/// Takes the specified indices into a new [`Buffer`] using AVX2 SIMD.
 #[allow(dead_code, unused_variables, reason = "TODO(connor): Implement this")]
 #[inline]
 pub fn take_avx2<T: NativePType, I: UnsignedPType>(buffer: &[T], indices: &[I]) -> Buffer<T> {

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

@@ -1,14 +1,16 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+//! Take function implementations on slices.
+
 use vortex_buffer::Buffer;
 use vortex_dtype::NativePType;
 use vortex_dtype::UnsignedPType;
 
 use crate::take::Take;
 
-mod avx2;
-mod portable;
+pub mod avx2;
+pub mod portable;
 
 /// Specialized implementation for non-nullable indices.
 impl<T: NativePType, I: UnsignedPType> Take<[I]> for &[T] {

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

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+//! Take function implementations on slices using `portable_simd`.
+
 #![cfg(vortex_nightly)]
 
 use std::mem::MaybeUninit;
@@ -16,6 +18,7 @@ use vortex_dtype::NativePType;
 use vortex_dtype::PType;
 use vortex_dtype::UnsignedPType;
 
+/// Takes the specified indices into a new [`Buffer`] using portable SIMD.
 #[inline]
 pub fn take_portable<T, I>(buffer: &[T], indices: &[I]) -> Buffer<T>
 where
@@ -44,7 +47,7 @@ where
 ///
 /// Returns a `Buffer<T>` where each element corresponds to `values[indices[i]]`.
 #[multiversion(targets("x86_64+avx2", "x86_64+avx", "aarch64+neon"))]
-fn take_portable_simd<T, I, const LANE_COUNT: usize>(values: &[T], indices: &[I]) -> Buffer<T>
+pub fn take_portable_simd<T, I, const LANE_COUNT: usize>(values: &[T], indices: &[I]) -> Buffer<T>
 where
     T: NativePType + simd::SimdElement,
     I: UnsignedPType + simd::SimdElement,