From 841d0e416c25b5928b3eb23de5076c4d1cffc288 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 14 Jul 2026 06:23:07 +0000 Subject: [PATCH] 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, 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 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 Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01DUFxXZLv4xxEqEYcoGWhfJ --- encodings/fastlanes/src/bit_transpose/mod.rs | 88 +++++---- .../src/arrays/bool/compute/filter.rs | 20 +- vortex-buffer/Cargo.toml | 4 + vortex-buffer/benches/cpu_dispatch.rs | 160 ++++++++++++++++ vortex-buffer/benches/vortex_bitbuffer.rs | 6 + vortex-buffer/src/bit/count_ones.rs | 66 +++++-- vortex-buffer/src/bit/select.rs | 97 ++++++---- vortex-buffer/src/dispatch.rs | 177 ++++++++++++++++++ vortex-buffer/src/lib.rs | 2 + vortex-mask/benches/rank.rs | 5 + vortex-mask/src/intersect_by_rank.rs | 35 ++-- 11 files changed, 556 insertions(+), 104 deletions(-) create mode 100644 vortex-buffer/benches/cpu_dispatch.rs create mode 100644 vortex-buffer/src/dispatch.rs diff --git a/encodings/fastlanes/src/bit_transpose/mod.rs b/encodings/fastlanes/src/bit_transpose/mod.rs index 99b01bc242c..26e8b595195 100644 --- a/encodings/fastlanes/src/bit_transpose/mod.rs +++ b/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 = 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 = 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)] diff --git a/vortex-array/src/arrays/bool/compute/filter.rs b/vortex-array/src/arrays/bool/compute/filter.rs index 2759c6e3077..98f52ed11ee 100644 --- a/vortex-array/src/arrays/bool/compute/filter.rs +++ b/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 = 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. diff --git a/vortex-buffer/Cargo.toml b/vortex-buffer/Cargo.toml index ae9d7e6cc05..f5cd04ad4fc 100644 --- a/vortex-buffer/Cargo.toml +++ b/vortex-buffer/Cargo.toml @@ -48,3 +48,7 @@ harness = false [[bench]] name = "vortex_bitbuffer" harness = false + +[[bench]] +name = "cpu_dispatch" +harness = false diff --git a/vortex-buffer/benches/cpu_dispatch.rs b/vortex-buffer/benches/cpu_dispatch.rs new file mode 100644 index 00000000000..6891dd22119 --- /dev/null +++ b/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` — 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 = 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 = + 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 = 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); +} diff --git a/vortex-buffer/benches/vortex_bitbuffer.rs b/vortex-buffer/benches/vortex_bitbuffer.rs index 923ef2aa98e..b67b3f1bf39 100644 --- a/vortex-buffer/benches/vortex_bitbuffer.rs +++ b/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(); } diff --git a/vortex-buffer/src/bit/count_ones.rs b/vortex-buffer/src/bit/count_ones.rs index df5844a2914..b977600cbe5 100644 --- a/vortex-buffer/src/bit/count_ones.rs +++ b/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,66 @@ 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 = 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") + && is_x86_feature_detected!("avx2") + { + 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, AVX-512VPOPCNTDQ, and AVX2. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx512f,avx512vpopcntdq,avx2")] +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")] +#[target_feature(enable = "avx2")] +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) } diff --git a/vortex-buffer/src/bit/select.rs b/vortex-buffer/src/bit/select.rs index 53eded23dd4..deae368ddfb 100644 --- a/vortex-buffer/src/bit/select.rs +++ b/vortex-buffer/src/bit/select.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use super::count_ones::align_offset_len; +use crate::dispatch::CpuKernel; /// Returns the position of the `nth` set bit (0-indexed) within the logical range /// `[offset, offset + len)` of the given byte slice. @@ -82,32 +83,47 @@ pub fn bit_select(bytes: &[u8], offset: usize, len: usize, nth: usize) -> Option /// /// If `chunk_index < chunks.len()`, the target bit is inside that chunk and `remaining` /// 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) -} +type ScanChunks = unsafe fn(&[[u8; 64]], usize, usize) -> (usize, usize, usize); -#[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")] #[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) }; +fn scan_chunks(chunks: &[[u8; 64]], remaining: usize, pos: usize) -> (usize, usize, usize) { + // Scans of a couple of chunks don't amortize the dispatch indirection: call the + // per-architecture unconditional kernel directly so it stays inlinable (see the + // size-gating note in the CpuKernel docs). + if chunks.len() <= 2 { + #[cfg(target_arch = "aarch64")] + return scan_chunks_neon(chunks, remaining, pos); + #[allow(unreachable_code)] + { + return scan_chunks_scalar(chunks, remaining, pos); + } } - scan_chunks_scalar(chunks, remaining, pos) + static KERNEL: CpuKernel = CpuKernel::new(|| { + #[cfg(target_arch = "x86_64")] + { + if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512vpopcntdq") { + return scan_chunks_avx512_vpopcnt; + } + } + #[cfg(target_arch = "aarch64")] + return scan_chunks_neon; + // The aarch64 arm above returns unconditionally (NEON needs no probe), making + // this portable default unreachable there. + #[allow(unreachable_code)] + { + scan_chunks_scalar + } + }); + // SAFETY: the selector only returns kernels that are safe or whose required CPU + // features were probed before selection. + unsafe { KERNEL.get()(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, @@ -185,7 +201,6 @@ unsafe fn scan_chunks_avx512_vpopcnt( (remaining, pos, chunks.len()) } -#[cfg(not(target_arch = "aarch64"))] #[inline] fn scan_chunks_scalar( chunks: &[[u8; 64]], @@ -280,21 +295,26 @@ fn scan_words_scalar( // ── In-chunk select ───────────────────────────────────────────────────── +type SelectInChunk = unsafe fn(&[u8; 64], usize) -> usize; + /// Position of the `nth` set bit inside a 64-byte chunk (0-indexed). #[inline] fn select_in_chunk(chunk: &[u8; 64], nth: usize) -> usize { - #[cfg(target_arch = "x86_64")] - { - if is_x86_feature_detected!("avx512f") - && is_x86_feature_detected!("avx512vpopcntdq") - && is_x86_feature_detected!("avx512vbmi2") + static KERNEL: CpuKernel = CpuKernel::new(|| { + #[cfg(target_arch = "x86_64")] { - // SAFETY: runtime detection guarantees the required target features. - return unsafe { select_in_chunk_vbmi2(chunk, nth) }; + if is_x86_feature_detected!("avx512f") + && is_x86_feature_detected!("avx512vpopcntdq") + && is_x86_feature_detected!("avx512vbmi2") + { + return select_in_chunk_vbmi2; + } } - } - - select_in_chunk_scalar(chunk, nth) + select_in_chunk_scalar + }); + // SAFETY: the selector only returns kernels that are safe or whose required CPU + // features were probed before selection. + unsafe { KERNEL.get()(chunk, nth) } } #[cfg(target_arch = "x86_64")] @@ -343,7 +363,6 @@ fn select_in_chunk_scalar(chunk: &[u8; 64], mut nth: usize) -> usize { unreachable!("select_in_chunk: nth exceeds popcount") } -#[cfg(not(target_arch = "aarch64"))] #[inline] fn count_ones_chunk(chunk: &[u8; 64]) -> usize { let words = chunk.as_chunks::<8>().0; @@ -359,17 +378,23 @@ fn count_ones_chunk(chunk: &[u8; 64]) -> usize { // ── In-word select ────────────────────────────────────────────────────── +type SelectInWord = unsafe fn(u64, usize) -> usize; + /// Position of the `nth` set bit inside a u64 (0-indexed, little-endian bit order). #[inline] fn select_in_word(word: u64, nth: usize) -> usize { - #[cfg(target_arch = "x86_64")] - { - if is_x86_feature_detected!("bmi2") { - // SAFETY: runtime detection guarantees the required target feature. - return unsafe { select_in_word_bmi2(word, nth) }; + static KERNEL: CpuKernel = CpuKernel::new(|| { + #[cfg(target_arch = "x86_64")] + { + if is_x86_feature_detected!("bmi2") { + return select_in_word_bmi2; + } } - } - select_in_word_scalar(word, nth) + select_in_word_scalar + }); + // SAFETY: the selector only returns kernels that are safe or whose required CPU + // features were probed before selection. + unsafe { KERNEL.get()(word, nth) } } /// BMI2: deposit a single bit at the nth set-bit position, then count trailing zeros. diff --git a/vortex-buffer/src/dispatch.rs b/vortex-buffer/src/dispatch.rs new file mode 100644 index 00000000000..6040bf0c885 --- /dev/null +++ b/vortex-buffer/src/dispatch.rs @@ -0,0 +1,177 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! One-time CPU-feature-based function dispatch. +//! +//! [`CpuKernel`] holds a function pointer chosen by a *selector* exactly once, on the +//! first call. Every later call is a relaxed atomic load, a never-taken predicted +//! branch, and an indirect call — measured indistinguishable from a direct call (see +//! `benches/cpu_dispatch.rs`). +//! +//! The selector is ordinary code that models both dispatch dimensions: +//! +//! * **Compile time** (x86_64 vs aarch64): one `#[cfg(target_arch = ...)]` block per +//! architecture, each *returning early* with its chosen kernel. +//! * **Runtime** (AVX-512 vs BMI2 vs ...): an if-chain of feature probes inside that +//! block. +//! * The **portable default** is the plain tail expression. Because the architecture +//! arms return early, no `#[cfg(not(any(...)))]` negation is ever needed. An arm +//! that returns unconditionally (e.g. NEON, architecturally guaranteed on aarch64) +//! makes the tail unreachable on that architecture — wrap just the tail in an +//! `#[allow(unreachable_code)]` block. +//! +//! When kernels are `unsafe` `#[target_feature]` functions, make `F` an +//! `unsafe fn(...)` pointer type: the bare kernel names then coerce directly (no +//! closure wrappers), and the one dispatched call is wrapped in `unsafe` with a +//! SAFETY comment stating that the selector probed the required features. +//! +//! Races are benign: the slot only ever holds valid function pointers of the same +//! type, and every candidate must compute the same result. +//! +//! # 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`. +//! +//! For the same reason, gate on input size *before* [`get`](CpuKernel::get) when tiny +//! inputs are common: call the portable kernel directly below the size where SIMD pays +//! off, so those calls stay inlinable and skip the dispatch entirely (see +//! `count_ones_aligned`). + +use core::mem::transmute_copy; +use core::ptr; +use core::sync::atomic::AtomicPtr; +use core::sync::atomic::Ordering; + +/// A function pointer selected by CPU-feature detection once, on first use. +/// +/// `F` must be a plain function-pointer type (`fn(...) -> ...`). The selector passed +/// to [`new`](Self::new) runs at most once per process (once per racing thread in the +/// worst case), on the first [`get`](Self::get), and its result is cached. +/// Non-capturing closures coerce to function pointers, so the selector can be written +/// inline in the `static`, like `LazyLock`. +/// +/// # Example +/// +/// ``` +/// use vortex_buffer::CpuKernel; +/// +/// /// Sums a slice, using the best kernel for the current CPU. +/// fn sum(values: &[u64]) -> u64 { +/// static KERNEL: CpuKernel u64> = CpuKernel::new(|| { +/// // Compile-time arm per architecture; runtime probes inside it. +/// #[cfg(target_arch = "x86_64")] +/// { +/// if std::arch::is_x86_feature_detected!("avx2") { +/// // return |values| unsafe { sum_avx2(values) }; +/// } +/// } +/// // Portable default: plain tail, no cfg(not(...)) needed. +/// |values| values.iter().sum() +/// }); +/// KERNEL.get()(values) +/// } +/// +/// assert_eq!(sum(&[1, 2, 3]), 6); +/// ``` +pub struct CpuKernel { + selected: AtomicPtr<()>, + select: fn() -> F, +} + +impl CpuKernel { + /// Create a kernel slot whose kernel is chosen by `select` on the first + /// [`get`](Self::get). + pub const fn new(select: fn() -> F) -> Self { + assert!( + size_of::() == size_of::<*mut ()>(), + "CpuKernel requires a function-pointer type" + ); + Self { + selected: AtomicPtr::new(ptr::null_mut()), + select, + } + } + + /// Return the selected kernel, running the selector on the first call. + /// + /// Steady state is a relaxed load plus a never-taken predicted branch. + #[inline] + pub fn get(&self) -> F { + let fn_ptr = self.selected.load(Ordering::Relaxed); + if fn_ptr.is_null() { + return self.select_slow(); + } + // SAFETY: non-null values in `selected` are always the bits of an `F` stored + // by `select_slow`, and `F` is pointer-sized (asserted in `new`). Function + // pointers are never null, so the null sentinel stays unambiguous. + unsafe { transmute_copy::<*mut (), F>(&fn_ptr) } + } + + #[cold] + fn select_slow(&self) -> F { + let kernel = (self.select)(); + // SAFETY: `F` is pointer-sized (asserted in `new`); a bitwise copy into a raw + // pointer preserves the function pointer for the transmute back in `get`. + let fn_ptr = unsafe { transmute_copy::(&kernel) }; + self.selected.store(fn_ptr, Ordering::Relaxed); + kernel + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + use super::CpuKernel; + + static SELECT_CALLS: AtomicUsize = AtomicUsize::new(0); + + fn add_one(x: u64) -> u64 { + static KERNEL: CpuKernel u64> = CpuKernel::new(|| { + SELECT_CALLS.fetch_add(1, Ordering::Relaxed); + |x| x + 1 + }); + KERNEL.get()(x) + } + + #[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 selector_early_returns_skip_the_default() { + static KERNEL: CpuKernel u64> = CpuKernel::new(|| { + if 1 + 1 == 2 { + return |x| x + 2; + } + |x| x + 100 + }); + assert_eq!(KERNEL.get()(0), 2); + } + + type XorFn = fn(&[u8; 4], &mut [u8; 4]); + + #[test] + fn supports_reference_arguments() { + static KERNEL: CpuKernel = CpuKernel::new(|| { + |input, output| { + for (o, i) in output.iter_mut().zip(input) { + *o ^= *i; + } + } + }); + let selected = KERNEL.get(); + let input = [1, 2, 3, 4]; + let mut output = [0, 0, 0, 0]; + selected(&input, &mut output); + assert_eq!(output, input); + } +} diff --git a/vortex-buffer/src/lib.rs b/vortex-buffer/src/lib.rs index 8319fffa387..ee113481353 100644 --- a/vortex-buffer/src/lib.rs +++ b/vortex-buffer/src/lib.rs @@ -52,6 +52,7 @@ pub use buffer::*; pub use buffer_mut::*; pub use bytes::*; pub use r#const::*; +pub use dispatch::*; pub use string::*; mod alignment; #[cfg(feature = "arrow")] @@ -62,6 +63,7 @@ mod buffer_mut; mod bytes; mod r#const; mod debug; +mod dispatch; mod macros; #[cfg(feature = "memmap2")] mod memmap2; diff --git a/vortex-mask/benches/rank.rs b/vortex-mask/benches/rank.rs index cd22f9ffd1f..fa2b820f75d 100644 --- a/vortex-mask/benches/rank.rs +++ b/vortex-mask/benches/rank.rs @@ -10,6 +10,11 @@ use vortex_buffer::BitBuffer; use vortex_mask::Mask; fn main() { + // Pre-resolve the one-time CpuKernel selections on the rank/select path so no + // benchmark iteration pays the first-call cost. + let warm = create_mask(512, 0.5); + let _ = warm.rank(warm.true_count() / 2); + divan::main(); } diff --git a/vortex-mask/src/intersect_by_rank.rs b/vortex-mask/src/intersect_by_rank.rs index 9b7da4a35dc..e5a8d6a08e5 100644 --- a/vortex-mask/src/intersect_by_rank.rs +++ b/vortex-mask/src/intersect_by_rank.rs @@ -9,6 +9,7 @@ use std::sync::Arc; use vortex_buffer::BitBuffer; use vortex_buffer::BitChunkIterator; use vortex_buffer::BufferMut; +use vortex_buffer::CpuKernel; use vortex_error::VortexExpect; use crate::Mask; @@ -396,22 +397,32 @@ fn intersect_bit_buffers_dispatch( mask_buffer: &BitBuffer, true_count: usize, ) -> Mask { - #[cfg(target_arch = "x86_64")] - if std::arch::is_x86_feature_detected!("bmi2") { - return intersect_bit_buffers::(self_buffer, mask_buffer, true_count); - } - - intersect_bit_buffers::(self_buffer, mask_buffer, true_count) + type IntersectBuffers = fn(&BitBuffer, &BitBuffer, usize) -> Mask; + static KERNEL: CpuKernel = CpuKernel::new(|| { + #[cfg(target_arch = "x86_64")] + { + if std::arch::is_x86_feature_detected!("bmi2") { + return intersect_bit_buffers::; + } + } + intersect_bit_buffers:: + }); + KERNEL.get()(self_buffer, mask_buffer, true_count) } #[inline] fn intersect_rank_indices_dispatch(self_buffer: &BitBuffer, mask_indices: &[usize]) -> Mask { - #[cfg(target_arch = "x86_64")] - if std::arch::is_x86_feature_detected!("bmi2") { - return intersect_bit_buffer_by_rank_indices::(self_buffer, mask_indices); - } - - intersect_bit_buffer_by_rank_indices::(self_buffer, mask_indices) + type IntersectRankIndices = fn(&BitBuffer, &[usize]) -> Mask; + static KERNEL: CpuKernel = CpuKernel::new(|| { + #[cfg(target_arch = "x86_64")] + { + if std::arch::is_x86_feature_detected!("bmi2") { + return intersect_bit_buffer_by_rank_indices::; + } + } + intersect_bit_buffer_by_rank_indices:: + }); + KERNEL.get()(self_buffer, mask_indices) } #[inline]