feat: add cpu_dispatch! macro for one-time CPU-feature kernel selection

Several hot kernels re-run cached CPU-feature checks on every call
(is_x86_feature_detected! is an atomic load + bit-test + branch chain per
check, per call). Add vortex_buffer::cpu_dispatch!, an ifunc-style
self-patching function pointer: the selector runs once on first call,
patches an AtomicPtr slot, and every later call is a single relaxed load
plus a predictable indirect call.

Convert the worst repeat-detection sites as a demonstration:
- fastlanes transpose_bits/untranspose_bits (2 feature checks per 128B call)
- vortex-buffer select_in_chunk (3 checks per 64B chunk), scan_chunks_impl,
  and select_in_word

aarch64 paths are intentionally left as direct #[cfg] calls: NEON is
architecturally guaranteed, so compile-time dispatch is already zero-cost
and inlinable there. The macro documents when to use it (coarse-grained
kernels with >1 runtime candidate) and when not to (per-element loops,
which should hoist + monomorphize like vortex-mask::intersect_by_rank).

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

@@ -45,56 +45,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")]
-    {
+    return transpose_bits_x86(input, output);
+    // SAFETY: NEON is architecturally guaranteed on aarch64.
+    #[cfg(target_arch = "aarch64")]
+    return unsafe { aarch64::transpose_bits_neon(input, output) };
+    #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
+    scalar::transpose_bits_scalar(input, output);
+}
+
+#[cfg(target_arch = "x86_64")]
+vortex_buffer::cpu_dispatch! {
+    fn transpose_bits_x86(input: &[u8; 128], output: &mut [u8; 128]);
+    select {
         // VBMI is fastest
         if x86::has_vbmi() {
-            unsafe { x86::transpose_bits_vbmi(input, output) };
-            return;
+            // SAFETY: VBMI support was just detected.
+            return |input, output| unsafe { x86::transpose_bits_vbmi(input, output) };
         }
         if x86::has_bmi2() {
-            unsafe { x86::transpose_bits_bmi2(input, output) };
-            return;
+            // SAFETY: BMI2 support was just detected.
+            return |input, output| unsafe { x86::transpose_bits_bmi2(input, output) };
         }
-        // Fall back to scalar
-        scalar::transpose_bits_scalar(input, output);
+        scalar::transpose_bits_scalar
     }
-    #[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);
 }
 
 /// 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")]
-    {
+    return untranspose_bits_x86(input, output);
+    // SAFETY: NEON is architecturally guaranteed on aarch64.
+    #[cfg(target_arch = "aarch64")]
+    return unsafe { aarch64::untranspose_bits_neon(input, output) };
+    #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
+    scalar::untranspose_bits_scalar(input, output);
+}
+
+#[cfg(target_arch = "x86_64")]
+vortex_buffer::cpu_dispatch! {
+    fn untranspose_bits_x86(input: &[u8; 128], output: &mut [u8; 128]);
+    select {
         // VBMI is fastest
         if x86::has_vbmi() {
-            unsafe { x86::untranspose_bits_vbmi(input, output) };
-            return;
+            // SAFETY: VBMI support was just detected.
+            return |input, output| unsafe { x86::untranspose_bits_vbmi(input, output) };
         }
         if x86::has_bmi2() {
-            unsafe { x86::untranspose_bits_bmi2(input, output) };
-            return;
+            // SAFETY: BMI2 support was just detected.
+            return |input, output| unsafe { x86::untranspose_bits_bmi2(input, output) };
         }
-        // Fall back to scalar
-        scalar::untranspose_bits_scalar(input, output);
+        scalar::untranspose_bits_scalar
     }
-    #[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);
 }
 
 #[cfg(test)]

vortex-buffer/src/bit/select.rs

@@ -94,14 +94,18 @@ fn scan_chunks_impl(chunks: &[[u8; 64]], remaining: usize, pos: usize) -> (usize
 }
 
 #[cfg(target_arch = "x86_64")]
-#[inline]
-fn scan_chunks_impl(chunks: &[[u8; 64]], remaining: usize, pos: usize) -> (usize, usize, usize) {
-    if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512vpopcntdq") {
-        // SAFETY: runtime detection guarantees the required target features.
-        return unsafe { scan_chunks_avx512_vpopcnt(chunks, remaining, pos) };
-    }
+crate::cpu_dispatch! {
+    fn scan_chunks_impl(chunks: &[[u8; 64]], remaining: usize, pos: usize) -> (usize, usize, usize);
+    select {
+        if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512vpopcntdq") {
+            // SAFETY: runtime detection guarantees the required target features.
+            return |chunks, remaining, pos| unsafe {
+                scan_chunks_avx512_vpopcnt(chunks, remaining, pos)
+            };
+        }
 
-    scan_chunks_scalar(chunks, remaining, pos)
+        scan_chunks_scalar
+    }
 }
 
 #[cfg(target_arch = "aarch64")]
