Add divan benchmark comparing scalar vs AVX2 vs Mojo take kernel

Adds `take_primitive_simd` benchmark that calls all three gather
implementations through identical `fn(&[T], &[u32]) -> Buffer<T>`
signatures on raw buffers. No Vortex Array overhead.

Results on AVX2 (65K values, random u32 indices, median):

  u32, n=100K: scalar=66.9µs, avx2=46.0µs (1.45x), mojo=44.0µs (1.52x)
  u64, n=100K: scalar=67.1µs, avx2=55.6µs (1.21x), mojo=55.4µs (1.21x)

Signed-off-by: Claude <noreply@anthropic.com>

https://claude.ai/code/session_01EVcJZP4ZmfvWRRg2CsgvST

Author: claude

vortex-array/Cargo.toml

@@ -165,6 +165,10 @@ harness = false
 name = "take_primitive"
 harness = false
 
+[[bench]]
+name = "take_primitive_simd"
+harness = false
+
 [[bench]]
 name = "take_struct"
 harness = false

vortex-array/benches/take_primitive_simd.rs

@@ -0,0 +1,80 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Raw buffer-level benchmarks for the primitive take (gather) kernel.
+//!
+//! Compares scalar, AVX2, and Mojo SIMD gather. All three are called through
+//! the same `fn(&[T], &[u32]) -> Buffer<T>` Rust signature on raw slices.
+//!
+//! Run with: `cargo bench -p vortex-array --bench take_primitive_simd`
+
+#![allow(clippy::cast_possible_truncation)]
+#![allow(clippy::unwrap_used)]
+
+use divan::Bencher;
+use rand::distr::Uniform;
+use rand::prelude::*;
+use vortex_array::arrays::primitive::{bench_take_avx2, bench_take_mojo, bench_take_scalar};
+
+fn main() {
+    divan::main();
+}
+
+const NUM_INDICES: &[usize] = &[1_000, 10_000, 100_000];
+const NUM_VALUES: usize = 65_536;
+
+fn make_u32_indices(num_indices: usize) -> Vec<u32> {
+    let rng = StdRng::seed_from_u64(42);
+    let range = Uniform::new(0u32, NUM_VALUES as u32).unwrap();
+    rng.sample_iter(range).take(num_indices).collect()
+}
+
+// ---------------------------------------------------------------------------
+// u32 values
+// ---------------------------------------------------------------------------
+
+#[divan::bench(args = NUM_INDICES, sample_count = 10_000)]
+fn gather_u32_scalar(bencher: Bencher, n: usize) {
+    let values: Vec<u32> = (0..NUM_VALUES as u32).collect();
+    let indices = make_u32_indices(n);
+    bencher.bench(|| divan::black_box(bench_take_scalar(&values, &indices)));
+}
+
+#[divan::bench(args = NUM_INDICES, sample_count = 10_000)]
+fn gather_u32_avx2(bencher: Bencher, n: usize) {
+    let values: Vec<u32> = (0..NUM_VALUES as u32).collect();
+    let indices = make_u32_indices(n);
+    bencher.bench(|| divan::black_box(bench_take_avx2(&values, &indices)));
+}
+
+#[divan::bench(args = NUM_INDICES, sample_count = 10_000)]
+fn gather_u32_mojo(bencher: Bencher, n: usize) {
+    let values: Vec<u32> = (0..NUM_VALUES as u32).collect();
+    let indices = make_u32_indices(n);
+    bencher.bench(|| divan::black_box(bench_take_mojo(&values, &indices)));
+}
+
+// ---------------------------------------------------------------------------
+// u64 values
+// ---------------------------------------------------------------------------
+
+#[divan::bench(args = NUM_INDICES, sample_count = 10_000)]
+fn gather_u64_scalar(bencher: Bencher, n: usize) {
+    let values: Vec<u64> = (0..NUM_VALUES as u64).map(|i| i * 100).collect();
+    let indices = make_u32_indices(n);
+    bencher.bench(|| divan::black_box(bench_take_scalar(&values, &indices)));
+}
+
+#[divan::bench(args = NUM_INDICES, sample_count = 10_000)]
+fn gather_u64_avx2(bencher: Bencher, n: usize) {
+    let values: Vec<u64> = (0..NUM_VALUES as u64).map(|i| i * 100).collect();
+    let indices = make_u32_indices(n);
+    bencher.bench(|| divan::black_box(bench_take_avx2(&values, &indices)));
+}
+
+#[divan::bench(args = NUM_INDICES, sample_count = 10_000)]
+fn gather_u64_mojo(bencher: Bencher, n: usize) {
+    let values: Vec<u64> = (0..NUM_VALUES as u64).map(|i| i * 100).collect();
+    let indices = make_u32_indices(n);
+    bencher.bench(|| divan::black_box(bench_take_mojo(&values, &indices)));
+}

