add take for slices, missing avx2 implementation

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

Author: connortsui20

Cargo.lock

@@ -412,6 +412,27 @@ impl<T> Buffer<T> {
         }
     }
 
+    /// Cast a `Buffer<T>` into a `Buffer<U>`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the type `U` does not have the same size and alignment as `T`.
+    pub fn cast_into<U>(self) -> Buffer<U> {
+        assert_eq!(size_of::<T>(), size_of::<U>(), "Buffer type size mismatch");
+        assert_eq!(
+            align_of::<T>(),
+            align_of::<U>(),
+            "Buffer type alignment mismatch"
+        );
+
+        Buffer {
+            bytes: self.bytes,
+            length: self.length,
+            alignment: self.alignment,
+            _marker: PhantomData,
+        }
+    }
+
     /// Try to convert self into `BufferMut<T>` if there is only a single strong reference.
     pub fn try_into_mut(self) -> Result<BufferMut<T>, Self> {
         self.bytes

vortex-buffer/src/buffer.rs

@@ -29,6 +29,7 @@ vortex-vector = { workspace = true }
 arrow-array = { workspace = true, optional = true }
 arrow-buffer = { workspace = true, optional = true }
 arrow-schema = { workspace = true, optional = true }
+multiversion = { workspace = true }
 num-traits = { workspace = true }
 
 [features]

vortex-compute/Cargo.toml

@@ -0,0 +1,16 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_buffer::Buffer;
+use vortex_dtype::NativePType;
+use vortex_dtype::UnsignedPType;
+
+use crate::take::Take;
+
+impl<T: NativePType, I: UnsignedPType> Take<[I]> for &Buffer<T> {
+    type Output = Buffer<T>;
+
+    fn take(self, indices: &[I]) -> Buffer<T> {
+        self.as_slice().take(indices)
+    }
+}

vortex-compute/src/take/buffer.rs

@@ -3,6 +3,9 @@
 
 //! Take function.
 
+mod buffer;
+mod slice;
+
 /// Function for taking based on indices (which can have different representations).
 pub trait Take<Indices: ?Sized> {
     /// The result type after performing the operation.

vortex-compute/src/take/mod.rs

@@ -0,0 +1,17 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#![cfg(any(target_arch = "x86_64", target_arch = "x86"))]
+
+use vortex_buffer::Buffer;
+use vortex_dtype::NativePType;
+use vortex_dtype::UnsignedPType;
+
+#[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> {
+    todo!(
+        "TODO(connor): Migrate the implementation in \
+            vortex-array/src/arrays/primitive/compute/take/avx2.rs"
+    )
+}

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

@@ -0,0 +1,41 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_buffer::Buffer;
+use vortex_dtype::NativePType;
+use vortex_dtype::UnsignedPType;
+
+use crate::take::Take;
+
+mod avx2;
+mod portable;
+
+/// Specialized implementation for non-nullable indices.
+impl<T: NativePType, I: UnsignedPType> Take<[I]> for &[T] {
+    type Output = Buffer<T>;
+
+    fn take(self, indices: &[I]) -> Buffer<T> {
+        #[cfg(vortex_nightly)]
+        {
+            return portable::take_portable(self, indices);
+        }
+
+        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+        {
+            if is_x86_feature_detected!("avx2") {
+                return avx2::take_avx2(self, indices);
+            }
+        }
+
+        take_scalar(self, indices)
+    }
+}
+
+#[allow(
+    unused,
+    reason = "Compiler may see this as unused based on enabled features"
+)]
+#[inline]
+fn take_scalar<T: NativePType, I: UnsignedPType>(buffer: &[T], indices: &[I]) -> Buffer<T> {
+    indices.iter().map(|idx| buffer[idx.as_()]).collect()
+}

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

@@ -0,0 +1,110 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#![cfg(vortex_nightly)]
+
+use std::mem::MaybeUninit;
+use std::mem::transmute;
+use std::simd;
+use std::simd::num::SimdUint;
+
+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;
+
+#[inline]
+pub fn take_portable<T, I>(buffer: &[T], indices: &[I]) -> Buffer<T>
+where
+    T: NativePType + simd::SimdElement,
+    I: UnsignedPType + simd::SimdElement,
+{
+    if T::PTYPE == PType::F16 {
+        // Since Rust does not actually support 16-bit floats, we first reinterpret the data as
+        // `u16` integers.
+        let u16_slice: &[u16] =
+            unsafe { std::slice::from_raw_parts(buffer.as_ptr() as *const u16, buffer.len()) };
+
+        let taken_u16 = take_portable_simd::<u16, I, SIMD_WIDTH>(u16_slice, indices);
+        let taken_f16 = taken_u16.cast_into::<T>();
+
+        taken_f16
+    } else {
+        take_portable_simd::<T, I, SIMD_WIDTH>(buffer, indices)
+    }
+}
+
+/// Takes elements from an array using SIMD indexing.
+///
+/// Performs a gather operation that takes values at specified indices and returns them in a new
+/// buffer. Uses SIMD instructions to process `LANE_COUNT` indices in parallel.
+///
+/// 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>
+where
+    T: NativePType + 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>>,
+{
+    let indices_len = indices.len();
+
+    let mut buffer = BufferMut::<T>::with_capacity_aligned(
+        indices_len,
+        Alignment::of::<simd::Simd<T, 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::<T, LANE_COUNT>::default(),
+        );
+
+        unsafe {
+            selection.store_select_unchecked(
+                transmute::<&mut [MaybeUninit<T>], &mut [T]>(&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;
+
+    #[test]
+    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]);
+    }
+}

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