@@ -281,20 +285,27 @@ fn scan_words_scalar(
 // ── In-chunk select ─────────────────────────────────────────────────────
 
 /// Position of the `nth` set bit inside a 64-byte chunk (0-indexed).
+#[cfg(not(target_arch = "x86_64"))]
 #[inline]
 fn select_in_chunk(chunk: &[u8; 64], nth: usize) -> usize {
-    #[cfg(target_arch = "x86_64")]
-    {
+    select_in_chunk_scalar(chunk, nth)
+}
+
+#[cfg(target_arch = "x86_64")]
+crate::cpu_dispatch! {
+    /// Position of the `nth` set bit inside a 64-byte chunk (0-indexed).
+    fn select_in_chunk(chunk: &[u8; 64], nth: usize) -> usize;
+    select {
         if is_x86_feature_detected!("avx512f")
             && is_x86_feature_detected!("avx512vpopcntdq")
             && is_x86_feature_detected!("avx512vbmi2")
         {
             // SAFETY: runtime detection guarantees the required target features.
-            return unsafe { select_in_chunk_vbmi2(chunk, nth) };
+            return |chunk, nth| unsafe { select_in_chunk_vbmi2(chunk, nth) };
         }
-    }
 
-    select_in_chunk_scalar(chunk, nth)
+        select_in_chunk_scalar
+    }
 }
 
 #[cfg(target_arch = "x86_64")]
@@ -360,16 +371,23 @@ fn count_ones_chunk(chunk: &[u8; 64]) -> usize {
 // ── In-word select ──────────────────────────────────────────────────────
 
 /// Position of the `nth` set bit inside a u64 (0-indexed, little-endian bit order).
+#[cfg(not(target_arch = "x86_64"))]
 #[inline]
 fn select_in_word(word: u64, nth: usize) -> usize {
-    #[cfg(target_arch = "x86_64")]
-    {
+    select_in_word_scalar(word, nth)
+}
+
+#[cfg(target_arch = "x86_64")]
+crate::cpu_dispatch! {
+    /// Position of the `nth` set bit inside a u64 (0-indexed, little-endian bit order).
+    fn select_in_word(word: u64, nth: usize) -> usize;
+    select {
         if is_x86_feature_detected!("bmi2") {
             // SAFETY: runtime detection guarantees the required target feature.
-            return unsafe { select_in_word_bmi2(word, nth) };
+            return |word, nth| unsafe { select_in_word_bmi2(word, nth) };
         }
+        select_in_word_scalar
     }
-    select_in_word_scalar(word, nth)
 }
 
 /// BMI2: deposit a single bit at the nth set-bit position, then count trailing zeros.

vortex-buffer/src/dispatch.rs

@@ -0,0 +1,135 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! One-time CPU-feature-based function dispatch.
+//!
+//! [`cpu_dispatch!`](crate::cpu_dispatch) defines a function whose implementation is
+//! selected by runtime CPU-feature detection exactly once, on the first call. Every
+//! later call is a single relaxed atomic load plus an indirect call — no repeated
+//! feature checks, no branch chains, no `LazyLock` initialized-state branch.
+//!
+//! This is the "ifunc" self-patching pattern: the function-pointer slot initially
+//! points at a resolver; the resolver picks the best implementation, patches the slot,
+//! and forwards the first call. Races are benign — the slot only ever holds valid
+//! function pointers of the same type, and every candidate computes the same result.
+//!
+//! # When to use it
+//!
+//! Use this for *coarse-grained* kernels (a whole chunk or buffer per call) that have
+//! more than one runtime candidate on the target architecture, e.g. x86_64 kernels
+//! with AVX-512 / BMI2 / scalar variants. Today that means x86_64: on aarch64 NEON is
+//! architecturally guaranteed, so a plain `#[cfg(target_arch = "aarch64")]` direct
+//! call is already zero-cost and inlinable — do not funnel it through a function
+//! pointer. If SVE/SVE2 kernels are added later, aarch64 call sites get the same
+//! one-time selection by adding `is_aarch64_feature_detected!` arms to the selector.
+//!
+//! # When NOT to use it
+//!
+//! Do not put the dispatched call inside a per-element hot loop: the indirect call
+//! blocks inlining. Hoist the decision to the outermost per-buffer entry point and
+//! monomorphize the loop instead, like the `Bmi2`/`Portable` type-parameter pattern in
+//! `vortex-mask::intersect_by_rank`.
+
+/// Define a function whose implementation is chosen by CPU-feature detection once, on
+/// first call, and cached in a self-patching function pointer thereafter.
+///
+/// The `select` block is the body of a resolver `fn() -> fn(args...) -> ret` that runs
+/// at most once per process (once per racing thread in the worst case). It must return
+/// the chosen implementation as a function pointer; non-capturing closures coerce, so
+/// `unsafe` `#[target_feature]` kernels can be wrapped inline after detection proves
+/// them sound.
+///
+/// # Example
+///
+/// ```
+/// vortex_buffer::cpu_dispatch! {
+///     /// Sums a slice, using the best kernel for the current CPU.
+///     pub fn sum(values: &[u64]) -> u64;
+///     select {
+///         #[cfg(target_arch = "x86_64")]
+///         if std::arch::is_x86_feature_detected!("avx2") {
+///             // return |values| unsafe { sum_avx2(values) };
+///         }
+///         |values: &[u64]| values.iter().sum()
+///     }
+/// }
+/// assert_eq!(sum(&[1, 2, 3]), 6);
+/// ```
+#[macro_export]
+macro_rules! cpu_dispatch {
+    (
+        $(#[$meta:meta])*
+        $vis:vis fn $name:ident($($arg:ident : $ty:ty),* $(,)?) $(-> $ret:ty)?;
+        select $select_body:block
+    ) => {
+        $(#[$meta])*
+        #[inline]
+        $vis fn $name($($arg: $ty),*) $(-> $ret)? {
+            type Impl = fn($($ty),*) $(-> $ret)?;
+
+            fn select() -> Impl $select_body
+
+            fn select_then_call($($arg: $ty),*) $(-> $ret)? {
+                let selected = select();
+                SELECTED.store(
+                    selected as *mut (),
+                    ::core::sync::atomic::Ordering::Relaxed,
+                );
+                selected($($arg),*)
+            }
+
+            static SELECTED: ::core::sync::atomic::AtomicPtr<()> =
+                ::core::sync::atomic::AtomicPtr::new(select_then_call as Impl as *mut ());
+
+            let fn_ptr = SELECTED.load(::core::sync::atomic::Ordering::Relaxed);
+            // SAFETY: `SELECTED` only ever holds values created by casting an `Impl`
+            // function pointer to `*mut ()`, so casting back is sound.
+            let selected = unsafe { ::core::mem::transmute::<*mut (), Impl>(fn_ptr) };
+            selected($($arg),*)
+        }
+    };
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::atomic::AtomicUsize;
+    use std::sync::atomic::Ordering;
+
+    static SELECT_CALLS: AtomicUsize = AtomicUsize::new(0);
+
+    crate::cpu_dispatch! {
+        fn add_one(x: u64) -> u64;
+        select {
+            SELECT_CALLS.fetch_add(1, Ordering::Relaxed);
+            |x| x + 1
+        }
+    }
+
+    crate::cpu_dispatch! {
+        fn xor_bytes(input: &[u8; 4], output: &mut [u8; 4]);
+        select {
+            fn xor_impl(input: &[u8; 4], output: &mut [u8; 4]) {
+                for (o, i) in output.iter_mut().zip(input) {
+                    *o ^= *i;
+                }
+            }
+            xor_impl
+        }
+    }
+
+    #[test]
+    fn selects_once_then_dispatches() {
+        assert_eq!(add_one(41), 42);
+        assert_eq!(add_one(1), 2);
+        assert_eq!(add_one(2), 3);
+        assert_eq!(SELECT_CALLS.load(Ordering::Relaxed), 1);
+    }
+
+    #[test]
+    fn supports_reference_arguments() {
+        let input = [1, 2, 3, 4];
+        let mut output = [0, 0, 0, 0];
+        xor_bytes(&input, &mut output);
+        assert_eq!(output, input);
+    }
+}

vortex-buffer/src/lib.rs

@@ -62,6 +62,7 @@ mod buffer_mut;
 mod bytes;
 mod r#const;
 mod debug;
+mod dispatch;
 mod macros;
 #[cfg(feature = "memmap2")]
 mod memmap2;