feat: arch intrinsic feature discovery pattern (#8750)

Currently we have a few places where we have arch/feature specific
kernels, they current use slower feature detection for each call.

This PR add a cache for the kernel to use and uses this pattern in the
codebase. I think we would later want to register different kernels in
the vortex session instead of this, but that can wait.

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: joseph-isaacs

encodings/fastlanes/src/bit_transpose/mod.rs

@@ -28,6 +28,10 @@ mod x86;
 mod validity;
 
 pub use validity::*;
+use vortex_buffer::CpuKernel;
+
+/// Signature shared by all 1024-bit transpose kernels.
+type TransposeKernel = unsafe fn(&[u8; 128], &mut [u8; 128]);
 
 /// Base indices for the first 64 output bytes (lanes 0-7).
 /// Each entry indicates the starting input byte index for that output byte group.
@@ -45,56 +49,64 @@ const TRANSPOSE_8X8: u64 = 0x0000_0000_F0F0_F0F0;
 
 /// Transpose 1024-bits into FastLanes layout.
 ///
-/// Dispatch to the best available implementation at runtime.
+/// Dispatch to the best available implementation, selected once on first call.
 #[inline]
 pub fn transpose_bits(input: &[u8; 128], output: &mut [u8; 128]) {
-    #[cfg(target_arch = "x86_64")]
-    {
-        // VBMI is fastest
-        if x86::has_vbmi() {
-            unsafe { x86::transpose_bits_vbmi(input, output) };
-            return;
+    static KERNEL: CpuKernel<TransposeKernel> = CpuKernel::new(|| {
+        #[cfg(target_arch = "x86_64")]
+        {
+            // VBMI is fastest
+            if x86::has_vbmi() {
+                return x86::transpose_bits_vbmi;
+            }
+            if x86::has_bmi2() {
+                return x86::transpose_bits_bmi2;
+            }
         }
-        if x86::has_bmi2() {
-            unsafe { x86::transpose_bits_bmi2(input, output) };
-            return;
+        // NEON is architecturally guaranteed on aarch64, so it needs no probe.
+        #[cfg(target_arch = "aarch64")]
+        return aarch64::transpose_bits_neon;
+        // The aarch64 arm above returns unconditionally, making this portable default
+        // unreachable there.
+        #[allow(unreachable_code)]
+        {
+            scalar::transpose_bits_scalar
         }
-        // Fall back to scalar
-        scalar::transpose_bits_scalar(input, output);
-    }
-    #[cfg(target_arch = "aarch64")]
-    {
-        unsafe { aarch64::transpose_bits_neon(input, output) };
-    }
-    #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
-    scalar::transpose_bits_scalar(input, output);
+    });
+    // SAFETY: the selector only returns kernels that are safe or whose required CPU
+    // features were probed before selection.
+    unsafe { KERNEL.get()(input, output) }
 }
 
 /// Untranspose 1024-bits from FastLanes layout.
 ///
-/// Dispatch untranspose to the best available implementation at runtime.
+/// Dispatch untranspose to the best available implementation, selected once on first call.
 #[inline]
 pub fn untranspose_bits(input: &[u8; 128], output: &mut [u8; 128]) {
-    #[cfg(target_arch = "x86_64")]
-    {
-        // VBMI is fastest
-        if x86::has_vbmi() {
-            unsafe { x86::untranspose_bits_vbmi(input, output) };
-            return;
+    static KERNEL: CpuKernel<TransposeKernel> = CpuKernel::new(|| {
+        #[cfg(target_arch = "x86_64")]
+        {
+            // VBMI is fastest
+            if x86::has_vbmi() {
+                return x86::untranspose_bits_vbmi;
+            }
+            if x86::has_bmi2() {
+                return x86::untranspose_bits_bmi2;
+            }
         }
-        if x86::has_bmi2() {
-            unsafe { x86::untranspose_bits_bmi2(input, output) };
-            return;
+        // NEON is architecturally guaranteed on aarch64, so it needs no probe.
+        #[cfg(target_arch = "aarch64")]
+        return aarch64::untranspose_bits_neon;
+        // The aarch64 arm above returns unconditionally, making this portable default
+        // unreachable there.
+        #[allow(unreachable_code)]
+        {
+            scalar::untranspose_bits_scalar
         }
-        // Fall back to scalar
-        scalar::untranspose_bits_scalar(input, output);
-    }
-    #[cfg(target_arch = "aarch64")]
-    {
-        unsafe { aarch64::untranspose_bits_neon(input, output) };
-    }
-    #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
-    scalar::untranspose_bits_scalar(input, output);
+    });
+    // SAFETY: the selector only returns kernels that are safe or whose required CPU
+    // features were probed before selection.
+    unsafe { KERNEL.get()(input, output) }
 }
 
 #[cfg(test)]

vortex-array/src/arrays/bool/compute/filter.rs

@@ -3,6 +3,7 @@
 
 use vortex_buffer::BitBuffer;
 use vortex_buffer::BitBufferMut;
+use vortex_buffer::CpuKernel;
 use vortex_buffer::get_bit;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
@@ -90,14 +91,19 @@ pub fn filter_bitbuffer_by_mask(
     mask_buf: &BitBuffer,
     true_count: usize,
 ) -> BitBuffer {
-    #[cfg(target_arch = "x86_64")]
-    {
-        if std::arch::is_x86_feature_detected!("bmi2") {
-            // SAFETY: BMI2 confirmed available; the inner function is compiled with BMI2.
-            return unsafe { filter_pext_bmi2(src, mask_buf, true_count) };
+    type FilterKernel = unsafe fn(&BitBuffer, &BitBuffer, usize) -> BitBuffer;
+    static KERNEL: CpuKernel<FilterKernel> = CpuKernel::new(|| {
+        #[cfg(target_arch = "x86_64")]
+        {
+            if std::arch::is_x86_feature_detected!("bmi2") {
+                return filter_pext_bmi2;
+            }
         }
-    }
-    filter_pext_fallback(src, mask_buf, true_count)
+        filter_pext_fallback
+    });
+    // SAFETY: the selector only returns kernels that are safe or whose required CPU
+    // features were probed before selection.
+    unsafe { KERNEL.get()(src, mask_buf, true_count) }
 }
 
 /// BMI2-native filter: entire function compiled with BMI2+POPCNT enabled.

vortex-buffer/Cargo.toml

@@ -49,6 +49,10 @@ harness = false
 name = "vortex_bitbuffer"
 harness = false
 
+[[bench]]
+name = "cpu_dispatch"
+harness = false
+
 [[bench]]
 name = "collect_bool"
 harness = false

vortex-buffer/benches/cpu_dispatch.rs

@@ -0,0 +1,160 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Overhead comparison of one-time CPU-feature dispatch mechanisms.
+//!
+//! Every variant ends in a call to the same `#[inline(never)]` kernel, so the
+//! difference between rows is pure dispatch overhead:
+//!
+//! * `direct`: plain call — the floor.
+//! * `cpu_kernel`: [`CpuKernel`] — relaxed load + null check + indirect call.
+//! * `cpu_kernel_unconditional`: a selector with no probes (the aarch64 shape) —
+//!   same steady state as `cpu_kernel`.
+//! * `lazy_lock`: `std::sync::LazyLock<fn>` — Once-state check + value load +
+//!   indirect call.
+//! * `feature_detect_each_call_1/_3`: the pre-`CpuKernel` pattern — 1 or 3
+//!   `is_x86_feature_detected!` cached lookups per call, then a direct call.
+
+use std::sync::LazyLock;
+
+use divan::Bencher;
+use divan::black_box;
+use divan::counter::ItemsCount;
+use vortex_buffer::CpuKernel;
+
+fn main() {
+    // Resolve every dispatcher and warm the std feature-detection cache so no
+    // benchmark iteration pays a one-time probe.
+    let seed = black_box(7);
+    let _ = via_direct(1, seed)
+        + via_cpu_kernel(1, seed)
+        + via_cpu_kernel_unconditional(1, seed)
+        + via_lazy_lock(1, seed);
+    #[cfg(target_arch = "x86_64")]
+    let _ = via_feature_detect_1(1, seed) + via_feature_detect_3(1, seed);
+
+    divan::main();
+}
+
+type Kernel = fn(u64, u64) -> u64;
+
+#[inline(never)]
+fn kernel_a(x: u64, seed: u64) -> u64 {
+    x.wrapping_mul(31).wrapping_add(seed)
+}
+
+#[inline(never)]
+fn kernel_b(x: u64, seed: u64) -> u64 {
+    x.wrapping_mul(33).wrapping_add(seed)
+}
+
+fn has_bmi2() -> bool {
+    #[cfg(target_arch = "x86_64")]
+    {
+        std::arch::is_x86_feature_detected!("bmi2")
+    }
+    #[cfg(not(target_arch = "x86_64"))]
+    {
+        false
+    }
+}
+
+// ── dispatch variants ───────────────────────────────────────────────────
+
+#[inline(never)]
+fn via_direct(x: u64, seed: u64) -> u64 {
+    kernel_a(x, seed)
+}
+
+#[cfg(target_arch = "x86_64")]
+#[inline(never)]
+fn via_feature_detect_1(x: u64, seed: u64) -> u64 {
+    if std::arch::is_x86_feature_detected!("bmi2") {
+        kernel_a(x, seed)
+    } else {
+        kernel_b(x, seed)
+    }
+}
+
+#[cfg(target_arch = "x86_64")]
+#[inline(never)]
+fn via_feature_detect_3(x: u64, seed: u64) -> u64 {
+    // Non-baseline features that are present on any recent x86_64 part, so all
+    // three cached lookups actually execute (mirrors the old select_in_chunk).
+    if std::arch::is_x86_feature_detected!("avx")
+        && std::arch::is_x86_feature_detected!("avx2")
+        && std::arch::is_x86_feature_detected!("bmi2")
+    {
+        kernel_a(x, seed)
+    } else {
+        kernel_b(x, seed)
+    }
+}
+
+static LAZY: LazyLock<Kernel> = LazyLock::new(|| if has_bmi2() { kernel_a } else { kernel_b });
+
+#[inline(never)]
+fn via_lazy_lock(x: u64, seed: u64) -> u64 {
+    (*LAZY)(x, seed)
+}
+
+static CPU_KERNEL: CpuKernel<Kernel> =
+    CpuKernel::new(|| if has_bmi2() { kernel_a } else { kernel_b });
+
+#[inline(never)]
+fn via_cpu_kernel(x: u64, seed: u64) -> u64 {
+    CPU_KERNEL.get()(x, seed)
+}
+
+static CPU_KERNEL_UNCONDITIONAL: CpuKernel<Kernel> = CpuKernel::new(|| kernel_a);
+
+#[inline(never)]
+fn via_cpu_kernel_unconditional(x: u64, seed: u64) -> u64 {
+    CPU_KERNEL_UNCONDITIONAL.get()(x, seed)
+}
+
+// ── benches ─────────────────────────────────────────────────────────────
+
+const CALLS: u64 = 1024;
+
+fn bench_via(bencher: Bencher, via: fn(u64, u64) -> u64) {
+    bencher.counter(ItemsCount::new(CALLS)).bench(|| {
+        let mut acc = 0u64;
+        for i in 0..CALLS {
+            acc = acc.wrapping_add(via(black_box(i), black_box(7)));
+        }
+        acc
+    })
+}
+
+#[divan::bench]
+fn direct(bencher: Bencher) {
+    bench_via(bencher, via_direct);
+}
+
+#[divan::bench]
+fn cpu_kernel_unconditional(bencher: Bencher) {
+    bench_via(bencher, via_cpu_kernel_unconditional);
+}
+
+#[divan::bench]
+fn cpu_kernel(bencher: Bencher) {
+    bench_via(bencher, via_cpu_kernel);
+}
+
+#[divan::bench]
+fn lazy_lock(bencher: Bencher) {
+    bench_via(bencher, via_lazy_lock);
+}
+
+#[cfg(target_arch = "x86_64")]
+#[divan::bench]
+fn feature_detect_each_call_1(bencher: Bencher) {
+    bench_via(bencher, via_feature_detect_1);
+}
+
+#[cfg(target_arch = "x86_64")]
+#[divan::bench]
+fn feature_detect_each_call_3(bencher: Bencher) {
+    bench_via(bencher, via_feature_detect_3);
+}

vortex-buffer/benches/vortex_bitbuffer.rs

@@ -19,6 +19,12 @@ fn main() {
         let _ = is_x86_feature_detected!("avx512vpopcntdq");
     }
 
+    // Pre-resolve the one-time CpuKernel selections for the count/select paths so
+    // no benchmark iteration pays the first-call cost.
+    let warm = BitBuffer::from_iter((0..512).map(|i| i % 3 == 0));
+    let _ = warm.true_count();
+    let _ = warm.select(1);
+
     divan::main();
 }