@@ -97,6 +97,42 @@ macro_rules! match_each_integer_pvector {
     }};
 }
 
+/// Matches on all unsigned type variants of [`PrimitiveVector`] and executes the same code for each
+/// of the unsigned variant branches.
+///
+/// This macro eliminates repetitive match statements when implementing operations that need to work
+/// uniformly across all unsigned type variants (`U8`, `U16`, `U32`, `U64`).
+///
+/// See [`match_each_pvector`] for similar usage.
+///
+/// [`PrimitiveVector`]: crate::primitive::PrimitiveVector
+///
+/// # Panics
+///
+/// Panics if the vector passed in to the macro is not an unsigned vector variant.
+#[macro_export]
+macro_rules! match_each_unsigned_pvector {
+    ($self:expr, | $vec:ident | $body:block) => {{
+        match $self {
+            $crate::primitive::PrimitiveVector::U8($vec) => $body,
+            $crate::primitive::PrimitiveVector::U16($vec) => $body,
+            $crate::primitive::PrimitiveVector::U32($vec) => $body,
+            $crate::primitive::PrimitiveVector::U64($vec) => $body,
+            $crate::primitive::PrimitiveVector::I8(_)
+            | $crate::primitive::PrimitiveVector::I16(_)
+            | $crate::primitive::PrimitiveVector::I32(_)
+            | $crate::primitive::PrimitiveVector::I64(_)
+            | $crate::primitive::PrimitiveVector::F16(_)
+            | $crate::primitive::PrimitiveVector::F32(_)
+            | $crate::primitive::PrimitiveVector::F64(_) => {
+                ::vortex_error::vortex_panic!(
+                    "Tried to match a non-unsigned vector in an unsigned match statement"
+                )
+            }
+        }
+    }};
+}
+
 /// Matches on all primitive type variants of [`PrimitiveVectorMut`] and executes the same code
 /// for each variant branch.
 ///
@@ -184,3 +220,39 @@ macro_rules! match_each_integer_pvector_mut {
         }
     }};
 }
+
+/// Matches on all unsigned type variants of [`PrimitiveVectorMut`] and executes the same code for
+/// each of the unsigned variant branches.
+///
+/// This macro eliminates repetitive match statements when implementing operations that need to work
+/// uniformly across all unsigned type variants (`U8`, `U16`, `U32`, `U64`).
+///
+/// See [`match_each_pvector_mut`] for similar usage.
+///
+/// [`PrimitiveVectorMut`]: crate::primitive::PrimitiveVectorMut
+///
+/// # Panics
+///
+/// Panics if the vector passed in to the macro is not an unsigned vector variant.
+#[macro_export]
+macro_rules! match_each_unsigned_pvector_mut {
+    ($self:expr, | $vec:ident | $body:block) => {{
+        match $self {
+            $crate::primitive::PrimitiveVectorMut::U8($vec) => $body,
+            $crate::primitive::PrimitiveVectorMut::U16($vec) => $body,
+            $crate::primitive::PrimitiveVectorMut::U32($vec) => $body,
+            $crate::primitive::PrimitiveVectorMut::U64($vec) => $body,
+            $crate::primitive::PrimitiveVectorMut::I8(_)
+            | $crate::primitive::PrimitiveVectorMut::I16(_)
+            | $crate::primitive::PrimitiveVectorMut::I32(_)
+            | $crate::primitive::PrimitiveVectorMut::I64(_)
+            | $crate::primitive::PrimitiveVectorMut::F16(_)
+            | $crate::primitive::PrimitiveVectorMut::F32(_)
+            | $crate::primitive::PrimitiveVectorMut::F64(_) => {
+                ::vortex_error::vortex_panic!(
+                    "Tried to match a non-unsigned mutable vector in an unsigned match statement"
+                )
+            }
+        }
+    }};
+}