Skip to content

Commit d64f919

Browse files
claudeclaude
committed
refactor: selector-style CpuKernel with early-return arch arms
Replace the Candidate-table API with a single selector fn per site, written as: one #[cfg(target_arch)] block per architecture that returns early (runtime feature probes as an if-chain inside it), followed by the portable default as the plain tail expression. Because the arch arms return early, no #[cfg(not(any(...)))] negation block is needed anywhere; an arm that returns unconditionally (aarch64/NEON) makes the tail unreachable on that architecture and carries a one-line #[allow(unreachable_code)]. The scalar kernels in bit/select.rs lose their cfg(not(aarch64)) guards since the shared tail now references them on every architecture. CpuDispatch (branchless resolver variant) is dropped: the benchmark shows its advantage over CpuKernel is at most ~0.2ns/call, below run-to-run noise. Benchmark updated to the selector API; fastest-of-3 per 1024 calls: direct 2.50us, one-feature-probe-per-call 2.44-2.58us, CpuKernel 2.74us (conditional and unconditional selectors alike), LazyLock<fn> 3.64us, three-probes-per-call 4.85us. 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
1 parent ee17659 commit d64f919

4 files changed

Lines changed: 169 additions & 386 deletions

File tree

‎encodings/fastlanes/src/bit_transpose/mod.rs‎

Lines changed: 46 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -28,71 +28,11 @@ mod x86;
2828
mod validity;
2929

3030
pub use validity::*;
31-
use vortex_buffer::Candidate;
3231
use vortex_buffer::CpuKernel;
3332

3433
/// Signature shared by all 1024-bit transpose kernels.
3534
type TransposeKernel = fn(&[u8; 128], &mut [u8; 128]);
3635

37-
// Runtime dimension: kernels that need a CPU probe, best first.
38-
#[cfg(target_arch = "x86_64")]
39-
const TRANSPOSE_CANDIDATES: &[Candidate<TransposeKernel>] = &[
40-
Candidate {
41-
available: x86::has_vbmi,
42-
kernel: |input, output| {
43-
// SAFETY: `available` confirms VBMI.
44-
unsafe { x86::transpose_bits_vbmi(input, output) }
45-
},
46-
},
47-
Candidate {
48-
available: x86::has_bmi2,
49-
kernel: |input, output| {
50-
// SAFETY: `available` confirms BMI2.
51-
unsafe { x86::transpose_bits_bmi2(input, output) }
52-
},
53-
},
54-
];
55-
#[cfg(not(target_arch = "x86_64"))]
56-
const TRANSPOSE_CANDIDATES: &[Candidate<TransposeKernel>] = &[];
57-
58-
#[cfg(target_arch = "x86_64")]
59-
const UNTRANSPOSE_CANDIDATES: &[Candidate<TransposeKernel>] = &[
60-
Candidate {
61-
available: x86::has_vbmi,
62-
kernel: |input, output| {
63-
// SAFETY: `available` confirms VBMI.
64-
unsafe { x86::untranspose_bits_vbmi(input, output) }
65-
},
66-
},
67-
Candidate {
68-
available: x86::has_bmi2,
69-
kernel: |input, output| {
70-
// SAFETY: `available` confirms BMI2.
71-
unsafe { x86::untranspose_bits_bmi2(input, output) }
72-
},
73-
},
74-
];
75-
#[cfg(not(target_arch = "x86_64"))]
76-
const UNTRANSPOSE_CANDIDATES: &[Candidate<TransposeKernel>] = &[];
77-
78-
// Compile-time dimension: the per-architecture unconditional kernel. NEON is
79-
// architecturally guaranteed on aarch64, so it needs no probe.
80-
#[cfg(target_arch = "aarch64")]
81-
const TRANSPOSE_FALLBACK: TransposeKernel = |input, output| {
82-
// SAFETY: NEON is architecturally guaranteed on aarch64.
83-
unsafe { aarch64::transpose_bits_neon(input, output) }
84-
};
85-
#[cfg(not(target_arch = "aarch64"))]
86-
const TRANSPOSE_FALLBACK: TransposeKernel = scalar::transpose_bits_scalar;
87-
88-
#[cfg(target_arch = "aarch64")]
89-
const UNTRANSPOSE_FALLBACK: TransposeKernel = |input, output| {
90-
// SAFETY: NEON is architecturally guaranteed on aarch64.
91-
unsafe { aarch64::untranspose_bits_neon(input, output) }
92-
};
93-
#[cfg(not(target_arch = "aarch64"))]
94-
const UNTRANSPOSE_FALLBACK: TransposeKernel = scalar::untranspose_bits_scalar;
95-
9636
/// Base indices for the first 64 output bytes (lanes 0-7).
9737
/// Each entry indicates the starting input byte index for that output byte group.
9838
/// 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]
@@ -110,19 +50,63 @@ const TRANSPOSE_8X8: u64 = 0x0000_0000_F0F0_F0F0;
11050
/// Transpose 1024-bits into FastLanes layout.
11151
///
11252
/// Dispatch to the best available implementation, selected once on first call.
53+
// The aarch64 arm returns unconditionally (NEON needs no probe), making the portable
54+
// tail unreachable there.
55+
#[allow(unreachable_code)]
11356
#[inline]
11457
pub fn transpose_bits(input: &[u8; 128], output: &mut [u8; 128]) {
11558
static KERNEL: CpuKernel<TransposeKernel> = CpuKernel::new();
116-
KERNEL.get(TRANSPOSE_CANDIDATES, TRANSPOSE_FALLBACK)(input, output)
59+
KERNEL.get(|| {
60+
#[cfg(target_arch = "x86_64")]
61+
{
62+
// VBMI is fastest
63+
if x86::has_vbmi() {
64+
// SAFETY: the probe above confirms VBMI.
65+
return |input, output| unsafe { x86::transpose_bits_vbmi(input, output) };
66+
}
67+
if x86::has_bmi2() {
68+
// SAFETY: the probe above confirms BMI2.
69+
return |input, output| unsafe { x86::transpose_bits_bmi2(input, output) };
70+
}
71+
}
72+
#[cfg(target_arch = "aarch64")]
73+
return |input, output| {
74+
// SAFETY: NEON is architecturally guaranteed on aarch64.
75+
unsafe { aarch64::transpose_bits_neon(input, output) }
76+
};
77+
scalar::transpose_bits_scalar
78+
})(input, output)
11779
}
11880

