refactor: replace cpu_dispatch! macro with macro-free CpuDispatch<F> slot

Drop the cpu_dispatch! macro in favor of a plain generic struct with the
same steady-state codegen. CpuDispatch<F> is a self-patching function
pointer slot: it is constructed pointing at a caller-written resolver of
the same signature, which detects CPU features once, set()s the chosen
kernel into the slot, and forwards the first call. Every later call is a
single relaxed atomic load plus an indirect call, with no null check or
LazyLock-style initialized branch (verified: the dispatch shim compiles
to `mov SELECTED(%rip), %rax; jmp *%rax`).

The pointer/fn-pointer transmutes are contained inside CpuDispatch, whose
constructor statically asserts F is pointer-sized; call sites are entirely
safe plain Rust with function-local statics.

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,6 +28,12 @@ mod x86;
 mod validity;
 
 pub use validity::*;
+#[cfg(target_arch = "x86_64")]
+use vortex_buffer::CpuDispatch;
+
+/// Signature shared by all 1024-bit transpose kernels.
+#[cfg(target_arch = "x86_64")]
+type TransposeKernel = 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.
@@ -49,28 +55,35 @@ const TRANSPOSE_8X8: u64 = 0x0000_0000_F0F0_F0F0;
 #[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.
+    {
+        static SELECTED: CpuDispatch<TransposeKernel> = CpuDispatch::new(resolve);
+
+        #[cold]
+        fn resolve(input: &[u8; 128], output: &mut [u8; 128]) {
+            // VBMI is fastest
+            let kernel: TransposeKernel = 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.set(kernel);
+            kernel(input, output);
+        }
+
+        SELECTED.get()(input, output)
+    }
     #[cfg(target_arch = "aarch64")]
-    return unsafe { aarch64::transpose_bits_neon(input, output) };
+    {
+        // 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);
-}
-
-#[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() {
-            // SAFETY: VBMI support was just detected.
-            return |input, output| unsafe { x86::transpose_bits_vbmi(input, output) };
-        }
-        if x86::has_bmi2() {
-            // SAFETY: BMI2 support was just detected.
-            return |input, output| unsafe { x86::transpose_bits_bmi2(input, output) };
-        }
-        scalar::transpose_bits_scalar
+    {
+        scalar::transpose_bits_scalar(input, output)
     }
 }
 
@@ -80,28 +93,35 @@ vortex_buffer::cpu_dispatch! {
 #[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.
+    {
+        static SELECTED: CpuDispatch<TransposeKernel> = CpuDispatch::new(resolve);
+
+        #[cold]
+        fn resolve(input: &[u8; 128], output: &mut [u8; 128]) {
+            // VBMI is fastest
+            let kernel: TransposeKernel = 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.set(kernel);
+            kernel(input, output);
+        }
+
+        SELECTED.get()(input, output)
+    }
     #[cfg(target_arch = "aarch64")]
-    return unsafe { aarch64::untranspose_bits_neon(input, output) };
+    {
+        // 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);
-}
-
-#[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() {
-            // SAFETY: VBMI support was just detected.
-            return |input, output| unsafe { x86::untranspose_bits_vbmi(input, output) };
-        }
-        if x86::has_bmi2() {
-            // SAFETY: BMI2 support was just detected.
-            return |input, output| unsafe { x86::untranspose_bits_bmi2(input, output) };
-        }
-        scalar::untranspose_bits_scalar
+    {
+        scalar::untranspose_bits_scalar(input, output)
     }
 }
 

vortex-buffer/src/bit/select.rs

@@ -2,6 +2,8 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use super::count_ones::align_offset_len;
+#[cfg(target_arch = "x86_64")]
+use crate::dispatch::CpuDispatch;
 
 /// Returns the position of the `nth` set bit (0-indexed) within the logical range
 /// `[offset, offset + len)` of the given byte slice.
