33
44//! One-time CPU-feature-based function dispatch.
55//!
6- //! [`CpuDispatch`] holds a function pointer that is re-pointed from a resolver to the
7- //! best kernel on the first call. Every later call is a single relaxed atomic load
8- //! plus an indirect call — no repeated feature checks, no branches, no `LazyLock`
9- //! initialized-state check.
6+ //! [`LazyCpuDispatch`] holds a function pointer chosen by runtime CPU-feature
7+ //! detection exactly once, on the first call. The selector — an ordinary
8+ //! `fn() -> F` containing only the feature checks — is stored at construction, so a
9+ //! call site is just a `static` and a [`get`](LazyCpuDispatch::get). Every call after
10+ //! the first is a relaxed atomic load, a never-taken predicted branch, and an
11+ //! indirect call.
1012//!
11- //! The slot is created pointing at a caller-written *resolver* with the same
12- //! signature as the kernels. On its first (and only) invocation the resolver detects
13- //! CPU features, [`set`](CpuDispatch::set)s the chosen kernel into the slot, and
14- //! forwards the call. Races are benign: the slot only ever holds valid function
15- //! pointers of the same type, and every candidate computes the same result.
13+ //! [`CpuDispatch`] is the branchless variant for the very hottest sites: steady state
14+ //! is a load plus an indirect call with no branch at all. The price is two lines of
15+ //! ceremony in a per-site resolver function (`set` the kernel, forward the first
16+ //! call). That ceremony is irreducible: branchless dispatch needs a distinct
17+ //! same-signature function per site that knows its slot's address, which generic
18+ //! code cannot synthesize — only a macro (or your hands) can write it.
19+ //!
20+ //! Races are benign in both: a slot only ever holds valid function pointers of the
21+ //! same type, and every candidate must compute the same result.
1622//!
1723//! # When to use it
1824//!
19- //! Use this for *coarse-grained* kernels (a whole chunk or buffer per call) that have
20- //! more than one runtime candidate on the target architecture, e.g. x86_64 kernels
21- //! with AVX-512 / BMI2 / scalar variants. Today that means x86_64: on aarch64 NEON is
22- //! architecturally guaranteed, so a plain `#[cfg(target_arch = "aarch64")]` direct
23- //! call is already zero-cost and inlinable — do not funnel it through a function
24- //! pointer. If SVE/SVE2 kernels are added later, aarch64 call sites get the same
25- //! one-time selection by adding `is_aarch64_feature_detected!` arms to the resolver.
25+ //! Use these for *coarse-grained* kernels (a whole chunk or buffer per call) that
26+ //! have more than one runtime candidate on the target architecture, e.g. x86_64
27+ //! kernels with AVX-512 / BMI2 / scalar variants. Today that means x86_64: on aarch64
28+ //! NEON is architecturally guaranteed, so a plain `#[cfg(target_arch = "aarch64")]`
29+ //! direct call is already zero-cost and inlinable — do not funnel it through a
30+ //! function pointer. If SVE/SVE2 kernels are added later, aarch64 call sites get the
31+ //! same one-time selection by adding `is_aarch64_feature_detected!` arms to the
32+ //! selector.
2633//!
2734//! # When NOT to use it
2835//!
3340
3441use core:: marker:: PhantomData ;
3542use core:: mem:: transmute_copy;
43+ use core:: ptr;
3644use core:: sync:: atomic:: AtomicPtr ;
3745use core:: sync:: atomic:: Ordering ;
3846
39- /// A self-patching function-pointer slot for one-time CPU-feature kernel selection.
47+ /// A function pointer selected by CPU-feature detection once, on first use.
48+ ///
49+ /// `F` must be a plain function-pointer type (`fn(...) -> ...`). The selector passed
50+ /// to [`new`](Self::new) contains only the feature checks and returns the chosen
51+ /// kernel; storage and forwarding are handled here. Non-capturing closures coerce to
52+ /// function pointers, so both the selector and `unsafe` `#[target_feature]` kernel
53+ /// wrappers can be written inline.
54+ ///
55+ /// # Example
56+ ///
57+ /// ```
58+ /// use vortex_buffer::LazyCpuDispatch;
59+ ///
60+ /// /// Sums a slice, using the best kernel for the current CPU.
61+ /// fn sum(values: &[u64]) -> u64 {
62+ /// static SELECTED: LazyCpuDispatch<fn(&[u64]) -> u64> = LazyCpuDispatch::new(|| {
63+ /// #[cfg(target_arch = "x86_64")]
64+ /// if std::arch::is_x86_feature_detected!("avx2") {
65+ /// // return |values| unsafe { sum_avx2(values) };
66+ /// }
67+ /// |values| values.iter().sum()
68+ /// });
69+ /// SELECTED.get()(values)
70+ /// }
71+ ///
72+ /// assert_eq!(sum(&[1, 2, 3]), 6);
73+ /// ```
74+ pub struct LazyCpuDispatch < F > {
75+ selected : AtomicPtr < ( ) > ,
76+ select : fn ( ) -> F ,
77+ }
78+
79+ impl < F : Copy > LazyCpuDispatch < F > {
80+ /// Create a slot whose kernel is chosen by `select` on the first
81+ /// [`get`](Self::get).
82+ pub const fn new ( select : fn ( ) -> F ) -> Self {
83+ assert ! (
84+ size_of:: <F >( ) == size_of:: <* mut ( ) >( ) ,
85+ "LazyCpuDispatch requires a function-pointer type"
86+ ) ;
87+ Self {
88+ selected : AtomicPtr :: new ( ptr:: null_mut ( ) ) ,
89+ select,
90+ }
91+ }
92+
93+ /// Return the selected kernel, running the selector on the first call.
94+ ///
95+ /// Steady state is a relaxed load plus a never-taken predicted branch.
96+ #[ inline]
97+ pub fn get ( & self ) -> F {
98+ let fn_ptr = self . selected . load ( Ordering :: Relaxed ) ;
99+ if fn_ptr. is_null ( ) {
100+ return self . select_slow ( ) ;
101+ }
102+ // SAFETY: non-null values in `selected` are always the bits of an `F` stored
103+ // by `select_slow`, and `F` is pointer-sized (asserted in `new`). Function
104+ // pointers are never null, so the null sentinel stays unambiguous.
105+ unsafe { transmute_copy :: < * mut ( ) , F > ( & fn_ptr) }
106+ }
107+
108+ #[ cold]
109+ fn select_slow ( & self ) -> F {
110+ let kernel = ( self . select ) ( ) ;
111+ // SAFETY: `F` is pointer-sized (asserted in `new`); a bitwise copy into a raw
112+ // pointer preserves the function pointer for the transmute back in `get`.
113+ let fn_ptr = unsafe { transmute_copy :: < F , * mut ( ) > ( & kernel) } ;
114+ self . selected . store ( fn_ptr, Ordering :: Relaxed ) ;
115+ kernel
116+ }
117+ }
118+
119+ /// A branchless self-patching function-pointer slot, for the hottest dispatch sites.
40120///
41121/// `F` must be a plain function-pointer type (`fn(...) -> ...`). Construct with a
42122/// resolver of that same signature which detects CPU features, stores the chosen
43- /// kernel via [`set`](Self::set), and forwards its arguments to it. Non-capturing
44- /// closures coerce to function pointers, so `unsafe` `#[target_feature]` kernels can
45- /// be wrapped inline after detection proves them sound .
123+ /// kernel via [`set`](Self::set), and forwards its arguments to it. Steady state is a
124+ /// relaxed load plus an indirect call — no branch, unlike [`LazyCpuDispatch`] — at
125+ /// the cost of writing that per-site resolver .
46126///
47127/// # Example
48128///
@@ -120,6 +200,40 @@ mod tests {
120200 use std:: sync:: atomic:: Ordering ;
121201
122202 use super :: CpuDispatch ;
203+ use super :: LazyCpuDispatch ;
204+
205+ static LAZY_SELECT_CALLS : AtomicUsize = AtomicUsize :: new ( 0 ) ;
206+
207+ static LAZY_ADD_ONE : LazyCpuDispatch < fn ( u64 ) -> u64 > = LazyCpuDispatch :: new ( || {
208+ LAZY_SELECT_CALLS . fetch_add ( 1 , Ordering :: Relaxed ) ;
209+ |x| x + 1
210+ } ) ;
211+
212+ type XorFn = fn ( & [ u8 ; 4 ] , & mut [ u8 ; 4 ] ) ;
213+
214+ static LAZY_XOR_BYTES : LazyCpuDispatch < XorFn > = LazyCpuDispatch :: new ( || {
215+ |input, output| {
216+ for ( o, i) in output. iter_mut ( ) . zip ( input) {
217+ * o ^= * i;
218+ }
219+ }
220+ } ) ;
221+
222+ #[ test]
223+ fn lazy_selects_once_then_dispatches ( ) {
224+ assert_eq ! ( LAZY_ADD_ONE . get( ) ( 41 ) , 42 ) ;
225+ assert_eq ! ( LAZY_ADD_ONE . get( ) ( 1 ) , 2 ) ;
226+ assert_eq ! ( LAZY_ADD_ONE . get( ) ( 2 ) , 3 ) ;
227+ assert_eq ! ( LAZY_SELECT_CALLS . load( Ordering :: Relaxed ) , 1 ) ;
228+ }
229+
230+ #[ test]
231+ fn lazy_supports_reference_arguments ( ) {
232+ let input = [ 1 , 2 , 3 , 4 ] ;
233+ let mut output = [ 0 , 0 , 0 , 0 ] ;
234+ LAZY_XOR_BYTES . get ( ) ( & input, & mut output) ;
235+ assert_eq ! ( output, input) ;
236+ }
123237
124238 static RESOLVE_CALLS : AtomicUsize = AtomicUsize :: new ( 0 ) ;
125239
@@ -132,33 +246,11 @@ mod tests {
132246 kernel ( x)
133247 }
134248
135- type XorFn = fn ( & [ u8 ; 4 ] , & mut [ u8 ; 4 ] ) ;
136-
137- static XOR_BYTES : CpuDispatch < XorFn > = CpuDispatch :: new ( resolve_xor) ;
138-
139- fn resolve_xor ( input : & [ u8 ; 4 ] , output : & mut [ u8 ; 4 ] ) {
140- let kernel: XorFn = |input, output| {
141- for ( o, i) in output. iter_mut ( ) . zip ( input) {
142- * o ^= * i;
143- }
144- } ;
145- XOR_BYTES . set ( kernel) ;
146- kernel ( input, output)
147- }
148-
149249 #[ test]
150250 fn resolves_once_then_dispatches ( ) {
151251 assert_eq ! ( ADD_ONE . get( ) ( 41 ) , 42 ) ;
152252 assert_eq ! ( ADD_ONE . get( ) ( 1 ) , 2 ) ;
153253 assert_eq ! ( ADD_ONE . get( ) ( 2 ) , 3 ) ;
154254 assert_eq ! ( RESOLVE_CALLS . load( Ordering :: Relaxed ) , 1 ) ;
155255 }
156-
157- #[ test]
158- fn supports_reference_arguments ( ) {
159- let input = [ 1 , 2 , 3 , 4 ] ;
160- let mut output = [ 0 , 0 , 0 , 0 ] ;
161- XOR_BYTES . get ( ) ( & input, & mut output) ;
162- assert_eq ! ( output, input) ;
163- }
164256}
0 commit comments