feat: model arch and CPU-feature dispatch dimensions in CpuKernel

Replace LazyCpuDispatch with CpuKernel, which models the two dispatch
dimensions separately: the compile-time dimension (x86_64 vs aarch64) as
cfg-selected const candidate tables and fallback kernels, and the runtime
dimension (AVX-512 vs BMI2 vs scalar) as an ordered list of
Candidate { available, kernel } probes, tried once on first call.

The candidate table and fallback are passed at get() rather than stored in
the static: the atomic slot gives the struct interior mutability, which
places the whole static in writable memory where LLVM cannot const-fold
fields. Passed as rodata consts instead, an empty candidate table (the
aarch64 shape today) folds get() into a direct call — verified by
disassembly, where LLVM merges the empty-table wrapper with the
direct-call wrapper into identical code.

Add a divan benchmark (benches/cpu_dispatch.rs) comparing dispatch
overhead per call against the same #[inline(never)] kernel. Fastest-of-3
on 1024 calls/iter: direct, CpuKernel (empty and populated), one
is_x86_feature_detected! per call, and branchless CpuDispatch are all
equal within noise (~2.74 us); LazyLock<fn> and the previous
three-feature-check-per-call pattern are both ~33% slower (~3.64 us).

Dispatch call sites in vortex-buffer bit/select.rs and fastlanes
bit_transpose now contain no cfg blocks in the function body: per-arch
kernels live in the candidate/fallback consts.

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: claude

encodings/fastlanes/src/bit_transpose/mod.rs

@@ -28,13 +28,71 @@ mod x86;
 mod validity;
 
 pub use validity::*;
-#[cfg(target_arch = "x86_64")]
-use vortex_buffer::LazyCpuDispatch;
+use vortex_buffer::Candidate;
+use vortex_buffer::CpuKernel;
 
 /// Signature shared by all 1024-bit transpose kernels.
-#[cfg(target_arch = "x86_64")]
 type TransposeKernel = fn(&[u8; 128], &mut [u8; 128]);
 
+// Runtime dimension: kernels that need a CPU probe, best first.
+#[cfg(target_arch = "x86_64")]
+const TRANSPOSE_CANDIDATES: &[Candidate<TransposeKernel>] = &[
+    Candidate {
+        available: x86::has_vbmi,
+        kernel: |input, output| {
+            // SAFETY: `available` confirms VBMI.
+            unsafe { x86::transpose_bits_vbmi(input, output) }
+        },
+    },
+    Candidate {
+        available: x86::has_bmi2,
+        kernel: |input, output| {
+            // SAFETY: `available` confirms BMI2.
+            unsafe { x86::transpose_bits_bmi2(input, output) }
+        },
+    },
+];
+#[cfg(not(target_arch = "x86_64"))]
+const TRANSPOSE_CANDIDATES: &[Candidate<TransposeKernel>] = &[];
+
+#[cfg(target_arch = "x86_64")]
+const UNTRANSPOSE_CANDIDATES: &[Candidate<TransposeKernel>] = &[
+    Candidate {
+        available: x86::has_vbmi,
+        kernel: |input, output| {
+            // SAFETY: `available` confirms VBMI.
+            unsafe { x86::untranspose_bits_vbmi(input, output) }
+        },
+    },
+    Candidate {
+        available: x86::has_bmi2,
+        kernel: |input, output| {
+            // SAFETY: `available` confirms BMI2.
+            unsafe { x86::untranspose_bits_bmi2(input, output) }
+        },
+    },
+];
+#[cfg(not(target_arch = "x86_64"))]
+const UNTRANSPOSE_CANDIDATES: &[Candidate<TransposeKernel>] = &[];
+
+// Compile-time dimension: the per-architecture unconditional kernel. NEON is
+// architecturally guaranteed on aarch64, so it needs no probe.
+#[cfg(target_arch = "aarch64")]
+const TRANSPOSE_FALLBACK: TransposeKernel = |input, output| {
+    // SAFETY: NEON is architecturally guaranteed on aarch64.
+    unsafe { aarch64::transpose_bits_neon(input, output) }
+};
+#[cfg(not(target_arch = "aarch64"))]
+const TRANSPOSE_FALLBACK: TransposeKernel = scalar::transpose_bits_scalar;
+
+#[cfg(target_arch = "aarch64")]
+const UNTRANSPOSE_FALLBACK: TransposeKernel = |input, output| {
+    // SAFETY: NEON is architecturally guaranteed on aarch64.
+    unsafe { aarch64::untranspose_bits_neon(input, output) }
+};
+#[cfg(not(target_arch = "aarch64"))]
+const UNTRANSPOSE_FALLBACK: TransposeKernel = scalar::untranspose_bits_scalar;
+
 /// Base indices for the first 64 output bytes (lanes 0-7).
 /// Each entry indicates the starting input byte index for that output byte group.
 /// Pattern: [0*2, 4*2, 2*2, 6*2, 1*2, 5*2, 3*2, 7*2] = [0, 8, 4, 12, 2, 10, 6, 14]