@@ -84,34 +86,43 @@ pub fn bit_select(bytes: &[u8], offset: usize, len: usize, nth: usize) -> Option
 /// is the rank *within* that chunk. Otherwise all chunks were consumed.
 #[inline]
 fn scan_chunks(chunks: &[[u8; 64]], remaining: usize, pos: usize) -> (usize, usize, usize) {
-    scan_chunks_impl(chunks, remaining, pos)
-}
-
-#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
-#[inline]
-fn scan_chunks_impl(chunks: &[[u8; 64]], remaining: usize, pos: usize) -> (usize, usize, usize) {
-    scan_chunks_scalar(chunks, remaining, pos)
-}
-
-#[cfg(target_arch = "x86_64")]
-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)
+    #[cfg(target_arch = "x86_64")]
+    {
+        type ScanChunks = fn(&[[u8; 64]], usize, usize) -> (usize, usize, usize);
+        static SELECTED: CpuDispatch<ScanChunks> = CpuDispatch::new(resolve);
+
+        #[cold]
+        fn resolve(chunks: &[[u8; 64]], remaining: usize, pos: usize) -> (usize, usize, usize) {
+            let kernel: ScanChunks = if is_x86_feature_detected!("avx512f")
+                && is_x86_feature_detected!("avx512vpopcntdq")
+            {
+                // SAFETY: runtime detection guarantees the required target features.
+                |chunks, remaining, pos| unsafe {
+                    scan_chunks_avx512_vpopcnt(chunks, remaining, pos)
+                }
+            } else {
+                scan_chunks_scalar
             };
+            SELECTED.set(kernel);
+            kernel(chunks, remaining, pos)
         }
 
-        scan_chunks_scalar
+        SELECTED.get()(chunks, remaining, pos)
+    }
+    #[cfg(target_arch = "aarch64")]
+    {
+        scan_chunks_neon(chunks, remaining, pos)
+    }
+    #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
+    {
+        scan_chunks_scalar(chunks, remaining, pos)
     }
 }
 
 #[cfg(target_arch = "aarch64")]
 #[allow(clippy::cast_possible_truncation)] // u64 → usize is lossless on aarch64 (64-bit)
 #[inline]
-fn scan_chunks_impl(
+fn scan_chunks_neon(
     chunks: &[[u8; 64]],
     mut remaining: usize,
     mut pos: usize,
@@ -285,26 +296,32 @@ 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 {
-    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 |chunk, nth| unsafe { select_in_chunk_vbmi2(chunk, nth) };
+    #[cfg(target_arch = "x86_64")]
+    {
+        static SELECTED: CpuDispatch<fn(&[u8; 64], usize) -> usize> = CpuDispatch::new(resolve);
+
+        #[cold]
+        fn resolve(chunk: &[u8; 64], nth: usize) -> usize {
+            let kernel: fn(&[u8; 64], usize) -> usize = if is_x86_feature_detected!("avx512f")
+                && is_x86_feature_detected!("avx512vpopcntdq")
+                && is_x86_feature_detected!("avx512vbmi2")
+            {
+                // SAFETY: runtime detection guarantees the required target features.
+                |chunk, nth| unsafe { select_in_chunk_vbmi2(chunk, nth) }
+            } else {
+                select_in_chunk_scalar
+            };
+            SELECTED.set(kernel);
+            kernel(chunk, nth)
         }
 
-        select_in_chunk_scalar
+        SELECTED.get()(chunk, nth)
+    }
+    #[cfg(not(target_arch = "x86_64"))]
+    {
+        select_in_chunk_scalar(chunk, nth)
     }
 }
 
@@ -371,22 +388,29 @@ 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 {
-    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 |word, nth| unsafe { select_in_word_bmi2(word, nth) };
+    #[cfg(target_arch = "x86_64")]
+    {
+        static SELECTED: CpuDispatch<fn(u64, usize) -> usize> = CpuDispatch::new(resolve);
+
+        #[cold]
+        fn resolve(word: u64, nth: usize) -> usize {
+            let kernel: fn(u64, usize) -> usize = if is_x86_feature_detected!("bmi2") {
+                // SAFETY: runtime detection guarantees the required target feature.
+                |word, nth| unsafe { select_in_word_bmi2(word, nth) }
+            } else {
+                select_in_word_scalar
+            };
+            SELECTED.set(kernel);
+            kernel(word, nth)
         }
-        select_in_word_scalar
+
+        SELECTED.get()(word, nth)
+    }
+    #[cfg(not(target_arch = "x86_64"))]
+    {
+        select_in_word_scalar(word, nth)
     }
 }