11981
/// Untranspose 1024-bits from FastLanes layout.
12082
///
12183
/// Dispatch untranspose to the best available implementation, selected once on first call.
84+
// The aarch64 arm returns unconditionally (NEON needs no probe), making the portable
85+
// tail unreachable there.
86+
#[allow(unreachable_code)]
12287
#[inline]
12388
pub fn untranspose_bits(input: &[u8; 128], output: &mut [u8; 128]) {
12489
static KERNEL: CpuKernel<TransposeKernel> = CpuKernel::new();
125-
KERNEL.get(UNTRANSPOSE_CANDIDATES, UNTRANSPOSE_FALLBACK)(input, output)
90+
KERNEL.get(|| {
91+
#[cfg(target_arch = "x86_64")]
92+
{
93+
// VBMI is fastest
94+
if x86::has_vbmi() {
95+
// SAFETY: the probe above confirms VBMI.
96+
return |input, output| unsafe { x86::untranspose_bits_vbmi(input, output) };
97+
}
98+
if x86::has_bmi2() {
99+
// SAFETY: the probe above confirms BMI2.
100+
return |input, output| unsafe { x86::untranspose_bits_bmi2(input, output) };
101+
}
102+
}
103+
#[cfg(target_arch = "aarch64")]
104+
return |input, output| {
105+
// SAFETY: NEON is architecturally guaranteed on aarch64.
106+
unsafe { aarch64::untranspose_bits_neon(input, output) }
107+
};
108+
scalar::untranspose_bits_scalar
109+
})(input, output)
126110
}
127111

128112
#[cfg(test)]

‎vortex-buffer/benches/cpu_dispatch.rs‎