@@ -54,63 +112,17 @@ const TRANSPOSE_8X8: u64 = 0x0000_0000_F0F0_F0F0;
 /// 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")]
-    {
-        static SELECTED: LazyCpuDispatch<TransposeKernel> = LazyCpuDispatch::new(|| {
-            // VBMI is fastest
-            if x86::has_vbmi() {
-                // SAFETY: VBMI support was just detected.
-                |input, output| unsafe { x86::transpose_bits_vbmi(input, output) }
-            } else if x86::has_bmi2() {
-                // SAFETY: BMI2 support was just detected.
-                |input, output| unsafe { x86::transpose_bits_bmi2(input, output) }
-            } else {
-                scalar::transpose_bits_scalar
-            }
-        });
-        SELECTED.get()(input, output)
-    }
-    #[cfg(target_arch = "aarch64")]
-    {
-        // SAFETY: NEON is architecturally guaranteed on aarch64.
-        unsafe { aarch64::transpose_bits_neon(input, output) }
-    }
-    #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
-    {
-        scalar::transpose_bits_scalar(input, output)
-    }
+    static KERNEL: CpuKernel<TransposeKernel> = CpuKernel::new();
+    KERNEL.get(TRANSPOSE_CANDIDATES, TRANSPOSE_FALLBACK)(input, output)
 }
 
 /// Untranspose 1024-bits from FastLanes layout.
 ///
 /// 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")]
-    {
-        static SELECTED: LazyCpuDispatch<TransposeKernel> = LazyCpuDispatch::new(|| {
-            // VBMI is fastest
-            if x86::has_vbmi() {
-                // SAFETY: VBMI support was just detected.
-                |input, output| unsafe { x86::untranspose_bits_vbmi(input, output) }
-            } else if x86::has_bmi2() {
-                // SAFETY: BMI2 support was just detected.
-                |input, output| unsafe { x86::untranspose_bits_bmi2(input, output) }
-            } else {
-                scalar::untranspose_bits_scalar
-            }
-        });
-        SELECTED.get()(input, output)
-    }
-    #[cfg(target_arch = "aarch64")]
-    {
-        // SAFETY: NEON is architecturally guaranteed on aarch64.
-        unsafe { aarch64::untranspose_bits_neon(input, output) }
-    }
-    #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
-    {
-        scalar::untranspose_bits_scalar(input, output)
-    }
+    static KERNEL: CpuKernel<TransposeKernel> = CpuKernel::new();
+    KERNEL.get(UNTRANSPOSE_CANDIDATES, UNTRANSPOSE_FALLBACK)(input, output)
 }
 
 #[cfg(test)]

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,187 @@
+// 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_empty_candidates`: [`CpuKernel`] with an empty candidate table
+//!   (the aarch64 shape) — should const-fold to `direct`.
+//! * `cpu_dispatch_branchless`: [`CpuDispatch`] — relaxed load + indirect call.
+//! * `cpu_kernel`: [`CpuKernel`] — relaxed load + null check + indirect call.
+//! * `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::Candidate;
+use vortex_buffer::CpuDispatch;
+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_empty_candidates(1, seed)
+        + via_cpu_dispatch(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)
+}
+
+const CPU_KERNEL_CANDIDATES: &[Candidate<Kernel>] = &[Candidate {
+    available: has_bmi2,
+    kernel: kernel_a,
+}];
+
+static CPU_KERNEL: CpuKernel<Kernel> = CpuKernel::new();
+
+#[inline(never)]
+fn via_cpu_kernel(x: u64, seed: u64) -> u64 {
+    CPU_KERNEL.get(CPU_KERNEL_CANDIDATES, kernel_b)(x, seed)
+}
+
+static CPU_KERNEL_EMPTY: CpuKernel<Kernel> = CpuKernel::new();
+
+#[inline(never)]
+fn via_cpu_kernel_empty_candidates(x: u64, seed: u64) -> u64 {
+    CPU_KERNEL_EMPTY.get(&[], kernel_a)(x, seed)
+}
+
+static CPU_DISPATCH: CpuDispatch<Kernel> = CpuDispatch::new(resolve);
+
+#[cold]
+fn resolve(x: u64, seed: u64) -> u64 {
+    let kernel: Kernel = if has_bmi2() { kernel_a } else { kernel_b };
+    CPU_DISPATCH.set(kernel);
+    kernel(x, seed)
+}
+
+#[inline(never)]
+fn via_cpu_dispatch(x: u64, seed: u64) -> u64 {
+    CPU_DISPATCH.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_empty_candidates(bencher: Bencher) {
+    bench_via(bencher, via_cpu_kernel_empty_candidates);
+}
+
+#[divan::bench]
+fn cpu_dispatch_branchless(bencher: Bencher) {
+    bench_via(bencher, via_cpu_dispatch);
+}
+
+#[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);
+}