vortex-array/src/arrays/primitive/compute/mod.rs

@@ -7,7 +7,7 @@ mod fill_null;
 mod mask;
 pub(crate) mod rules;
 mod slice;
-mod take;
+pub(crate) mod take;
 
 #[cfg(test)]
 mod tests {

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

@@ -121,7 +121,10 @@ where
 /// The caller must ensure the `avx2` feature is enabled.
 #[target_feature(enable = "avx2")]
 #[doc(hidden)]
-unsafe fn take_avx2<V: NativePType, I: UnsignedPType>(buffer: &[V], indices: &[I]) -> Buffer<V> {
+pub(super) unsafe fn take_avx2<V: NativePType, I: UnsignedPType>(
+    buffer: &[V],
+    indices: &[I],
+) -> Buffer<V> {
     macro_rules! dispatch_avx2 {
         ($indices:ty, $values:ty) => {
             { let result = dispatch_avx2!($indices, $values, cast: $values); result }

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

@@ -144,6 +144,53 @@ fn take_primitive_scalar<T: NativePType, I: IntegerPType>(
     result.freeze()
 }
 
+// ---------------------------------------------------------------------------
+// Benchmark-visible helpers — expose the raw scalar and Mojo gather kernels
+// with identical signatures so benchmarks can compare them directly.
+// ---------------------------------------------------------------------------
+
+/// Scalar gather: `result[i] = buffer[indices[i]]`. No SIMD.
+#[doc(hidden)]
+pub fn bench_take_scalar<T: NativePType, I: IntegerPType>(
+    buffer: &[T],
+    indices: &[I],
+) -> Buffer<T> {
+    take_primitive_scalar(buffer, indices)
+}
+
+/// AVX2 gather via hand-written intrinsics. Falls back to scalar on non-x86 or when AVX2
+/// is unavailable at runtime.
+#[doc(hidden)]
+pub fn bench_take_avx2<T: NativePType, I: crate::dtype::UnsignedPType>(
+    buffer: &[T],
+    indices: &[I],
+) -> Buffer<T> {
+    #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
+    {
+        if is_x86_feature_detected!("avx2") {
+            // SAFETY: We just checked AVX2 is available.
+            return unsafe { avx2::take_avx2(buffer, indices) };
+        }
+    }
+    take_primitive_scalar(buffer, indices)
+}
+
+/// SIMD gather via the Mojo AOT kernel. Falls back to scalar when Mojo is not available.
+#[doc(hidden)]
+pub fn bench_take_mojo<T: NativePType, I: crate::dtype::UnsignedPType>(
+    buffer: &[T],
+    indices: &[I],
+) -> Buffer<T> {
+    #[cfg(vortex_mojo)]
+    {
+        mojo::take_mojo(buffer, indices)
+    }
+    #[cfg(not(vortex_mojo))]
+    {
+        take_primitive_scalar(buffer, indices)
+    }
+}
+
 #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
 #[cfg(test)]
 mod test {

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

@@ -80,7 +80,10 @@ impl TakeImpl for TakeKernelMojo {
 }
 
 /// Dispatch to the appropriate Mojo kernel based on value byte width and index type.
-fn take_mojo<V: NativePType, I: UnsignedPType>(values: &[V], indices: &[I]) -> Buffer<V> {
+pub(super) fn take_mojo<V: NativePType, I: UnsignedPType>(
+    values: &[V],
+    indices: &[I],
+) -> Buffer<V> {
     let len = indices.len();
     let mut buffer = BufferMut::<V>::with_capacity(len);
 

vortex-array/src/arrays/primitive/mod.rs

@@ -13,6 +13,12 @@ pub(crate) mod compute;
 
 mod vtable;
 pub use compute::rules::PrimitiveMaskedValidityRule;
+#[doc(hidden)]
+pub use compute::take::bench_take_avx2;
+#[doc(hidden)]
+pub use compute::take::bench_take_mojo;
+#[doc(hidden)]
+pub use compute::take::bench_take_scalar;
 pub use vtable::Primitive;
 
 mod native_value;