Lines changed: 9 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,9 @@
77
//! difference between rows is pure dispatch overhead:
88
//!
99
//! * `direct`: plain call — the floor.
10-
//! * `cpu_kernel_empty_candidates`: [`CpuKernel`] with an empty candidate table
11-
//! (the aarch64 shape) — should const-fold to `direct`.
12-
//! * `cpu_dispatch_branchless`: [`CpuDispatch`] — relaxed load + indirect call.
1310
//! * `cpu_kernel`: [`CpuKernel`] — relaxed load + null check + indirect call.
11+
//! * `cpu_kernel_unconditional`: a selector with no probes (the aarch64 shape) —
12+
//! same steady state as `cpu_kernel`.
1413
//! * `lazy_lock`: `std::sync::LazyLock<fn>` — Once-state check + value load +
1514
//! indirect call.
1615
//! * `feature_detect_each_call_1/_3`: the pre-`CpuKernel` pattern — 1 or 3
@@ -21,8 +20,6 @@ use std::sync::LazyLock;
2120
use divan::Bencher;
2221
use divan::black_box;
2322
use divan::counter::ItemsCount;
24-
use vortex_buffer::Candidate;
25-
use vortex_buffer::CpuDispatch;
2623
use vortex_buffer::CpuKernel;
2724

