feat: add CpuKernel for one-time CPU-feature kernel dispatch

Hot kernels re-ran cached CPU-feature checks on every call: each
is_x86_feature_detected! is an atomic load + bit test per call, and sites
like select_in_chunk chained three of them per 64-byte chunk. aarch64
paths were separate cfg blocks at every call site.

Add vortex_buffer::CpuKernel<F>, a LazyLock-like slot for function
pointers: the selector passed to CpuKernel::new runs once on the first
get(), and every later call is a relaxed atomic load, a never-taken
predicted branch, and an indirect call. No macros; the two
transmute_copy calls between F and the atomic pointer are contained in
the type, whose constructor statically asserts F is pointer-sized.

Selectors model both dispatch dimensions in plain code: one
cfg(target_arch) block per architecture that returns early (runtime
feature probes as an if-chain inside), then the portable default as the
plain tail - no cfg(not(any(...))) negation blocks anywhere. An arm that
returns unconditionally (aarch64 NEON needs no probe) makes the tail
unreachable there, silenced by an allow(unreachable_code) scoped to just
the tail block. Kernel types are unsafe fn pointers so #[target_feature]
kernels coerce directly by name with a single narrow unsafe call per
site.

Tiny inputs are gated before the dispatch so they keep the old fully
inlined code path: count_ones_aligned calls the scalar kernel directly
below 32 bytes (where SIMD never engaged anyway), and scan_chunks calls
the per-architecture unconditional kernel directly for two chunks or
fewer. Without these gates the dispatch indirection dominates
sub-SIMD-threshold workloads: CodSpeed flagged true_count[128] (~-30%)
and rank_single[(1024, 0.1)]; cachegrind measured the rank regression at
+27 instructions per call and parity after the gates (37.70M vs 37.71M
baseline instructions for 100k rank calls).

Converted sites: fastlanes transpose_bits/untranspose_bits, vortex-buffer
scan_chunks/select_in_chunk/select_in_word/count_ones_aligned,
vortex-array filter_bitbuffer_by_mask, and vortex-mask
intersect_bit_buffers_dispatch/intersect_rank_indices_dispatch.
Intentionally left: primitive take (already one-time via LazyLock dyn
kernel), intersect_mask_driven_dispatch (generic over an iterator type,
so its instantiations cannot share one fn-pointer slot), and the
compile-time IS_CONST_LANE_WIDTH constant.

The count_ones any-length wrappers carry #[target_feature] so the SIMD
kernels inline into them, avoiding an extra call level, and the
vortex_bitbuffer and rank bench mains pre-resolve the kernel slots
(matching the existing CPUID pre-warm precedent) so no single-shot
benchmark measurement includes first-call selection.

benches/cpu_dispatch.rs compares dispatch overhead per call against the
same #[inline(never)] kernel. Fastest-of-3 per 1024 calls: direct 2.44us,
CpuKernel 2.44us (indistinguishable, confirmed by disassembly: load +
fused null test + indirect jmp), LazyLock<fn> 3.04us, and three cached
feature checks per call 4.54us. All touched crates also cargo check
clean on aarch64-unknown-linux-gnu.

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DUFxXZLv4xxEqEYcoGWhfJ

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

@@ -48,3 +48,7 @@ harness = false
 [[bench]]
 name = "vortex_bitbuffer"
 harness = false
+
+[[bench]]
+name = "cpu_dispatch"
+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();
 }