oxifft_codegen_impl/gen_simd/
runtime_dispatch.rs1use proc_macro2::TokenStream;
38use quote::{format_ident, quote};
39use syn::{
40 parse::{Parse, ParseStream},
41 LitInt, Token,
42};
43
44pub use super::multi_transform::Precision;
45
46#[derive(Debug, Clone, Copy)]
52pub struct DispatcherConfig {
53 pub size: usize,
55 pub precision: Precision,
57}
58
59pub const ISA_SCALAR: u8 = 0;
65pub const ISA_SSE2: u8 = 1;
67pub const ISA_AVX: u8 = 2;
69pub const ISA_AVX2_FMA: u8 = 3;
71pub const ISA_AVX512: u8 = 4;
73pub const ISA_NEON: u8 = 5;
75pub const ISA_UNDETECTED: u8 = 255;
77
78#[must_use]
89pub fn detect_host_isa() -> u8 {
90 #[cfg(target_arch = "x86_64")]
91 {
92 if is_x86_feature_detected!("avx512f") {
93 return ISA_AVX512;
94 }
95 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
96 return ISA_AVX2_FMA;
97 }
98 if is_x86_feature_detected!("avx") {
99 return ISA_AVX;
100 }
101 if is_x86_feature_detected!("sse2") {
102 return ISA_SSE2;
103 }
104 return ISA_SCALAR;
105 }
106
107 #[cfg(target_arch = "aarch64")]
108 {
109 if std::arch::is_aarch64_feature_detected!("neon") {
110 return ISA_NEON;
111 }
112 return ISA_SCALAR;
113 }
114
115 #[allow(unreachable_code)]
117 ISA_SCALAR
118}
119
120fn build_detect_x86_body() -> TokenStream {
126 quote! {
127 #[cfg(target_arch = "x86_64")]
128 {
129 #[cfg(feature = "avx512")]
130 if is_x86_feature_detected!("avx512f") {
131 return ISA_AVX512_LEVEL;
132 }
133 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
134 return ISA_AVX2_FMA_LEVEL;
135 }
136 if is_x86_feature_detected!("avx") {
137 return ISA_AVX_LEVEL;
138 }
139 if is_x86_feature_detected!("sse2") {
140 return ISA_SSE2_LEVEL;
141 }
142 return ISA_SCALAR_LEVEL;
143 }
144 }
145}
146
147fn build_detect_aarch64_body() -> TokenStream {
149 quote! {
150 #[cfg(target_arch = "aarch64")]
151 {
152 if std::arch::is_aarch64_feature_detected!("neon") {
153 return ISA_NEON_LEVEL;
154 }
155 return ISA_SCALAR_LEVEL;
156 }
157 }
158}
159
160fn build_x86_64_branches(config: DispatcherConfig) -> TokenStream {
168 let size = config.size;
169 let ty_str = config.precision.type_str();
170 let ty_tokens: TokenStream = ty_str
171 .parse()
172 .unwrap_or_else(|_| unreachable!("ty_str is always f32 or f64"));
173 let avx512_fn = format_ident!("codelet_simd_{}_avx512_{}", size, ty_str);
174 let avx2_fn = format_ident!("codelet_simd_{}_avx2_{}", size, ty_str);
175 let sse2_fn = format_ident!("codelet_simd_{}_sse2_{}", size, ty_str);
176
177 if size == 16 {
178 if config.precision == Precision::F32 {
179 return quote! {
180 #[cfg(all(target_arch = "x86_64", feature = "avx512"))]
181 {
182 if cached_level == ISA_AVX512_LEVEL {
183 let data_len = data.len() * 2;
186 let data_ptr = data.as_mut_ptr().cast::<#ty_tokens>();
187 let data_inner = unsafe { core::slice::from_raw_parts_mut(data_ptr, data_len) };
188 unsafe { super::#avx512_fn(data_inner, sign); }
189 return;
190 }
191 }
192 };
193 }
194 return quote! {};
196 }
197
198 let avx_branch = if config.precision == Precision::F64 {
200 let avx_f64_fn = format_ident!("codelet_simd_{}_avx_f64", size);
201 quote! {
202 if cached_level == ISA_AVX_LEVEL {
203 let data_len = data.len() * 2;
205 let data_ptr = data.as_mut_ptr().cast::<#ty_tokens>();
206 let data_inner = unsafe { core::slice::from_raw_parts_mut(data_ptr, data_len) };
207 unsafe { super::#avx_f64_fn(data_inner, sign); }
208 return;
209 }
210 }
211 } else {
212 quote! {}
213 };
214
215 quote! {
216 #[cfg(target_arch = "x86_64")]
217 {
218 #[cfg(feature = "avx512")]
219 if cached_level == ISA_AVX512_LEVEL {
220 let data_len = data.len() * 2;
222 let data_ptr = data.as_mut_ptr().cast::<#ty_tokens>();
223 let data_inner = unsafe { core::slice::from_raw_parts_mut(data_ptr, data_len) };
224 unsafe { super::#avx512_fn(data_inner, sign); }
225 return;
226 }
227 if cached_level == ISA_AVX2_FMA_LEVEL {
228 let data_len = data.len() * 2;
230 let data_ptr = data.as_mut_ptr().cast::<#ty_tokens>();
231 let data_inner = unsafe { core::slice::from_raw_parts_mut(data_ptr, data_len) };
232 unsafe { super::#avx2_fn(data_inner, sign); }
233 return;
234 }
235 #avx_branch
236 if cached_level == ISA_SSE2_LEVEL {
237 let data_len = data.len() * 2;
239 let data_ptr = data.as_mut_ptr().cast::<#ty_tokens>();
240 let data_inner = unsafe { core::slice::from_raw_parts_mut(data_ptr, data_len) };
241 unsafe { super::#sse2_fn(data_inner, sign); }
242 return;
243 }
244 }
245 }
246}
247
248fn build_aarch64_branch(config: DispatcherConfig) -> TokenStream {
252 if config.size == 16 {
253 return quote! {};
254 }
255 let ty_str = config.precision.type_str();
256 let ty_tokens: TokenStream = ty_str
257 .parse()
258 .unwrap_or_else(|_| unreachable!("ty_str is always f32 or f64"));
259 let neon_fn = format_ident!("codelet_simd_{}_neon_{}", config.size, ty_str);
260 quote! {
261 #[cfg(target_arch = "aarch64")]
262 {
263 if cached_level == ISA_NEON_LEVEL {
264 let data_len = data.len() * 2;
266 let data_ptr = data.as_mut_ptr().cast::<#ty_tokens>();
267 let data_inner = unsafe { core::slice::from_raw_parts_mut(data_ptr, data_len) };
268 unsafe { super::#neon_fn(data_inner, sign); }
269 return;
270 }
271 }
272 }
273}
274
275#[allow(clippy::too_many_lines)] pub fn generate_dispatcher(config: DispatcherConfig) -> Result<TokenStream, syn::Error> {
297 let size = config.size;
298 if !matches!(size, 2 | 4 | 8 | 16) {
299 return Err(syn::Error::new(
300 proc_macro2::Span::call_site(),
301 format!(
302 "gen_dispatcher_codelet: unsupported size {size} (expected one of 2, 4, 8, 16)"
303 ),
304 ));
305 }
306
307 let ty_str = config.precision.type_str();
308 let ty_upper = ty_str.to_uppercase();
309 let size_str = size.to_string();
310
311 let static_name = format_ident!("DETECTED_ISA_{}_{}", size_str, ty_upper);
313 let detect_fn = format_ident!("detect_isa_{}_{}", size_str, ty_str);
315 let cached_fn = format_ident!("codelet_simd_{}_cached_{}", size_str, ty_str);
317 let scalar_fn = format_ident!("codelet_simd_{}_scalar", size);
319
320 let detect_x86_body = build_detect_x86_body();
321 let detect_aarch64_body = build_detect_aarch64_body();
322 let x86_64_branches = build_x86_64_branches(config);
323 let aarch64_branch = build_aarch64_branch(config);
324
325 let ty_tokens: TokenStream = ty_str
326 .parse()
327 .unwrap_or_else(|_| unreachable!("ty_str is always f32 or f64"));
328
329 let fn_doc = format!(
330 "Cached runtime ISA dispatcher for size-{size} DFT ({ty_str}).\n\n\
331 On first call, probes CPU features and stores the ISA level in a\n\
332 thread-safe `AtomicU8` static. Subsequent calls read the cache with\n\
333 `Relaxed` ordering (benign-racy: all threads converge on the same answer).\n\n\
334 Dispatch priority on `x86_64`: AVX-512F > AVX2+FMA > AVX > SSE2 > scalar.\n\
335 Dispatch priority on `aarch64`: NEON > scalar.\n\
336 Other architectures fall through to the scalar codelet."
337 );
338
339 let size_lit = size;
340
341 Ok(quote! {
342 const ISA_SCALAR_LEVEL: u8 = 0;
344 const ISA_SSE2_LEVEL: u8 = 1;
345 const ISA_AVX_LEVEL: u8 = 2;
346 const ISA_AVX2_FMA_LEVEL: u8 = 3;
347 const ISA_AVX512_LEVEL: u8 = 4;
348 const ISA_NEON_LEVEL: u8 = 5;
349 const ISA_UNDETECTED_LEVEL: u8 = 255;
350
351 static #static_name: core::sync::atomic::AtomicU8 =
355 core::sync::atomic::AtomicU8::new(ISA_UNDETECTED_LEVEL);
356
357 fn #detect_fn() -> u8 {
359 #detect_x86_body
360 #detect_aarch64_body
361 #[allow(unreachable_code)]
362 ISA_SCALAR_LEVEL
363 }
364
365 #[doc = #fn_doc]
366 #[inline]
367 pub fn #cached_fn(
368 data: &mut [crate::kernel::Complex<#ty_tokens>],
369 sign: i32,
370 ) {
371 debug_assert!(
372 data.len() >= #size_lit,
373 "codelet_simd_{}_cached_{}: need >= {} elements, got {}",
374 #size_lit,
375 stringify!(#ty_tokens),
376 #size_lit,
377 data.len(),
378 );
379
380 let cached_level = {
382 let level = #static_name.load(core::sync::atomic::Ordering::Relaxed);
383 if level == ISA_UNDETECTED_LEVEL {
384 let detected = #detect_fn();
385 #static_name.store(detected, core::sync::atomic::Ordering::Relaxed);
387 detected
388 } else {
389 level
390 }
391 };
392
393 #x86_64_branches
399 #aarch64_branch
400
401 super::#scalar_fn(data, sign);
404 }
405 })
406}
407
408struct MacroArgs {
414 size: usize,
415 precision: Precision,
416}
417
418impl Parse for MacroArgs {
419 fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
420 let mut size: Option<usize> = None;
421 let mut precision: Option<Precision> = None;
422
423 while !input.is_empty() {
424 let key: syn::Ident = input.parse()?;
425 let _eq: Token![=] = input.parse()?;
426 match key.to_string().as_str() {
427 "size" => {
428 let lit: LitInt = input.parse()?;
429 size = Some(lit.base10_parse::<usize>().map_err(|_| {
430 syn::Error::new(lit.span(), "expected an integer literal for `size`")
431 })?);
432 }
433 "ty" => {
434 let ident: syn::Ident = input.parse()?;
435 precision = Some(match ident.to_string().as_str() {
436 "f32" => Precision::F32,
437 "f64" => Precision::F64,
438 other => {
439 return Err(syn::Error::new(
440 ident.span(),
441 format!("unknown ty `{other}`, expected f32 or f64"),
442 ));
443 }
444 });
445 }
446 other => {
447 return Err(syn::Error::new(
448 key.span(),
449 format!("unknown key `{other}`, expected one of: size, ty"),
450 ));
451 }
452 }
453 if input.peek(Token![,]) {
454 let _: Token![,] = input.parse()?;
455 }
456 }
457
458 let size = size.ok_or_else(|| {
459 syn::Error::new(proc_macro2::Span::call_site(), "missing `size` argument")
460 })?;
461 let precision = precision.ok_or_else(|| {
462 syn::Error::new(proc_macro2::Span::call_site(), "missing `ty` argument")
463 })?;
464
465 Ok(Self { size, precision })
466 }
467}
468
469pub fn generate_from_macro(input: TokenStream) -> Result<TokenStream, syn::Error> {
478 let args: MacroArgs = syn::parse2(input)?;
479 generate_dispatcher(DispatcherConfig {
480 size: args.size,
481 precision: args.precision,
482 })
483}
484
485#[cfg(test)]
490mod tests {
491 use super::*;
492
493 #[test]
496 fn test_dispatcher_config_valid_f32() {
497 let config = DispatcherConfig {
498 size: 4,
499 precision: Precision::F32,
500 };
501 assert_eq!(config.size, 4);
502 assert_eq!(config.precision, Precision::F32);
503 }
504
505 #[test]
506 fn test_dispatcher_config_valid_f64() {
507 let config = DispatcherConfig {
508 size: 8,
509 precision: Precision::F64,
510 };
511 assert_eq!(config.size, 8);
512 assert_eq!(config.precision, Precision::F64);
513 }
514
515 #[test]
518 fn test_isa_constants_are_ordered() {
519 const _: () = {
521 assert!(ISA_SCALAR < ISA_SSE2);
522 assert!(ISA_SSE2 < ISA_AVX);
523 assert!(ISA_AVX < ISA_AVX2_FMA);
524 assert!(ISA_AVX2_FMA < ISA_AVX512);
525 assert!(ISA_NEON != ISA_SCALAR);
526 assert!(ISA_UNDETECTED == 255);
527 };
528 }
529
530 #[test]
533 fn test_generate_dispatcher_nonempty() {
534 let ts = generate_dispatcher(DispatcherConfig {
535 size: 4,
536 precision: Precision::F32,
537 })
538 .expect("should generate for size 4 f32");
539 assert!(!ts.is_empty(), "TokenStream must not be empty");
540 }
541
542 #[test]
543 fn test_generate_dispatcher_nonempty_f64() {
544 let ts = generate_dispatcher(DispatcherConfig {
545 size: 8,
546 precision: Precision::F64,
547 })
548 .expect("should generate for size 8 f64");
549 assert!(!ts.is_empty(), "TokenStream must not be empty");
550 }
551
552 #[test]
553 fn test_generate_dispatcher_contains_is_x86_feature_detected() {
554 let ts = generate_dispatcher(DispatcherConfig {
555 size: 4,
556 precision: Precision::F32,
557 })
558 .expect("should generate");
559 let s = ts.to_string();
560 assert!(
561 s.contains("is_x86_feature_detected"),
562 "generated code must contain is_x86_feature_detected! macro; got snippet: {}",
563 &s[..s.len().min(500)]
564 );
565 }
566
567 #[test]
568 fn test_generate_dispatcher_contains_atomic_u8() {
569 let ts = generate_dispatcher(DispatcherConfig {
570 size: 4,
571 precision: Precision::F32,
572 })
573 .expect("should generate");
574 let s = ts.to_string();
575 assert!(
576 s.contains("AtomicU8"),
577 "generated code must contain AtomicU8 static; got snippet: {}",
578 &s[..s.len().min(500)]
579 );
580 }
581
582 #[test]
583 fn test_generate_dispatcher_contains_isa_undetected() {
584 let ts = generate_dispatcher(DispatcherConfig {
585 size: 4,
586 precision: Precision::F32,
587 })
588 .expect("should generate");
589 let s = ts.to_string();
590 assert!(
591 s.contains("ISA_UNDETECTED_LEVEL") || s.contains("255"),
592 "generated code must reference ISA_UNDETECTED_LEVEL sentinel"
593 );
594 }
595
596 #[test]
597 fn test_generate_dispatcher_function_name_size4_f32() {
598 let ts = generate_dispatcher(DispatcherConfig {
599 size: 4,
600 precision: Precision::F32,
601 })
602 .expect("should generate");
603 let s = ts.to_string();
604 assert!(
605 s.contains("codelet_simd_4_cached_f32"),
606 "expected cached dispatcher name in output; snippet: {}",
607 &s[..s.len().min(400)]
608 );
609 }
610
611 #[test]
612 fn test_generate_dispatcher_function_name_size8_f64() {
613 let ts = generate_dispatcher(DispatcherConfig {
614 size: 8,
615 precision: Precision::F64,
616 })
617 .expect("should generate");
618 let s = ts.to_string();
619 assert!(
620 s.contains("codelet_simd_8_cached_f64"),
621 "expected cached dispatcher name in output"
622 );
623 }
624
625 #[test]
626 fn test_generate_dispatcher_all_valid_sizes() {
627 for &size in &[2_usize, 4, 8, 16] {
628 for &prec in &[Precision::F32, Precision::F64] {
629 let result = generate_dispatcher(DispatcherConfig {
630 size,
631 precision: prec,
632 });
633 assert!(
634 result.is_ok(),
635 "size={size} prec={prec:?} should succeed, got: {:?}",
636 result.err()
637 );
638 }
639 }
640 }
641
642 #[test]
643 fn test_generate_dispatcher_unsupported_size_returns_error() {
644 let result = generate_dispatcher(DispatcherConfig {
645 size: 3,
646 precision: Precision::F32,
647 });
648 assert!(result.is_err(), "size 3 must return Err");
649 }
650
651 #[test]
652 fn test_generate_dispatcher_unsupported_size_6_returns_error() {
653 let result = generate_dispatcher(DispatcherConfig {
654 size: 6,
655 precision: Precision::F64,
656 });
657 assert!(result.is_err(), "size 6 must return Err");
658 }
659
660 #[test]
663 fn test_dispatcher_isa_detection() {
664 let isa = detect_host_isa();
668 assert_ne!(
669 isa, ISA_UNDETECTED,
670 "detect_host_isa must never return ISA_UNDETECTED (255)"
671 );
672 assert!(
674 matches!(
675 isa,
676 ISA_SCALAR | ISA_SSE2 | ISA_AVX | ISA_AVX2_FMA | ISA_AVX512 | ISA_NEON
677 ),
678 "detect_host_isa returned unknown level {isa}"
679 );
680 }
681
682 #[test]
683 fn test_detect_host_isa_is_deterministic() {
684 let first = detect_host_isa();
685 let second = detect_host_isa();
686 assert_eq!(first, second, "detect_host_isa must be deterministic");
687 }
688
689 #[test]
692 fn test_generate_from_macro_size4_f32() {
693 let input: TokenStream = "size = 4, ty = f32".parse().expect("valid token stream");
694 let result = generate_from_macro(input);
695 assert!(
696 result.is_ok(),
697 "size=4 ty=f32 must succeed: {:?}",
698 result.err()
699 );
700 let s = result.expect("TokenStream").to_string();
701 assert!(
702 s.contains("codelet_simd_4_cached_f32"),
703 "must contain cached dispatcher name"
704 );
705 }
706
707 #[test]
708 fn test_generate_from_macro_size8_f64() {
709 let input: TokenStream = "size = 8, ty = f64".parse().expect("valid token stream");
710 let result = generate_from_macro(input);
711 assert!(
712 result.is_ok(),
713 "size=8 ty=f64 must succeed: {:?}",
714 result.err()
715 );
716 let s = result.expect("TokenStream").to_string();
717 assert!(
718 s.contains("codelet_simd_8_cached_f64"),
719 "must contain cached dispatcher name"
720 );
721 }
722
723 #[test]
724 fn test_generate_from_macro_size2_f64() {
725 let input: TokenStream = "size = 2, ty = f64".parse().expect("valid token stream");
726 let result = generate_from_macro(input);
727 assert!(result.is_ok(), "size=2 ty=f64 must succeed");
728 }
729
730 #[test]
731 fn test_generate_from_macro_size16_f32() {
732 let input: TokenStream = "size = 16, ty = f32".parse().expect("valid token stream");
733 let result = generate_from_macro(input);
734 assert!(result.is_ok(), "size=16 ty=f32 must succeed");
735 }
736
737 #[test]
738 fn test_generate_from_macro_missing_size_returns_error() {
739 let input: TokenStream = "ty = f32".parse().expect("valid token stream");
740 let result = generate_from_macro(input);
741 assert!(result.is_err(), "missing size must return error");
742 }
743
744 #[test]
745 fn test_generate_from_macro_missing_ty_returns_error() {
746 let input: TokenStream = "size = 4".parse().expect("valid token stream");
747 let result = generate_from_macro(input);
748 assert!(result.is_err(), "missing ty must return error");
749 }
750
751 #[test]
752 fn test_generate_from_macro_unknown_ty_returns_error() {
753 let input: TokenStream = "size = 4, ty = f16".parse().expect("valid token stream");
754 let result = generate_from_macro(input);
755 assert!(result.is_err(), "unknown ty must return error");
756 }
757
758 #[test]
759 fn test_generate_from_macro_unknown_key_returns_error() {
760 let input: TokenStream = "size = 4, ty = f32, isa = avx2"
761 .parse()
762 .expect("valid token stream");
763 let result = generate_from_macro(input);
764 assert!(result.is_err(), "unknown key must return error");
765 }
766
767 #[test]
768 fn test_generate_from_macro_unsupported_size_returns_error() {
769 let input: TokenStream = "size = 5, ty = f32".parse().expect("valid token stream");
770 let result = generate_from_macro(input);
771 assert!(result.is_err(), "size=5 must return error");
772 }
773}