2825
fn main() {
@@ -31,8 +28,7 @@ fn main() {
3128
let seed = black_box(7);
3229
let _ = via_direct(1, seed)
3330
+ via_cpu_kernel(1, seed)
34-
+ via_cpu_kernel_empty_candidates(1, seed)
35-
+ via_cpu_dispatch(1, seed)
31+
+ via_cpu_kernel_unconditional(1, seed)
3632
+ via_lazy_lock(1, seed);
3733
#[cfg(target_arch = "x86_64")]
3834
let _ = via_feature_detect_1(1, seed) + via_feature_detect_3(1, seed);
@@ -102,37 +98,18 @@ fn via_lazy_lock(x: u64, seed: u64) -> u64 {
10298
(*LAZY)(x, seed)
10399
}
104100

105-
const CPU_KERNEL_CANDIDATES: &[Candidate<Kernel>] = &[Candidate {
106-
available: has_bmi2,
107-
kernel: kernel_a,
108-
}];
109-
110101
static CPU_KERNEL: CpuKernel<Kernel> = CpuKernel::new();
111102

112103
#[inline(never)]
113104
fn via_cpu_kernel(x: u64, seed: u64) -> u64 {
114-
CPU_KERNEL.get(CPU_KERNEL_CANDIDATES, kernel_b)(x, seed)
115-
}
116-
117-
static CPU_KERNEL_EMPTY: CpuKernel<Kernel> = CpuKernel::new();
118-
119-
#[inline(never)]
120-
fn via_cpu_kernel_empty_candidates(x: u64, seed: u64) -> u64 {
121-
CPU_KERNEL_EMPTY.get(&[], kernel_a)(x, seed)
105+
CPU_KERNEL.get(|| if has_bmi2() { kernel_a } else { kernel_b })(x, seed)
122106
}
123107

124-
static CPU_DISPATCH: CpuDispatch<Kernel> = CpuDispatch::new(resolve);
125-
126-
#[cold]
127-
fn resolve(x: u64, seed: u64) -> u64 {
128-
let kernel: Kernel = if has_bmi2() { kernel_a } else { kernel_b };
129-
CPU_DISPATCH.set(kernel);
130-
kernel(x, seed)
131-
}
108+
static CPU_KERNEL_UNCONDITIONAL: CpuKernel<Kernel> = CpuKernel::new();
132109

133110
#[inline(never)]
134-
fn via_cpu_dispatch(x: u64, seed: u64) -> u64 {
135-
CPU_DISPATCH.get()(x, seed)
111+
fn via_cpu_kernel_unconditional(x: u64, seed: u64) -> u64 {
112+
CPU_KERNEL_UNCONDITIONAL.get(|| kernel_a)(x, seed)
136113
}
137114

138115
// ── benches ─────────────────────────────────────────────────────────────
@@ -155,13 +132,8 @@ fn direct(bencher: Bencher) {
155132
}
156133

157134
#[divan::bench]
158-
fn cpu_kernel_empty_candidates(bencher: Bencher) {
159-
bench_via(bencher, via_cpu_kernel_empty_candidates);
160-
}
161-
162-
#[divan::bench]
163-
fn cpu_dispatch_branchless(bencher: Bencher) {
164-
bench_via(bencher, via_cpu_dispatch);
135+
fn cpu_kernel_unconditional(bencher: Bencher) {
136+
bench_via(bencher, via_cpu_kernel_unconditional);
165137
}
166138

167139
#[divan::bench]

‎vortex-buffer/src/bit/select.rs‎

Lines changed: 40 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
// SPDX-FileCopyrightText: Copyright the Vortex contributors
33

44
use super::count_ones::align_offset_len;
5-
use crate::dispatch::Candidate;
65
use crate::dispatch::CpuKernel;
76

87
/// Returns the position of the `nth` set bit (0-indexed) within the logical range
@@ -86,30 +85,26 @@ pub fn bit_select(bytes: &[u8], offset: usize, len: usize, nth: usize) -> Option
8685
/// is the rank *within* that chunk. Otherwise all chunks were consumed.
8786
type ScanChunks = fn(&[[u8; 64]], usize, usize) -> (usize, usize, usize);
8887

89-
// Runtime dimension: kernels that need a CPU probe, best first.
90-
#[cfg(target_arch = "x86_64")]
91-
const SCAN_CHUNKS_CANDIDATES: &[Candidate<ScanChunks>] = &[Candidate {
92-
available: || {
93-
is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512vpopcntdq")
94-
},
95-
kernel: |chunks, remaining, pos| {
96-
// SAFETY: `available` confirms the required target features.
97-
unsafe { scan_chunks_avx512_vpopcnt(chunks, remaining, pos) }
98-
},
99-
}];
100-
#[cfg(not(target_arch = "x86_64"))]
101-
const SCAN_CHUNKS_CANDIDATES: &[Candidate<ScanChunks>] = &[];
102-
103-
// Compile-time dimension: the per-architecture unconditional kernel.
104-
#[cfg(target_arch = "aarch64")]
105-
const SCAN_CHUNKS_FALLBACK: ScanChunks = scan_chunks_neon;
106-
#[cfg(not(target_arch = "aarch64"))]
107-
const SCAN_CHUNKS_FALLBACK: ScanChunks = scan_chunks_scalar;
108-
88+
// The aarch64 arm returns unconditionally (NEON needs no probe), making the portable
89+
// tail unreachable there.
90+
#[allow(unreachable_code)]
10991
#[inline]
11092
fn scan_chunks(chunks: &[[u8; 64]], remaining: usize, pos: usize) -> (usize, usize, usize) {
11193
static KERNEL: CpuKernel<ScanChunks> = CpuKernel::new();
112-
KERNEL.get(SCAN_CHUNKS_CANDIDATES, SCAN_CHUNKS_FALLBACK)(chunks, remaining, pos)
94+
KERNEL.get(|| {
95+
#[cfg(target_arch = "x86_64")]
96+
{
97+
if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512vpopcntdq") {
98+
// SAFETY: the probe above confirms the required target features.
99+
return |chunks, remaining, pos| unsafe {
100+
scan_chunks_avx512_vpopcnt(chunks, remaining, pos)
101+
};
102+
}
103+
}
104+
#[cfg(target_arch = "aarch64")]
105+
return scan_chunks_neon;
106+
scan_chunks_scalar
107+
})(chunks, remaining, pos)
113108
}
114109

115110
#[cfg(target_arch = "aarch64")]
@@ -193,7 +188,6 @@ unsafe fn scan_chunks_avx512_vpopcnt(
193188
(remaining, pos, chunks.len())
194189
}
195190

196-
#[cfg(not(target_arch = "aarch64"))]
197191
#[inline]
198192
fn scan_chunks_scalar(
199193
chunks: &[[u8; 64]],
@@ -290,26 +284,23 @@ fn scan_words_scalar(
290284

291285
type SelectInChunk = fn(&[u8; 64], usize) -> usize;
292286

293-
#[cfg(target_arch = "x86_64")]
294-
const SELECT_IN_CHUNK_CANDIDATES: &[Candidate<SelectInChunk>] = &[Candidate {
295-
available: || {
296-
is_x86_feature_detected!("avx512f")
297-
&& is_x86_feature_detected!("avx512vpopcntdq")
298-
&& is_x86_feature_detected!("avx512vbmi2")
299-
},
300-
kernel: |chunk, nth| {
301-
// SAFETY: `available` confirms the required target features.
302-
unsafe { select_in_chunk_vbmi2(chunk, nth) }
303-
},
304-
}];
305-
#[cfg(not(target_arch = "x86_64"))]
306-
const SELECT_IN_CHUNK_CANDIDATES: &[Candidate<SelectInChunk>] = &[];
307-
308287
/// Position of the `nth` set bit inside a 64-byte chunk (0-indexed).
309288
#[inline]
310289
fn select_in_chunk(chunk: &[u8; 64], nth: usize) -> usize {
311290
static KERNEL: CpuKernel<SelectInChunk> = CpuKernel::new();
312-
KERNEL.get(SELECT_IN_CHUNK_CANDIDATES, select_in_chunk_scalar)(chunk, nth)
291+
KERNEL.get(|| {
292+
#[cfg(target_arch = "x86_64")]
293+
{
294+
if is_x86_feature_detected!("avx512f")
295+
&& is_x86_feature_detected!("avx512vpopcntdq")
296+
&& is_x86_feature_detected!("avx512vbmi2")
297+
{
298+
// SAFETY: the probe above confirms the required target features.
299+
return |chunk, nth| unsafe { select_in_chunk_vbmi2(chunk, nth) };
300+
}
301+
}
302+
select_in_chunk_scalar
303+
})(chunk, nth)
313304
}
314305

315306
#[cfg(target_arch = "x86_64")]
@@ -358,7 +349,6 @@ fn select_in_chunk_scalar(chunk: &[u8; 64], mut nth: usize) -> usize {
358349
unreachable!("select_in_chunk: nth exceeds popcount")
359350
}
360351

361-
#[cfg(not(target_arch = "aarch64"))]
362352
#[inline]
363353
fn count_ones_chunk(chunk: &[u8; 64]) -> usize {
364354
let words = chunk.as_chunks::<8>().0;
@@ -376,22 +366,20 @@ fn count_ones_chunk(chunk: &[u8; 64]) -> usize {
376366

377367
type SelectInWord = fn(u64, usize) -> usize;
378368

379-
#[cfg(target_arch = "x86_64")]
380-
const SELECT_IN_WORD_CANDIDATES: &[Candidate<SelectInWord>] = &[Candidate {
381-
available: || is_x86_feature_detected!("bmi2"),
382-
kernel: |word, nth| {
383-
// SAFETY: `available` confirms the required target feature.
384-
unsafe { select_in_word_bmi2(word, nth) }
385-
},
386-
}];
387-
#[cfg(not(target_arch = "x86_64"))]
388-
const SELECT_IN_WORD_CANDIDATES: &[Candidate<SelectInWord>] = &[];
389-
390369
/// Position of the `nth` set bit inside a u64 (0-indexed, little-endian bit order).
391370
#[inline]
392371
fn select_in_word(word: u64, nth: usize) -> usize {
393372
static KERNEL: CpuKernel<SelectInWord> = CpuKernel::new();
394-
KERNEL.get(SELECT_IN_WORD_CANDIDATES, select_in_word_scalar)(word, nth)
373+
KERNEL.get(|| {
374+
#[cfg(target_arch = "x86_64")]
375+
{
376+
if is_x86_feature_detected!("bmi2") {
377+
// SAFETY: the probe above confirms the required target feature.
378+
return |word, nth| unsafe { select_in_word_bmi2(word, nth) };
379+
}
380+
}
381+
select_in_word_scalar
382+
})(word, nth)
395383
}
396384

397385
/// BMI2: deposit a single bit at the nth set-bit position, then count trailing zeros.

0 commit comments

Comments
 (0)