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 scalar path: count_ones_aligned calls the scalar kernel directly
below 32 bytes (where SIMD never engaged anyway), and scan_chunks
returns immediately for an empty chunk slice. Without these gates the
dispatch indirection dominates sub-SIMD-threshold workloads (CodSpeed
flagged true_count[128] at about -30% instructions).

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.

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/src/bit/count_ones.rs

@@ -4,6 +4,8 @@
 #[cfg(target_arch = "x86_64")]
 use vortex_error::VortexExpect;
 
+use crate::dispatch::CpuKernel;
+
 #[inline]
 pub fn count_ones(bytes: &[u8], offset: usize, len: usize) -> usize {
     if bytes.is_empty() {
@@ -70,24 +72,61 @@ fn mask_byte(byte: u8, bit_offset: usize, bit_len: usize) -> u8 {
     shifted & mask
 }
 
+type CountOnes = unsafe fn(&[u8]) -> usize;
+
 #[inline]
 fn count_ones_aligned(bytes: &[u8]) -> usize {
-    #[cfg(target_arch = "x86_64")]
-    {
-        if bytes.len() >= 64
-            && is_x86_feature_detected!("avx512f")
-            && is_x86_feature_detected!("avx512vpopcntdq")
+    // SIMD kernels only pay off from 32 bytes. Below that, call the scalar kernel
+    // directly: it stays inlinable and skips the dispatch indirection, which would
+    // otherwise dominate the couple of word popcounts.
+    if bytes.len() < 32 {
+        return count_ones_aligned_scalar(bytes);
+    }
+
+    static KERNEL: CpuKernel<CountOnes> = CpuKernel::new(|| {
+        #[cfg(target_arch = "x86_64")]
         {
-            // SAFETY: Runtime detection guarantees the required target features.
-            return unsafe { count_ones_aligned_avx512(bytes) };
+            if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512vpopcntdq") {
+                return count_ones_aligned_avx512_any_len;
+            }
+            if is_x86_feature_detected!("avx2") {
+                return count_ones_aligned_avx2_any_len;
+            }
         }
+        count_ones_aligned_scalar
+    });
+    // SAFETY: the selector only returns kernels that are safe or whose required CPU
+    // features were probed before selection.
+    unsafe { KERNEL.get()(bytes) }
+}
 
-        if bytes.len() >= 32 && is_x86_feature_detected!("avx2") {
-            // SAFETY: Runtime detection guarantees the required target features.
-            return unsafe { count_ones_aligned_avx2(bytes) };
-        }
+/// Length-aware AVX-512 entry: SIMD only pays off from 64 (AVX-512) / 32 (AVX2) bytes.
+///
+/// # Safety
+/// Requires AVX-512F and AVX-512VPOPCNTDQ (which imply AVX2).
+#[cfg(target_arch = "x86_64")]
+unsafe fn count_ones_aligned_avx512_any_len(bytes: &[u8]) -> usize {
+    if bytes.len() >= 64 {
+        // SAFETY: the caller guarantees the required target features.
+        return unsafe { count_ones_aligned_avx512(bytes) };
+    }
+    if bytes.len() >= 32 {
+        // SAFETY: every AVX-512F CPU also supports AVX2.
+        return unsafe { count_ones_aligned_avx2(bytes) };
     }
+    count_ones_aligned_scalar(bytes)
+}
 
+/// Length-aware AVX2 entry: SIMD only pays off from 32 bytes.
+///
+/// # Safety
+/// Requires AVX2.
+#[cfg(target_arch = "x86_64")]
+unsafe fn count_ones_aligned_avx2_any_len(bytes: &[u8]) -> usize {
+    if bytes.len() >= 32 {
+        // SAFETY: the caller guarantees AVX2.
+        return unsafe { count_ones_aligned_avx2(bytes) };
+    }
     count_ones_aligned_scalar(bytes)
 }