From 0095a4f3eacb09a1f651efe9a26bc8696dc832f3 Mon Sep 17 00:00:00 2001 From: Henrik Date: Wed, 9 Sep 2026 08:03:05 +0200 Subject: [PATCH 01/22] Factor out the policy-free half of RadixN planning The scalar planner spells out two pieces of arithmetic around its own base-choice rules, and the SIMD planners are about to need the same two. Both come out into building blocks that carry no policy: PrimeFactors::get_power_of(value) replaces the hand-written find_map over get_other_factors(). get_power_of_two and get_power_of_three already existed, this covers 5, 7 and anything later. RadixFactor::split_cross_len(len) is the 7/6/5/3-then-4s-last split of a cross-FFT length. It returns None where the scalar planner asserted, and plan.rs keeps the panic by calling .expect() on it. What is left in the planner is only its own decisions: which base, and whether Radix4 takes the length. Nothing about the algorithm set is shared, so a planner can use these without being tied to any other planner's choices. Recipes are unchanged: FftPlannerScalar dumped for every length from 1 to 20000, f32 and f64, all 40000 identical before and after. --- src/common.rs | 38 +++++++++++++++++++++++++++++++++++ src/math_utils.rs | 20 +++++++++++++++++++ src/plan.rs | 50 ++++++++--------------------------------------- 3 files changed, 66 insertions(+), 42 deletions(-) diff --git a/src/common.rs b/src/common.rs index 2dd1db12..d59dc6ea 100644 --- a/src/common.rs +++ b/src/common.rs @@ -281,4 +281,42 @@ impl RadixFactor { RadixFactor::Factor7 => 7, } } + + /// Split a RadixN cross-FFT length into the layers that make it up, or None if it has a + /// factor no layer can handle. + /// + /// Every planner picks its base by its own rules, but once the base is divided out the rest + /// is the same arithmetic for all of them, so this is deliberately free of any policy. + pub fn split_cross_len(mut cross_len: usize) -> Option> { + let mut factors = Vec::new(); + while cross_len % 7 == 0 { + cross_len /= 7; + factors.push(RadixFactor::Factor7); + } + while cross_len % 6 == 0 { + cross_len /= 6; + factors.push(RadixFactor::Factor6); + } + while cross_len % 5 == 0 { + cross_len /= 5; + factors.push(RadixFactor::Factor5); + } + while cross_len % 3 == 0 { + cross_len /= 3; + factors.push(RadixFactor::Factor3); + } + if !cross_len.is_power_of_two() { + return None; + } + + // benchmarking suggests that we want to add the 4s *last*, i suspect because 4 is a + // better-than-usual value for the transpose + let cross_bits = cross_len.trailing_zeros(); + if cross_bits % 2 == 1 { + factors.push(RadixFactor::Factor2); + } + factors.extend(std::iter::repeat(RadixFactor::Factor4).take(cross_bits as usize / 2)); + + Some(factors.into_boxed_slice()) + } } diff --git a/src/math_utils.rs b/src/math_utils.rs index 57d1cd2f..164b72f5 100644 --- a/src/math_utils.rs +++ b/src/math_utils.rs @@ -185,6 +185,26 @@ impl PrimeFactors { pub fn get_other_factors(&self) -> &[PrimeFactor] { &self.other_factors } + /// How many times `value` divides this number, or zero if it isn't a factor at all. + /// `value` is assumed to be prime. + #[allow(unused)] + pub fn get_power_of(&self, value: usize) -> u32 { + match value { + 2 => self.power_two, + 3 => self.power_three, + _ => self + .other_factors + .iter() + .find_map(|f| { + if f.value == value { + Some(f.count) + } else { + None + } + }) + .unwrap_or(0), + } + } #[allow(unused)] pub fn is_power_of_three(&self) -> bool { self.power_three > 0 && self.power_two == 0 && self.other_factors.len() == 0 diff --git a/src/plan.rs b/src/plan.rs index 38052f3d..56a790d2 100644 --- a/src/plan.rs +++ b/src/plan.rs @@ -506,18 +506,10 @@ impl FftPlannerScalar { } fn design_radixn(&mut self, factors: PrimeFactors) -> Arc { - let p2 = factors.get_power_of_two(); - let p3 = factors.get_power_of_three(); - let p5 = factors - .get_other_factors() - .iter() - .find_map(|f| if f.value == 5 { Some(f.count) } else { None }) // if we had rustc 1.62, we could use (f.value == 5).then_some(f.count) - .unwrap_or(0); - let p7 = factors - .get_other_factors() - .iter() - .find_map(|f| if f.value == 7 { Some(f.count) } else { None }) - .unwrap_or(0); + let p2 = factors.get_power_of(2); + let p3 = factors.get_power_of(3); + let p5 = factors.get_power_of(5); + let p7 = factors.get_power_of(7); let base_len: usize = if factors.has_factors_gt(MAX_RADIXN_FACTOR) { // If we have factors larger than RadixN can handle, we *must* use the product of those factors as our base @@ -563,7 +555,7 @@ impl FftPlannerScalar { // now that we know the base length, divide it out get what radix4 needs to compute let base_fft = self.design_fft_for_len(base_len); - let mut cross_len = factors.get_product() / base_len; + let cross_len = factors.get_product() / base_len; // see if we can use radix4 let cross_bits = cross_len.trailing_zeros(); @@ -574,36 +566,10 @@ impl FftPlannerScalar { // we weren't able to use radix4, so fall back to RadixN // theoretically we could do this with the p2, p3, p5 etc values above, but our choice of base knocked them out of sync - let mut factors = Vec::new(); - while cross_len % 7 == 0 { - cross_len /= 7; - factors.push(RadixFactor::Factor7); - } - while cross_len % 6 == 0 { - cross_len /= 6; - factors.push(RadixFactor::Factor6); - } - while cross_len % 5 == 0 { - cross_len /= 5; - factors.push(RadixFactor::Factor5); - } - while cross_len % 3 == 0 { - cross_len /= 3; - factors.push(RadixFactor::Factor3); - } - assert!(cross_len.is_power_of_two()); - - // benchmarking suggests that we want to add the 4s *last*, i suspect because 4 is a better-than-usual value for the transpose - let cross_bits = cross_len.trailing_zeros(); - if cross_bits % 2 == 1 { - factors.push(RadixFactor::Factor2); - } - factors.extend(std::iter::repeat(RadixFactor::Factor4).take(cross_bits as usize / 2)); + let factors = RadixFactor::split_cross_len(cross_len) + .expect("Every factor RadixN can't handle should have gone into the base"); - Arc::new(Recipe::RadixN { - factors: factors.into_boxed_slice(), - base_fft, - }) + Arc::new(Recipe::RadixN { factors, base_fft }) } // Returns Some(instance) if we have a butterfly available for this size. Returns None if there is no butterfly available for this size From daeea1eb45c41b3a7fb650a209b1c3e4c0511f17 Mon Sep 17 00:00:00 2001 From: Henrik Date: Wed, 9 Sep 2026 08:03:41 +0200 Subject: [PATCH 02/22] Make the generated prime butterflies reachable from sibling modules The RadixN drivers about to be added reuse the size-7 butterfly for their radix-7 cross-FFT layer, so the generated structs and their constructors have to be visible outside their own module. Changes the shared template and regenerates all three backends, so the autogeneration check keeps passing. The neon, sse and wasm_simd modules are themselves private, so this exports nothing new from the crate. --- src/neon/neon_prime_butterflies.rs | 64 +++++++++---------- src/sse/sse_prime_butterflies.rs | 64 +++++++++---------- src/wasm_simd/wasm_simd_prime_butterflies.rs | 64 +++++++++---------- .../src/templates/prime_template.hbs.rs | 8 +-- 4 files changed, 100 insertions(+), 100 deletions(-) diff --git a/src/neon/neon_prime_butterflies.rs b/src/neon/neon_prime_butterflies.rs index 49165478..e0da16a4 100644 --- a/src/neon/neon_prime_butterflies.rs +++ b/src/neon/neon_prime_butterflies.rs @@ -78,7 +78,7 @@ fn make_twiddles(len: usize, direction: FftDirection }) } -struct NeonF32Butterfly7 { +pub struct NeonF32Butterfly7 { direction: FftDirection, twiddles_re: [float32x4_t; 3], twiddles_im: [float32x4_t; 3], @@ -89,7 +89,7 @@ boilerplate_fft_neon_f32_butterfly!(NeonF32Butterfly7, 7, |this: &NeonF32Butterf impl NeonF32Butterfly7 { /// Safety: The current machine must support the neon instruction set #[target_feature(enable = "neon")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(7, direction); Self { @@ -182,7 +182,7 @@ impl NeonF32Butterfly7 { } } -struct NeonF64Butterfly7 { +pub struct NeonF64Butterfly7 { direction: FftDirection, twiddles_re: [float64x2_t; 3], twiddles_im: [float64x2_t; 3], @@ -193,7 +193,7 @@ boilerplate_fft_neon_f64_butterfly!(NeonF64Butterfly7, 7, |this: &NeonF64Butterf impl NeonF64Butterfly7 { /// Safety: The current machine must support the neon instruction set #[target_feature(enable = "neon")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(7, direction); unsafe {Self { @@ -257,7 +257,7 @@ impl NeonF64Butterfly7 { } } -struct NeonF32Butterfly11 { +pub struct NeonF32Butterfly11 { direction: FftDirection, twiddles_re: [float32x4_t; 5], twiddles_im: [float32x4_t; 5], @@ -268,7 +268,7 @@ boilerplate_fft_neon_f32_butterfly!(NeonF32Butterfly11, 11, |this: &NeonF32Butte impl NeonF32Butterfly11 { /// Safety: The current machine must support the neon instruction set #[target_feature(enable = "neon")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(11, direction); Self { @@ -411,7 +411,7 @@ impl NeonF32Butterfly11 { } } -struct NeonF64Butterfly11 { +pub struct NeonF64Butterfly11 { direction: FftDirection, twiddles_re: [float64x2_t; 5], twiddles_im: [float64x2_t; 5], @@ -422,7 +422,7 @@ boilerplate_fft_neon_f64_butterfly!(NeonF64Butterfly11, 11, |this: &NeonF64Butte impl NeonF64Butterfly11 { /// Safety: The current machine must support the neon instruction set #[target_feature(enable = "neon")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(11, direction); unsafe {Self { @@ -528,7 +528,7 @@ impl NeonF64Butterfly11 { } } -struct NeonF32Butterfly13 { +pub struct NeonF32Butterfly13 { direction: FftDirection, twiddles_re: [float32x4_t; 6], twiddles_im: [float32x4_t; 6], @@ -539,7 +539,7 @@ boilerplate_fft_neon_f32_butterfly!(NeonF32Butterfly13, 13, |this: &NeonF32Butte impl NeonF32Butterfly13 { /// Safety: The current machine must support the neon instruction set #[target_feature(enable = "neon")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(13, direction); Self { @@ -713,7 +713,7 @@ impl NeonF32Butterfly13 { } } -struct NeonF64Butterfly13 { +pub struct NeonF64Butterfly13 { direction: FftDirection, twiddles_re: [float64x2_t; 6], twiddles_im: [float64x2_t; 6], @@ -724,7 +724,7 @@ boilerplate_fft_neon_f64_butterfly!(NeonF64Butterfly13, 13, |this: &NeonF64Butte impl NeonF64Butterfly13 { /// Safety: The current machine must support the neon instruction set #[target_feature(enable = "neon")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(13, direction); unsafe {Self { @@ -857,7 +857,7 @@ impl NeonF64Butterfly13 { } } -struct NeonF32Butterfly17 { +pub struct NeonF32Butterfly17 { direction: FftDirection, twiddles_re: [float32x4_t; 8], twiddles_im: [float32x4_t; 8], @@ -868,7 +868,7 @@ boilerplate_fft_neon_f32_butterfly!(NeonF32Butterfly17, 17, |this: &NeonF32Butte impl NeonF32Butterfly17 { /// Safety: The current machine must support the neon instruction set #[target_feature(enable = "neon")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(17, direction); Self { @@ -1116,7 +1116,7 @@ impl NeonF32Butterfly17 { } } -struct NeonF64Butterfly17 { +pub struct NeonF64Butterfly17 { direction: FftDirection, twiddles_re: [float64x2_t; 8], twiddles_im: [float64x2_t; 8], @@ -1127,7 +1127,7 @@ boilerplate_fft_neon_f64_butterfly!(NeonF64Butterfly17, 17, |this: &NeonF64Butte impl NeonF64Butterfly17 { /// Safety: The current machine must support the neon instruction set #[target_feature(enable = "neon")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(17, direction); unsafe {Self { @@ -1326,7 +1326,7 @@ impl NeonF64Butterfly17 { } } -struct NeonF32Butterfly19 { +pub struct NeonF32Butterfly19 { direction: FftDirection, twiddles_re: [float32x4_t; 9], twiddles_im: [float32x4_t; 9], @@ -1337,7 +1337,7 @@ boilerplate_fft_neon_f32_butterfly!(NeonF32Butterfly19, 19, |this: &NeonF32Butte impl NeonF32Butterfly19 { /// Safety: The current machine must support the neon instruction set #[target_feature(enable = "neon")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(19, direction); Self { @@ -1628,7 +1628,7 @@ impl NeonF32Butterfly19 { } } -struct NeonF64Butterfly19 { +pub struct NeonF64Butterfly19 { direction: FftDirection, twiddles_re: [float64x2_t; 9], twiddles_im: [float64x2_t; 9], @@ -1639,7 +1639,7 @@ boilerplate_fft_neon_f64_butterfly!(NeonF64Butterfly19, 19, |this: &NeonF64Butte impl NeonF64Butterfly19 { /// Safety: The current machine must support the neon instruction set #[target_feature(enable = "neon")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(19, direction); unsafe {Self { @@ -1877,7 +1877,7 @@ impl NeonF64Butterfly19 { } } -struct NeonF32Butterfly23 { +pub struct NeonF32Butterfly23 { direction: FftDirection, twiddles_re: [float32x4_t; 11], twiddles_im: [float32x4_t; 11], @@ -1888,7 +1888,7 @@ boilerplate_fft_neon_f32_butterfly!(NeonF32Butterfly23, 23, |this: &NeonF32Butte impl NeonF32Butterfly23 { /// Safety: The current machine must support the neon instruction set #[target_feature(enable = "neon")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(23, direction); Self { @@ -2277,7 +2277,7 @@ impl NeonF32Butterfly23 { } } -struct NeonF64Butterfly23 { +pub struct NeonF64Butterfly23 { direction: FftDirection, twiddles_re: [float64x2_t; 11], twiddles_im: [float64x2_t; 11], @@ -2288,7 +2288,7 @@ boilerplate_fft_neon_f64_butterfly!(NeonF64Butterfly23, 23, |this: &NeonF64Butte impl NeonF64Butterfly23 { /// Safety: The current machine must support the neon instruction set #[target_feature(enable = "neon")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(23, direction); unsafe {Self { @@ -2616,7 +2616,7 @@ impl NeonF64Butterfly23 { } } -struct NeonF32Butterfly29 { +pub struct NeonF32Butterfly29 { direction: FftDirection, twiddles_re: [float32x4_t; 14], twiddles_im: [float32x4_t; 14], @@ -2627,7 +2627,7 @@ boilerplate_fft_neon_f32_butterfly!(NeonF32Butterfly29, 29, |this: &NeonF32Butte impl NeonF32Butterfly29 { /// Safety: The current machine must support the neon instruction set #[target_feature(enable = "neon")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(29, direction); Self { @@ -3193,7 +3193,7 @@ impl NeonF32Butterfly29 { } } -struct NeonF64Butterfly29 { +pub struct NeonF64Butterfly29 { direction: FftDirection, twiddles_re: [float64x2_t; 14], twiddles_im: [float64x2_t; 14], @@ -3204,7 +3204,7 @@ boilerplate_fft_neon_f64_butterfly!(NeonF64Butterfly29, 29, |this: &NeonF64Butte impl NeonF64Butterfly29 { /// Safety: The current machine must support the neon instruction set #[target_feature(enable = "neon")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(29, direction); unsafe {Self { @@ -3697,7 +3697,7 @@ impl NeonF64Butterfly29 { } } -struct NeonF32Butterfly31 { +pub struct NeonF32Butterfly31 { direction: FftDirection, twiddles_re: [float32x4_t; 15], twiddles_im: [float32x4_t; 15], @@ -3708,7 +3708,7 @@ boilerplate_fft_neon_f32_butterfly!(NeonF32Butterfly31, 31, |this: &NeonF32Butte impl NeonF32Butterfly31 { /// Safety: The current machine must support the neon instruction set #[target_feature(enable = "neon")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(31, direction); Self { @@ -4341,7 +4341,7 @@ impl NeonF32Butterfly31 { } } -struct NeonF64Butterfly31 { +pub struct NeonF64Butterfly31 { direction: FftDirection, twiddles_re: [float64x2_t; 15], twiddles_im: [float64x2_t; 15], @@ -4352,7 +4352,7 @@ boilerplate_fft_neon_f64_butterfly!(NeonF64Butterfly31, 31, |this: &NeonF64Butte impl NeonF64Butterfly31 { /// Safety: The current machine must support the neon instruction set #[target_feature(enable = "neon")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(31, direction); unsafe {Self { diff --git a/src/sse/sse_prime_butterflies.rs b/src/sse/sse_prime_butterflies.rs index 0de599a6..6d606748 100644 --- a/src/sse/sse_prime_butterflies.rs +++ b/src/sse/sse_prime_butterflies.rs @@ -78,7 +78,7 @@ fn make_twiddles(len: usize, direction: FftDirection }) } -struct SseF32Butterfly7 { +pub struct SseF32Butterfly7 { direction: FftDirection, twiddles_re: [__m128; 3], twiddles_im: [__m128; 3], @@ -89,7 +89,7 @@ boilerplate_fft_sse_f32_butterfly!(SseF32Butterfly7, 7, |this: &SseF32Butterfly7 impl SseF32Butterfly7 { /// Safety: The current machine must support the sse4.1 instruction set #[target_feature(enable = "sse4.1")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(7, direction); Self { @@ -182,7 +182,7 @@ impl SseF32Butterfly7 { } } -struct SseF64Butterfly7 { +pub struct SseF64Butterfly7 { direction: FftDirection, twiddles_re: [__m128d; 3], twiddles_im: [__m128d; 3], @@ -193,7 +193,7 @@ boilerplate_fft_sse_f64_butterfly!(SseF64Butterfly7, 7, |this: &SseF64Butterfly7 impl SseF64Butterfly7 { /// Safety: The current machine must support the sse4.1 instruction set #[target_feature(enable = "sse4.1")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(7, direction); unsafe {Self { @@ -257,7 +257,7 @@ impl SseF64Butterfly7 { } } -struct SseF32Butterfly11 { +pub struct SseF32Butterfly11 { direction: FftDirection, twiddles_re: [__m128; 5], twiddles_im: [__m128; 5], @@ -268,7 +268,7 @@ boilerplate_fft_sse_f32_butterfly!(SseF32Butterfly11, 11, |this: &SseF32Butterfl impl SseF32Butterfly11 { /// Safety: The current machine must support the sse4.1 instruction set #[target_feature(enable = "sse4.1")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(11, direction); Self { @@ -411,7 +411,7 @@ impl SseF32Butterfly11 { } } -struct SseF64Butterfly11 { +pub struct SseF64Butterfly11 { direction: FftDirection, twiddles_re: [__m128d; 5], twiddles_im: [__m128d; 5], @@ -422,7 +422,7 @@ boilerplate_fft_sse_f64_butterfly!(SseF64Butterfly11, 11, |this: &SseF64Butterfl impl SseF64Butterfly11 { /// Safety: The current machine must support the sse4.1 instruction set #[target_feature(enable = "sse4.1")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(11, direction); unsafe {Self { @@ -528,7 +528,7 @@ impl SseF64Butterfly11 { } } -struct SseF32Butterfly13 { +pub struct SseF32Butterfly13 { direction: FftDirection, twiddles_re: [__m128; 6], twiddles_im: [__m128; 6], @@ -539,7 +539,7 @@ boilerplate_fft_sse_f32_butterfly!(SseF32Butterfly13, 13, |this: &SseF32Butterfl impl SseF32Butterfly13 { /// Safety: The current machine must support the sse4.1 instruction set #[target_feature(enable = "sse4.1")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(13, direction); Self { @@ -713,7 +713,7 @@ impl SseF32Butterfly13 { } } -struct SseF64Butterfly13 { +pub struct SseF64Butterfly13 { direction: FftDirection, twiddles_re: [__m128d; 6], twiddles_im: [__m128d; 6], @@ -724,7 +724,7 @@ boilerplate_fft_sse_f64_butterfly!(SseF64Butterfly13, 13, |this: &SseF64Butterfl impl SseF64Butterfly13 { /// Safety: The current machine must support the sse4.1 instruction set #[target_feature(enable = "sse4.1")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(13, direction); unsafe {Self { @@ -857,7 +857,7 @@ impl SseF64Butterfly13 { } } -struct SseF32Butterfly17 { +pub struct SseF32Butterfly17 { direction: FftDirection, twiddles_re: [__m128; 8], twiddles_im: [__m128; 8], @@ -868,7 +868,7 @@ boilerplate_fft_sse_f32_butterfly!(SseF32Butterfly17, 17, |this: &SseF32Butterfl impl SseF32Butterfly17 { /// Safety: The current machine must support the sse4.1 instruction set #[target_feature(enable = "sse4.1")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(17, direction); Self { @@ -1116,7 +1116,7 @@ impl SseF32Butterfly17 { } } -struct SseF64Butterfly17 { +pub struct SseF64Butterfly17 { direction: FftDirection, twiddles_re: [__m128d; 8], twiddles_im: [__m128d; 8], @@ -1127,7 +1127,7 @@ boilerplate_fft_sse_f64_butterfly!(SseF64Butterfly17, 17, |this: &SseF64Butterfl impl SseF64Butterfly17 { /// Safety: The current machine must support the sse4.1 instruction set #[target_feature(enable = "sse4.1")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(17, direction); unsafe {Self { @@ -1326,7 +1326,7 @@ impl SseF64Butterfly17 { } } -struct SseF32Butterfly19 { +pub struct SseF32Butterfly19 { direction: FftDirection, twiddles_re: [__m128; 9], twiddles_im: [__m128; 9], @@ -1337,7 +1337,7 @@ boilerplate_fft_sse_f32_butterfly!(SseF32Butterfly19, 19, |this: &SseF32Butterfl impl SseF32Butterfly19 { /// Safety: The current machine must support the sse4.1 instruction set #[target_feature(enable = "sse4.1")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(19, direction); Self { @@ -1628,7 +1628,7 @@ impl SseF32Butterfly19 { } } -struct SseF64Butterfly19 { +pub struct SseF64Butterfly19 { direction: FftDirection, twiddles_re: [__m128d; 9], twiddles_im: [__m128d; 9], @@ -1639,7 +1639,7 @@ boilerplate_fft_sse_f64_butterfly!(SseF64Butterfly19, 19, |this: &SseF64Butterfl impl SseF64Butterfly19 { /// Safety: The current machine must support the sse4.1 instruction set #[target_feature(enable = "sse4.1")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(19, direction); unsafe {Self { @@ -1877,7 +1877,7 @@ impl SseF64Butterfly19 { } } -struct SseF32Butterfly23 { +pub struct SseF32Butterfly23 { direction: FftDirection, twiddles_re: [__m128; 11], twiddles_im: [__m128; 11], @@ -1888,7 +1888,7 @@ boilerplate_fft_sse_f32_butterfly!(SseF32Butterfly23, 23, |this: &SseF32Butterfl impl SseF32Butterfly23 { /// Safety: The current machine must support the sse4.1 instruction set #[target_feature(enable = "sse4.1")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(23, direction); Self { @@ -2277,7 +2277,7 @@ impl SseF32Butterfly23 { } } -struct SseF64Butterfly23 { +pub struct SseF64Butterfly23 { direction: FftDirection, twiddles_re: [__m128d; 11], twiddles_im: [__m128d; 11], @@ -2288,7 +2288,7 @@ boilerplate_fft_sse_f64_butterfly!(SseF64Butterfly23, 23, |this: &SseF64Butterfl impl SseF64Butterfly23 { /// Safety: The current machine must support the sse4.1 instruction set #[target_feature(enable = "sse4.1")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(23, direction); unsafe {Self { @@ -2616,7 +2616,7 @@ impl SseF64Butterfly23 { } } -struct SseF32Butterfly29 { +pub struct SseF32Butterfly29 { direction: FftDirection, twiddles_re: [__m128; 14], twiddles_im: [__m128; 14], @@ -2627,7 +2627,7 @@ boilerplate_fft_sse_f32_butterfly!(SseF32Butterfly29, 29, |this: &SseF32Butterfl impl SseF32Butterfly29 { /// Safety: The current machine must support the sse4.1 instruction set #[target_feature(enable = "sse4.1")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(29, direction); Self { @@ -3193,7 +3193,7 @@ impl SseF32Butterfly29 { } } -struct SseF64Butterfly29 { +pub struct SseF64Butterfly29 { direction: FftDirection, twiddles_re: [__m128d; 14], twiddles_im: [__m128d; 14], @@ -3204,7 +3204,7 @@ boilerplate_fft_sse_f64_butterfly!(SseF64Butterfly29, 29, |this: &SseF64Butterfl impl SseF64Butterfly29 { /// Safety: The current machine must support the sse4.1 instruction set #[target_feature(enable = "sse4.1")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(29, direction); unsafe {Self { @@ -3697,7 +3697,7 @@ impl SseF64Butterfly29 { } } -struct SseF32Butterfly31 { +pub struct SseF32Butterfly31 { direction: FftDirection, twiddles_re: [__m128; 15], twiddles_im: [__m128; 15], @@ -3708,7 +3708,7 @@ boilerplate_fft_sse_f32_butterfly!(SseF32Butterfly31, 31, |this: &SseF32Butterfl impl SseF32Butterfly31 { /// Safety: The current machine must support the sse4.1 instruction set #[target_feature(enable = "sse4.1")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(31, direction); Self { @@ -4341,7 +4341,7 @@ impl SseF32Butterfly31 { } } -struct SseF64Butterfly31 { +pub struct SseF64Butterfly31 { direction: FftDirection, twiddles_re: [__m128d; 15], twiddles_im: [__m128d; 15], @@ -4352,7 +4352,7 @@ boilerplate_fft_sse_f64_butterfly!(SseF64Butterfly31, 31, |this: &SseF64Butterfl impl SseF64Butterfly31 { /// Safety: The current machine must support the sse4.1 instruction set #[target_feature(enable = "sse4.1")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(31, direction); unsafe {Self { diff --git a/src/wasm_simd/wasm_simd_prime_butterflies.rs b/src/wasm_simd/wasm_simd_prime_butterflies.rs index 11a775f9..7fbd2b43 100644 --- a/src/wasm_simd/wasm_simd_prime_butterflies.rs +++ b/src/wasm_simd/wasm_simd_prime_butterflies.rs @@ -78,7 +78,7 @@ fn make_twiddles(len: usize, direction: FftDirection }) } -struct WasmSimdF32Butterfly7 { +pub struct WasmSimdF32Butterfly7 { direction: FftDirection, twiddles_re: [WasmVector32; 3], twiddles_im: [WasmVector32; 3], @@ -89,7 +89,7 @@ boilerplate_fft_wasm_simd_f32_butterfly!(WasmSimdF32Butterfly7, 7, |this: &WasmS impl WasmSimdF32Butterfly7 { /// Safety: The current machine must support the simd128 instruction set #[target_feature(enable = "simd128")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(7, direction); Self { @@ -182,7 +182,7 @@ impl WasmSimdF32Butterfly7 { } } -struct WasmSimdF64Butterfly7 { +pub struct WasmSimdF64Butterfly7 { direction: FftDirection, twiddles_re: [WasmVector64; 3], twiddles_im: [WasmVector64; 3], @@ -193,7 +193,7 @@ boilerplate_fft_wasm_simd_f64_butterfly!(WasmSimdF64Butterfly7, 7, |this: &WasmS impl WasmSimdF64Butterfly7 { /// Safety: The current machine must support the simd128 instruction set #[target_feature(enable = "simd128")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(7, direction); unsafe {Self { @@ -257,7 +257,7 @@ impl WasmSimdF64Butterfly7 { } } -struct WasmSimdF32Butterfly11 { +pub struct WasmSimdF32Butterfly11 { direction: FftDirection, twiddles_re: [WasmVector32; 5], twiddles_im: [WasmVector32; 5], @@ -268,7 +268,7 @@ boilerplate_fft_wasm_simd_f32_butterfly!(WasmSimdF32Butterfly11, 11, |this: &Was impl WasmSimdF32Butterfly11 { /// Safety: The current machine must support the simd128 instruction set #[target_feature(enable = "simd128")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(11, direction); Self { @@ -411,7 +411,7 @@ impl WasmSimdF32Butterfly11 { } } -struct WasmSimdF64Butterfly11 { +pub struct WasmSimdF64Butterfly11 { direction: FftDirection, twiddles_re: [WasmVector64; 5], twiddles_im: [WasmVector64; 5], @@ -422,7 +422,7 @@ boilerplate_fft_wasm_simd_f64_butterfly!(WasmSimdF64Butterfly11, 11, |this: &Was impl WasmSimdF64Butterfly11 { /// Safety: The current machine must support the simd128 instruction set #[target_feature(enable = "simd128")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(11, direction); unsafe {Self { @@ -528,7 +528,7 @@ impl WasmSimdF64Butterfly11 { } } -struct WasmSimdF32Butterfly13 { +pub struct WasmSimdF32Butterfly13 { direction: FftDirection, twiddles_re: [WasmVector32; 6], twiddles_im: [WasmVector32; 6], @@ -539,7 +539,7 @@ boilerplate_fft_wasm_simd_f32_butterfly!(WasmSimdF32Butterfly13, 13, |this: &Was impl WasmSimdF32Butterfly13 { /// Safety: The current machine must support the simd128 instruction set #[target_feature(enable = "simd128")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(13, direction); Self { @@ -713,7 +713,7 @@ impl WasmSimdF32Butterfly13 { } } -struct WasmSimdF64Butterfly13 { +pub struct WasmSimdF64Butterfly13 { direction: FftDirection, twiddles_re: [WasmVector64; 6], twiddles_im: [WasmVector64; 6], @@ -724,7 +724,7 @@ boilerplate_fft_wasm_simd_f64_butterfly!(WasmSimdF64Butterfly13, 13, |this: &Was impl WasmSimdF64Butterfly13 { /// Safety: The current machine must support the simd128 instruction set #[target_feature(enable = "simd128")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(13, direction); unsafe {Self { @@ -857,7 +857,7 @@ impl WasmSimdF64Butterfly13 { } } -struct WasmSimdF32Butterfly17 { +pub struct WasmSimdF32Butterfly17 { direction: FftDirection, twiddles_re: [WasmVector32; 8], twiddles_im: [WasmVector32; 8], @@ -868,7 +868,7 @@ boilerplate_fft_wasm_simd_f32_butterfly!(WasmSimdF32Butterfly17, 17, |this: &Was impl WasmSimdF32Butterfly17 { /// Safety: The current machine must support the simd128 instruction set #[target_feature(enable = "simd128")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(17, direction); Self { @@ -1116,7 +1116,7 @@ impl WasmSimdF32Butterfly17 { } } -struct WasmSimdF64Butterfly17 { +pub struct WasmSimdF64Butterfly17 { direction: FftDirection, twiddles_re: [WasmVector64; 8], twiddles_im: [WasmVector64; 8], @@ -1127,7 +1127,7 @@ boilerplate_fft_wasm_simd_f64_butterfly!(WasmSimdF64Butterfly17, 17, |this: &Was impl WasmSimdF64Butterfly17 { /// Safety: The current machine must support the simd128 instruction set #[target_feature(enable = "simd128")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(17, direction); unsafe {Self { @@ -1326,7 +1326,7 @@ impl WasmSimdF64Butterfly17 { } } -struct WasmSimdF32Butterfly19 { +pub struct WasmSimdF32Butterfly19 { direction: FftDirection, twiddles_re: [WasmVector32; 9], twiddles_im: [WasmVector32; 9], @@ -1337,7 +1337,7 @@ boilerplate_fft_wasm_simd_f32_butterfly!(WasmSimdF32Butterfly19, 19, |this: &Was impl WasmSimdF32Butterfly19 { /// Safety: The current machine must support the simd128 instruction set #[target_feature(enable = "simd128")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(19, direction); Self { @@ -1628,7 +1628,7 @@ impl WasmSimdF32Butterfly19 { } } -struct WasmSimdF64Butterfly19 { +pub struct WasmSimdF64Butterfly19 { direction: FftDirection, twiddles_re: [WasmVector64; 9], twiddles_im: [WasmVector64; 9], @@ -1639,7 +1639,7 @@ boilerplate_fft_wasm_simd_f64_butterfly!(WasmSimdF64Butterfly19, 19, |this: &Was impl WasmSimdF64Butterfly19 { /// Safety: The current machine must support the simd128 instruction set #[target_feature(enable = "simd128")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(19, direction); unsafe {Self { @@ -1877,7 +1877,7 @@ impl WasmSimdF64Butterfly19 { } } -struct WasmSimdF32Butterfly23 { +pub struct WasmSimdF32Butterfly23 { direction: FftDirection, twiddles_re: [WasmVector32; 11], twiddles_im: [WasmVector32; 11], @@ -1888,7 +1888,7 @@ boilerplate_fft_wasm_simd_f32_butterfly!(WasmSimdF32Butterfly23, 23, |this: &Was impl WasmSimdF32Butterfly23 { /// Safety: The current machine must support the simd128 instruction set #[target_feature(enable = "simd128")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(23, direction); Self { @@ -2277,7 +2277,7 @@ impl WasmSimdF32Butterfly23 { } } -struct WasmSimdF64Butterfly23 { +pub struct WasmSimdF64Butterfly23 { direction: FftDirection, twiddles_re: [WasmVector64; 11], twiddles_im: [WasmVector64; 11], @@ -2288,7 +2288,7 @@ boilerplate_fft_wasm_simd_f64_butterfly!(WasmSimdF64Butterfly23, 23, |this: &Was impl WasmSimdF64Butterfly23 { /// Safety: The current machine must support the simd128 instruction set #[target_feature(enable = "simd128")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(23, direction); unsafe {Self { @@ -2616,7 +2616,7 @@ impl WasmSimdF64Butterfly23 { } } -struct WasmSimdF32Butterfly29 { +pub struct WasmSimdF32Butterfly29 { direction: FftDirection, twiddles_re: [WasmVector32; 14], twiddles_im: [WasmVector32; 14], @@ -2627,7 +2627,7 @@ boilerplate_fft_wasm_simd_f32_butterfly!(WasmSimdF32Butterfly29, 29, |this: &Was impl WasmSimdF32Butterfly29 { /// Safety: The current machine must support the simd128 instruction set #[target_feature(enable = "simd128")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(29, direction); Self { @@ -3193,7 +3193,7 @@ impl WasmSimdF32Butterfly29 { } } -struct WasmSimdF64Butterfly29 { +pub struct WasmSimdF64Butterfly29 { direction: FftDirection, twiddles_re: [WasmVector64; 14], twiddles_im: [WasmVector64; 14], @@ -3204,7 +3204,7 @@ boilerplate_fft_wasm_simd_f64_butterfly!(WasmSimdF64Butterfly29, 29, |this: &Was impl WasmSimdF64Butterfly29 { /// Safety: The current machine must support the simd128 instruction set #[target_feature(enable = "simd128")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(29, direction); unsafe {Self { @@ -3697,7 +3697,7 @@ impl WasmSimdF64Butterfly29 { } } -struct WasmSimdF32Butterfly31 { +pub struct WasmSimdF32Butterfly31 { direction: FftDirection, twiddles_re: [WasmVector32; 15], twiddles_im: [WasmVector32; 15], @@ -3708,7 +3708,7 @@ boilerplate_fft_wasm_simd_f32_butterfly!(WasmSimdF32Butterfly31, 31, |this: &Was impl WasmSimdF32Butterfly31 { /// Safety: The current machine must support the simd128 instruction set #[target_feature(enable = "simd128")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles(31, direction); Self { @@ -4341,7 +4341,7 @@ impl WasmSimdF32Butterfly31 { } } -struct WasmSimdF64Butterfly31 { +pub struct WasmSimdF64Butterfly31 { direction: FftDirection, twiddles_re: [WasmVector64; 15], twiddles_im: [WasmVector64; 15], @@ -4352,7 +4352,7 @@ boilerplate_fft_wasm_simd_f64_butterfly!(WasmSimdF64Butterfly31, 31, |this: &Was impl WasmSimdF64Butterfly31 { /// Safety: The current machine must support the simd128 instruction set #[target_feature(enable = "simd128")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles(31, direction); unsafe {Self { diff --git a/tools/gen_simd_butterflies/src/templates/prime_template.hbs.rs b/tools/gen_simd_butterflies/src/templates/prime_template.hbs.rs index 28b8fd57..2f3fb16b 100644 --- a/tools/gen_simd_butterflies/src/templates/prime_template.hbs.rs +++ b/tools/gen_simd_butterflies/src/templates/prime_template.hbs.rs @@ -69,7 +69,7 @@ fn make_twiddles(len: usize, direction: FftDirection } {{#each lengths}} -struct {{this.struct_name_32}} { +pub struct {{this.struct_name_32}} { direction: FftDirection, twiddles_re: [{{../arch.vector_f32}}; {{this.twiddle_len}}], twiddles_im: [{{../arch.vector_f32}}; {{this.twiddle_len}}], @@ -80,7 +80,7 @@ boilerplate_fft_{{../arch.name_snakecase}}_f32_butterfly!({{this.struct_name_32} impl {{this.struct_name_32}} { /// Safety: The current machine must support the {{../arch.cpu_feature_name}} instruction set #[target_feature(enable = "{{../arch.cpu_feature_name}}")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f32::(); let twiddles = make_twiddles({{this.len}}, direction); Self { @@ -123,7 +123,7 @@ impl {{this.struct_name_32}} { } } -struct {{this.struct_name_64}} { +pub struct {{this.struct_name_64}} { direction: FftDirection, twiddles_re: [{{../arch.vector_f64}}; {{this.twiddle_len}}], twiddles_im: [{{../arch.vector_f64}}; {{this.twiddle_len}}], @@ -134,7 +134,7 @@ boilerplate_fft_{{../arch.name_snakecase}}_f64_butterfly!({{this.struct_name_64} impl {{this.struct_name_64}} { /// Safety: The current machine must support the {{../arch.cpu_feature_name}} instruction set #[target_feature(enable = "{{../arch.cpu_feature_name}}")] - unsafe fn new(direction: FftDirection) -> Self { + pub unsafe fn new(direction: FftDirection) -> Self { assert_f64::(); let twiddles = make_twiddles({{this.len}}, direction); unsafe {Self { From 1007be3f8f684eb4c7e1c10ce820a8d4d480a04c Mon Sep 17 00:00:00 2001 From: Henrik Date: Wed, 9 Sep 2026 08:09:11 +0200 Subject: [PATCH 03/22] Add a vectorised RadixN, and plan it for mixed-factor lengths on NEON Mirrors src/algorithm/radixn.rs: one flat transpose down to a base FFT, then in-place cross-FFT layers over a single packed twiddle array. The difference is that the layers use column butterflies, so a whole vector of columns goes through each butterfly call. The algorithm lives in src/simd_radixn.rs as SimdRadixN, generic over a RadixNVector trait, so the other two SIMD backends can reuse it by supplying two trait impls. NEON is the first, at 224 lines of impl plus a type alias. - generic over f32 and f64. Twiddles are stored as vectors, and the trait pairs each vector type with its radix 3, 5, 6 and 7 butterflies. Radix 2 and 4 were already vector-generic. - the base may be a composite recipe with its own scratch, which the existing boilerplate macros hardcode to zero, so Fft is implemented on SimdRadixN directly instead. - an f32 vector holds two complex numbers, so every cross-FFT layer needs an even column count. An odd base folds in a spare factor of two, and odd lengths, which have none to spare, keep the mixed radix path. src/simd_planner.rs holds the planning arithmetic, again so the other two backends get it unchanged: design_radixn, design_butterfly_product, and complex_per_vector. Each planner owns a private Recipe enum, so these hand back plain numbers and the caller builds its own recipe. The planner's dispatch chain ends up in the same order src/plan.rs uses, with the butterfly pair search ahead of RadixN. That moves the pair search out of the final else, so it now sees lengths with trailing_zeros() >= 6 that it never used to. Two of them change plan, and both get faster. NEON, forward, 10*len buffer, M1, ns/iter, mean of two runs: len dtype old new speedup 320 f32 8730 7444 1.17x 320 f64 12235 10806 1.13x 576 f32 15690 13124 1.20x 576 f64 22265 19226 1.16x Recipes audited over every length from 1 to 20000 for f32 and f64: 26323 change, 26290 of them by gaining a RadixN, and all 33 of the rest are those same two plans propagating through nested designs. Nothing else moves. Measured on an M1, f64, at 1008, 1050, 1080, 1296, 10368 and 100800: 1.46x to 2.03x faster than the previous planner, and 1.13x to 1.53x faster than the best MixedRadix tree this backend could build before. --- src/lib.rs | 8 + src/neon/mod.rs | 1 + src/neon/neon_common.rs | 59 ++++ src/neon/neon_planner.rs | 147 ++++++--- src/neon/neon_radixn.rs | 224 +++++++++++++ src/simd_planner.rs | 178 ++++++++++ src/simd_radixn.rs | 679 +++++++++++++++++++++++++++++++++++++++ 7 files changed, 1256 insertions(+), 40 deletions(-) create mode 100644 src/neon/neon_radixn.rs create mode 100644 src/simd_planner.rs create mode 100644 src/simd_radixn.rs diff --git a/src/lib.rs b/src/lib.rs index 2bcb30dd..7d0cb8ce 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -130,6 +130,14 @@ mod math_utils; mod plan; mod twiddles; +// The SIMD backends all share one RadixN, generic over the vector type each of them provides +#[cfg(all(target_arch = "aarch64", feature = "neon"))] +mod simd_radixn; + +// ...and the planner arithmetic that goes with it +#[cfg(all(target_arch = "aarch64", feature = "neon"))] +mod simd_planner; + use num_complex::Complex; use num_traits::Zero; diff --git a/src/neon/mod.rs b/src/neon/mod.rs index 20b936cb..05e96e23 100644 --- a/src/neon/mod.rs +++ b/src/neon/mod.rs @@ -7,6 +7,7 @@ mod neon_vector; pub mod neon_butterflies; pub mod neon_prime_butterflies; pub mod neon_radix4; +pub mod neon_radixn; mod neon_utils; diff --git a/src/neon/neon_common.rs b/src/neon/neon_common.rs index df31e9fe..384b64a8 100644 --- a/src/neon/neon_common.rs +++ b/src/neon/neon_common.rs @@ -144,6 +144,65 @@ macro_rules! boilerplate_fft_neon_oop { }; } +// The `RadixNVector::fft_helper_*` methods, which are the same forwarding calls for every NEON +// vector type: they hand the chunk loop to the target-feature-enabled wrappers below. +macro_rules! neon_radixn_fft_helpers { + () => { + #[inline(always)] + unsafe fn fft_helper_immut( + input: &[E], + output: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&[E], &mut [E], &mut [E]), + ) { + super::neon_common::neon_fft_helper_immut( + input, + output, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + #[inline(always)] + unsafe fn fft_helper_outofplace( + input: &mut [E], + output: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&mut [E], &mut [E], &mut [E]), + ) { + super::neon_common::neon_fft_helper_outofplace( + input, + output, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + #[inline(always)] + unsafe fn fft_helper_inplace( + buffer: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&mut [E], &mut [E]), + ) { + super::neon_common::neon_fft_helper_inplace( + buffer, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + }; +} + // A wrapper for the FFT helper functions that make sure the entire thing happens with the benefit of the NEON target feature, // so that things like loading twiddle factor registers etc can be lifted out of the loop #[target_feature(enable = "neon")] diff --git a/src/neon/neon_planner.rs b/src/neon/neon_planner.rs index b5282f1c..56840a9b 100644 --- a/src/neon/neon_planner.rs +++ b/src/neon/neon_planner.rs @@ -4,17 +4,24 @@ use std::collections::HashMap; use std::sync::Arc; -use crate::{common::FftNum, fft_cache::FftCache, FftDirection}; +use crate::{ + common::{FftNum, RadixFactor}, + fft_cache::FftCache, + FftDirection, +}; use crate::algorithm::*; use crate::neon::neon_butterflies::*; use crate::neon::neon_prime_butterflies; use crate::neon::neon_radix4::*; +use crate::neon::neon_radixn::*; use crate::Fft; use crate::math_utils::{PrimeFactor, PrimeFactors}; +use crate::simd_planner::{self, RadixNPlan}; const MIN_RADIX4_BITS: u32 = 6; // smallest size to consider radix 4 an option is 2^6 = 64 + const MAX_RADER_PRIME_FACTOR: usize = 23; // don't use Raders if the inner fft length has prime factor larger than this /// A Recipe is a structure that describes the design of a FFT, without actually creating it. @@ -50,6 +57,10 @@ pub enum Recipe { k: u32, base_fft: Arc, }, + RadixN { + factors: Box<[RadixFactor]>, + base_fft: Arc, + }, Butterfly1, Butterfly2, Butterfly3, @@ -74,6 +85,9 @@ impl Recipe { match self { Recipe::Dft(length) => *length, Recipe::Radix4 { k, base_fft } => base_fft.len() * (1 << (k * 2)), + Recipe::RadixN { factors, base_fft } => { + base_fft.len() * factors.iter().map(|f| f.radix()).product::() + } Recipe::Butterfly1 => 1, Recipe::Butterfly2 => 2, Recipe::Butterfly3 => 3, @@ -268,6 +282,16 @@ impl FftPlannerNeon { panic!("Not f32 or f64"); } } + Recipe::RadixN { factors, base_fft } => { + let base_fft = self.build_fft(&base_fft, direction); + if id_t == id_f32 { + Arc::new(NeonRadixN::::new(factors, base_fft)) as Arc> + } else if id_t == id_f64 { + Arc::new(NeonRadixN::::new(factors, base_fft)) as Arc> + } else { + panic!("Not f32 or f64"); + } + } Recipe::Butterfly1 => { if id_t == id_f32 { Arc::new(NeonF32Butterfly1::new(direction)) as Arc> @@ -445,49 +469,56 @@ impl FftPlannerNeon { fft_instance } else if factors.is_prime() { self.design_prime(len) + } else if len.trailing_zeros() >= MIN_RADIX4_BITS + && factors.get_other_factors().is_empty() + && factors.get_power_of_three() < 2 + { + // pure powers of two, and 3 * 2^k, are Radix4's job. It's a specialised RadixN, and + // measurably faster than the generic driver on the shapes it covers. + self.design_radix4(factors) + } else if let Some(butterfly_product) = self.design_butterfly_product(len) { + butterfly_product + } else if let Some(radixn) = self.design_radixn(&factors) { + radixn } else if len.trailing_zeros() >= MIN_RADIX4_BITS { - if factors.get_other_factors().is_empty() && factors.get_power_of_three() < 2 { - self.design_radix4(factors) - } else { - let non_power_of_two = factors - .remove_factors(PrimeFactor { - value: 2, - count: len.trailing_zeros(), - }) - .unwrap(); - let power_of_two = PrimeFactors::compute(1 << len.trailing_zeros()); - self.design_mixed_radix(power_of_two, non_power_of_two) - } + // RadixN couldn't take this one, so fall back to peeling the power of two off the + // front and mixed-radixing the rest. + let non_power_of_two = factors + .remove_factors(PrimeFactor { + value: 2, + count: len.trailing_zeros(), + }) + .unwrap(); + let power_of_two = PrimeFactors::compute(1 << len.trailing_zeros()); + self.design_mixed_radix(power_of_two, non_power_of_two) } else { - // Can we do this as a mixed radix with just two butterflies? - // Loop through and find all combinations - // If more than one is found, keep the one where the factors are closer together. - // For example length 20 where 10x2 and 5x4 are possible, we use 5x4. - let mut bf_left = 0; - let mut bf_right = 0; - // If the length is below 14, or over 1024 we don't need to try this. - if len > 13 && len <= 1024 { - for (n, bf_l) in self.all_butterflies.iter().enumerate() { - if len % bf_l == 0 { - let bf_r = len / bf_l; - if self.all_butterflies.iter().skip(n).any(|&m| m == bf_r) { - bf_left = *bf_l; - bf_right = bf_r; - } - } - } - if bf_left > 0 { - let fact_l = PrimeFactors::compute(bf_left); - let fact_r = PrimeFactors::compute(bf_right); - return self.design_mixed_radix(fact_l, fact_r); - } - } - // Not possible with just butterflies, go with the general solution. let (left_factors, right_factors) = factors.partition_factors(); self.design_mixed_radix(left_factors, right_factors) } } + // Can we do this as a mixed radix with just two butterflies? + fn design_butterfly_product(&mut self, len: usize) -> Option> { + let (bf_left, bf_right) = + simd_planner::design_butterfly_product(len, &self.all_butterflies)?; + + let fact_l = PrimeFactors::compute(bf_left); + let fact_r = PrimeFactors::compute(bf_right); + Some(self.design_mixed_radix(fact_l, fact_r)) + } + + // Design a RadixN, or the Radix4 that some of its shapes are better served by. Returns None + // when RadixN can't cover this length, and the caller falls back to mixed radix. + fn design_radixn(&mut self, factors: &PrimeFactors) -> Option> { + let plan = simd_planner::design_radixn(factors, simd_planner::complex_per_vector::())?; + + let base_fft = self.design_fft_for_len(plan.base_len()); + Some(match plan { + RadixNPlan::Radix4 { k, .. } => Arc::new(Recipe::Radix4 { k, base_fft }), + RadixNPlan::RadixN { factors, .. } => Arc::new(Recipe::RadixN { factors, base_fft }), + }) + } + fn design_mixed_radix( &mut self, left_factors: PrimeFactors, @@ -643,6 +674,13 @@ mod unit_tests { } } + fn is_radixn(plan: &Recipe) -> bool { + match plan { + &Recipe::RadixN { .. } => true, + _ => false, + } + } + fn is_mixedradixsmall(plan: &Recipe) -> bool { match plan { &Recipe::MixedRadixSmall { .. } => true, @@ -721,7 +759,19 @@ mod unit_tests { #[test] fn test_plan_neon_mixedradix() { - // Products of several different primes should become MixedRadix + // Products of several primes that are all too big for a RadixN cross-FFT layer should + // become MixedRadix + let mut planner = FftPlannerNeon::::new().unwrap(); + for len in [11 * 11 * 13, 11 * 13 * 17, 17 * 19 * 23, 11 * 13 * 17 * 19] { + let plan = planner.design_fft_for_len(len); + assert!(is_mixedradix(&plan), "Expected MixedRadix, got {:?}", plan); + assert_eq!(plan.len(), len, "Recipe reports wrong length"); + } + } + + #[test] + fn test_plan_neon_radixn() { + // Products of several small primes should become RadixN let mut planner = FftPlannerNeon::::new().unwrap(); for pow2 in 2..5 { for pow3 in 2..5 { @@ -732,7 +782,7 @@ mod unit_tests { * 5usize.pow(pow5) * 7usize.pow(pow7); let plan = planner.design_fft_for_len(len); - assert!(is_mixedradix(&plan), "Expected MixedRadix, got {:?}", plan); + assert!(is_radixn(&plan), "Expected RadixN, got {:?}", plan); assert_eq!(plan.len(), len, "Recipe reports wrong length"); } } @@ -740,11 +790,28 @@ mod unit_tests { } } + #[test] + fn test_plan_neon_radixn_f32_needs_an_even_base() { + // An f32 vector holds two complex numbers, so RadixN needs an even column count and can + // never take an odd length. Those have to keep falling back to mixed radix. + let mut planner32 = FftPlannerNeon::::new().unwrap(); + let mut planner64 = FftPlannerNeon::::new().unwrap(); + for len in [1215, 10125, 3125] { + let plan32 = planner32.design_fft_for_len(len); + assert!(!is_radixn(&plan32), "Expected no RadixN, got {:?}", plan32); + assert_eq!(plan32.len(), len, "Recipe reports wrong length"); + + let plan64 = planner64.design_fft_for_len(len); + assert!(is_radixn(&plan64), "Expected RadixN, got {:?}", plan64); + assert_eq!(plan64.len(), len, "Recipe reports wrong length"); + } + } + #[test] fn test_plan_neon_mixedradixsmall() { // Products of two "small" lengths < 31 that have a common divisor >1, and isn't a power of 2 should be MixedRadixSmall let mut planner = FftPlannerNeon::::new().unwrap(); - for len in [5 * 20, 5 * 25].iter() { + for len in [5 * 20, 6 * 9, 12 * 15, 10 * 15].iter() { let plan = planner.design_fft_for_len(*len); assert!( is_mixedradixsmall(&plan), diff --git a/src/neon/neon_radixn.rs b/src/neon/neon_radixn.rs new file mode 100644 index 00000000..d714bdae --- /dev/null +++ b/src/neon/neon_radixn.rs @@ -0,0 +1,224 @@ +//! The NEON side of `SimdRadixN`. +//! +//! The algorithm itself lives in `src/simd_radixn.rs`, shared by every SIMD backend. All that is +//! left here is the `RadixNVector` impl for each vector type: the NEON loads, stores and vector +//! math, and the element-type-specific butterfly structs for radix 3, 5, 6 and 7. + +use std::arch::aarch64::{float32x4_t, float64x2_t}; + +use num_complex::Complex; + +use crate::simd_radixn::{RadixNVector, SimdRadixN}; +use crate::FftDirection; + +use super::neon_butterflies::{ + NeonF32Butterfly3, NeonF32Butterfly5, NeonF32Butterfly6, NeonF64Butterfly3, NeonF64Butterfly5, + NeonF64Butterfly6, +}; +use super::neon_prime_butterflies::{NeonF32Butterfly7, NeonF64Butterfly7}; +use super::neon_vector::{NeonArray, NeonArrayMut, NeonVector, Rotation90}; +use super::NeonNum; + +/// FFT algorithm for lengths that factor into small radixes, NEON accelerated version. +/// This is designed to be used via a Planner, and not created directly. +pub type NeonRadixN = SimdRadixN<::VectorType, T>; + +impl RadixNVector for float64x2_t { + const COMPLEX_PER_VECTOR: usize = 1; + + type ScalarType = f64; + type Rotation = Rotation90; + + type Butterfly3 = NeonF64Butterfly3; + type Butterfly5 = NeonF64Butterfly5; + type Butterfly6 = NeonF64Butterfly6; + type Butterfly7 = NeonF64Butterfly7; + + #[inline(always)] + unsafe fn load(data: &[Complex], index: usize) -> Self { + data.load_complex(index) + } + #[inline(always)] + unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { + data.store_complex(value, index) + } + + #[inline(always)] + unsafe fn mul_complex(left: Self, right: Self) -> Self { + NeonVector::mul_complex(left, right) + } + #[inline(always)] + unsafe fn make_mixedradix_twiddle_chunk( + x: usize, + y: usize, + len: usize, + direction: FftDirection, + ) -> Self { + NeonVector::make_mixedradix_twiddle_chunk(x, y, len, direction) + } + + #[inline(always)] + unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { + NeonVector::make_rotate90(direction) + } + #[inline(always)] + unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { + NeonF64Butterfly3::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { + NeonF64Butterfly5::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { + NeonF64Butterfly6::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { + NeonF64Butterfly7::new(direction) + } + + #[inline(always)] + unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { + NeonVector::column_butterfly2(rows) + } + #[inline(always)] + unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { + bf.perform_fft_direct(rows[0], rows[1], rows[2]) + } + #[inline(always)] + unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { + NeonVector::column_butterfly4(rows, rotation) + } + #[inline(always)] + unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { + bf.perform_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4]) + } + #[inline(always)] + unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { + bf.perform_fft_direct(rows) + } + #[inline(always)] + unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { + bf.perform_fft_direct(rows) + } + + neon_radixn_fft_helpers!(); +} + +impl RadixNVector for float32x4_t { + const COMPLEX_PER_VECTOR: usize = 2; + + type ScalarType = f32; + type Rotation = Rotation90; + + type Butterfly3 = NeonF32Butterfly3; + type Butterfly5 = NeonF32Butterfly5; + type Butterfly6 = NeonF32Butterfly6; + type Butterfly7 = NeonF32Butterfly7; + + #[inline(always)] + unsafe fn load(data: &[Complex], index: usize) -> Self { + data.load_complex(index) + } + #[inline(always)] + unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { + data.store_complex(value, index) + } + + #[inline(always)] + unsafe fn mul_complex(left: Self, right: Self) -> Self { + NeonVector::mul_complex(left, right) + } + #[inline(always)] + unsafe fn make_mixedradix_twiddle_chunk( + x: usize, + y: usize, + len: usize, + direction: FftDirection, + ) -> Self { + NeonVector::make_mixedradix_twiddle_chunk(x, y, len, direction) + } + + #[inline(always)] + unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { + NeonVector::make_rotate90(direction) + } + #[inline(always)] + unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { + NeonF32Butterfly3::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { + NeonF32Butterfly5::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { + NeonF32Butterfly6::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { + NeonF32Butterfly7::new(direction) + } + + #[inline(always)] + unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { + NeonVector::column_butterfly2(rows) + } + #[inline(always)] + unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { + bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2]) + } + #[inline(always)] + unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { + NeonVector::column_butterfly4(rows, rotation) + } + #[inline(always)] + unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { + bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4]) + } + #[inline(always)] + unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { + bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4], rows[5]) + } + #[inline(always)] + unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { + bf.perform_parallel_fft_direct(rows) + } + + neon_radixn_fft_helpers!(); +} + +#[cfg(test)] +mod unit_tests { + use super::*; + use crate::simd_radixn::test_bodies; + + #[test] + fn test_neon_radixn_f64() { + // f64 fits one complex per vector, so every base length is legal + test_bodies::factor_pairs::(&[1, 2, 3, 4, 5, 6]); + } + + #[test] + fn test_neon_radixn_f32() { + // f32 fits two complex per vector, so the base length has to be even + test_bodies::factor_pairs::(&[2, 4, 6]); + } + + #[test] + fn test_neon_radixn_composite_base() { + test_bodies::composite_base::(); + } + + #[test] + fn test_neon_radixn_large_recipes() { + test_bodies::large_recipes::(); + } + + #[test] + #[ignore] + fn test_neon_radixn_six_layers() { + test_bodies::six_layers::(); + } +} diff --git a/src/simd_planner.rs b/src/simd_planner.rs new file mode 100644 index 00000000..f404520a --- /dev/null +++ b/src/simd_planner.rs @@ -0,0 +1,178 @@ +//! The FFT design decisions that every SIMD planner makes the same way. +//! +//! `FftPlannerNeon`, `FftPlannerSse` and `FftPlannerWasmSimd` each own a private `Recipe` enum +//! and a private cache, so they can't share the planner itself. What they can share is the +//! arithmetic that picks a plan, which is a pure function of the length's prime factors and of +//! how many complex numbers fit in one of the backend's vectors. These functions do that part +//! and hand back plain numbers; the caller turns them into its own recipes. +//! +//! The scalar planner in `src/plan.rs` deliberately stays out of this. Sharing the choice logic +//! would tie the SIMD backends to the scalar planner's algorithm set, and the two have never been +//! required to match: each backend has its own butterflies, its own Radix4 bases, and its own +//! answer for primes. Its `design_radixn` already differs in ways that change plans, not just +//! style, so folding it in would be a planner change to measure, not a deduplication. + +use crate::common::RadixFactor; +use crate::math_utils::PrimeFactors; +use crate::FftNum; + +use std::any::TypeId; + +const MAX_RADIXN_FACTOR: usize = 7; // The largest butterfly factor that the RadixN algorithm can handle + +/// How many complex numbers fit in one SIMD vector, which is what decides the column-count +/// constraints on RadixN and Radix4. +/// +/// Every backend here has 128 bit vectors, so this only depends on the float type. +pub fn complex_per_vector() -> usize { + if TypeId::of::() == TypeId::of::() { + 2 + } else { + 1 + } +} + +/// What `design_radixn` decided, in terms the caller turns into its own recipe. +pub enum RadixNPlan { + Radix4 { + k: u32, + base_len: usize, + }, + RadixN { + factors: Box<[RadixFactor]>, + base_len: usize, + }, +} + +impl RadixNPlan { + /// The base FFT length, which both variants carry, so the caller can build it once. + pub fn base_len(&self) -> usize { + match self { + RadixNPlan::Radix4 { base_len, .. } => *base_len, + RadixNPlan::RadixN { base_len, .. } => *base_len, + } + } +} + +/// Can we do this as a mixed radix with just two butterflies? +/// +/// Loops through and finds all combinations. If more than one is found, keeps the one where the +/// factors are closer together. For example length 20, where 10x2 and 5x4 are possible, gives 5x4. +/// +/// `all_butterflies` is the sorted list of lengths the backend has a butterfly for. +pub fn design_butterfly_product(len: usize, all_butterflies: &[usize]) -> Option<(usize, usize)> { + // If the length is below 14, or over 1024 we don't need to try this. + if len <= 13 || len > 1024 { + return None; + } + + let mut bf_left = 0; + let mut bf_right = 0; + for (n, bf_l) in all_butterflies.iter().enumerate() { + if len % bf_l == 0 { + let bf_r = len / bf_l; + if all_butterflies.iter().skip(n).any(|&m| m == bf_r) { + bf_left = *bf_l; + bf_right = bf_r; + } + } + } + if bf_left == 0 { + return None; + } + + Some((bf_left, bf_right)) +} + +/// Design a RadixN: fold any factors too big for a cross-FFT layer into the base, pick a base +/// for what's left, and turn the rest into a list of radixes. Mirrors `design_radixn` in +/// `src/plan.rs`, which the scalar planner uses for the same job. +/// +/// Returns None when RadixN can't cover this length, which happens for f32 when no legal base +/// is available. The caller falls back to mixed radix in that case. +pub fn design_radixn(factors: &PrimeFactors, complex_per_vector: usize) -> Option { + // With no factors small enough for a cross-FFT layer, the base would have to be the whole + // length and there would be nothing left for RadixN to do. + if !factors.has_factors_leq(MAX_RADIXN_FACTOR) { + return None; + } + + let len = factors.get_product(); + let p2 = factors.get_power_of(2); + let p3 = factors.get_power_of(3); + let p5 = factors.get_power_of(5); + let p7 = factors.get_power_of(7); + + let mut base_len: usize = if factors.has_factors_gt(MAX_RADIXN_FACTOR) { + // Factors larger than a cross-FFT layer can handle *must* go in the base + factors.product_above(MAX_RADIXN_FACTOR) + } else if p7 == 0 && p5 == 0 && p3 < 2 { + // pure powers of two, and 3 * 2^k. Use the same bases design_radix4 does, so that the + // Radix4 escape below hands these over unchanged. + if p3 == 0 { + if p2 % 2 == 1 { + 32 + } else { + 16 + } + } else if p2 % 2 == 1 { + 24 + } else { + 12 + } + } else if p2 > 0 && p3 > 0 { + // a mixed bag of 2s and 3s + match p2.saturating_sub(p3) { + 0 => 6, + 1 => 12, + _ => 24, + } + } else if p3 > 2 { + 27 + } else if p3 > 1 { + 9 + } else if p7 > 0 { + 7 + } else { + assert!(p5 > 0); + 5 + }; + + // An f32 vector holds two complex numbers, so every cross-FFT layer needs an even column + // count. The column count starts at base_len, so an odd base is unusable: fold one of the + // length's factors of two into it instead. If there isn't one to spare, RadixN can't do + // this length at all. + if base_len % complex_per_vector != 0 { + if (len / base_len) % 2 != 0 { + return None; + } + base_len *= 2; + } + + if base_len >= len || len % base_len != 0 { + return None; + } + + let cross_len = len / base_len; + + // Radix4 is faster than the generic driver on pure powers of four, so hand those over. It + // needs twice the column count RadixN does, hence the extra check on the base. + let cross_bits = cross_len.trailing_zeros(); + if cross_len.is_power_of_two() + && cross_bits % 2 == 0 + && base_len % (2 * complex_per_vector) == 0 + { + return Some(RadixNPlan::Radix4 { + k: cross_bits / 2, + base_len, + }); + } + + // Split what's left into cross-FFT layers. Every factor too big for a layer went into the + // base above, so the split can't fail, and the same expect guards it in `src/plan.rs`. + Some(RadixNPlan::RadixN { + factors: RadixFactor::split_cross_len(cross_len) + .expect("Every factor RadixN can't handle should have gone into the base"), + base_len, + }) +} diff --git a/src/simd_radixn.rs b/src/simd_radixn.rs new file mode 100644 index 00000000..c800142d --- /dev/null +++ b/src/simd_radixn.rs @@ -0,0 +1,679 @@ +//! The body of the SIMD `RadixN` implementations, shared by every SIMD backend. +//! +//! This mirrors `src/algorithm/radixn.rs`: one flat transpose down to a base FFT, then a stack of +//! in-place cross-FFT layers over a single packed twiddle array. The only difference is that the +//! cross-FFT layers use SIMD column butterflies instead of the scalar ones, so a whole vector of +//! columns is processed per butterfly call. +//! +//! Everything here is generic over `RadixNVector`, which is the handful of vector operations the +//! algorithm needs. A backend implements that trait once per vector type and gets the algorithm, +//! so `SimdRadixN` is the only copy of it. +//! +//! Because a column butterfly consumes `COMPLEX_PER_VECTOR` columns at a time, the column count at +//! every layer has to be a whole number of vectors. The column count starts at `base_len` and only +//! ever grows by whole factors, so requiring `base_len % COMPLEX_PER_VECTOR == 0` is enough. That +//! is 1 for f64 (no restriction) and 2 for f32. + +use std::any::TypeId; +use std::sync::Arc; + +use num_complex::Complex; + +use crate::array_utils::{factor_transpose, workaround_transmute_mut, TransposeFactor}; +use crate::common::{FftNum, RadixFactor}; +use crate::{Direction, Fft, FftDirection, Length}; + +/// Everything `SimdRadixN` needs from a backend's vector type. +/// +/// Radix 2 and 4 are vector-generic in every backend, so they map straight onto the backend's +/// vector trait. Radix 3, 5, 6 and 7 only exist as element-type-specific structs holding +/// precomputed twiddles, so this trait pairs each vector type with its own set of them. For f64 a +/// vector is one complex number and the plain `perform_fft_direct` is already a single column; for +/// f32 a vector is two complex numbers and `perform_parallel_fft_direct` does two columns at once. +/// +/// Safety: every method here requires the current machine to support the backend's SIMD +/// instruction set. +pub trait RadixNVector: Copy + Send + Sync + Sized { + const COMPLEX_PER_VECTOR: usize; + + /// The scalar this vector holds. Always the same type as the `T` of the `SimdRadixN` using it. + type ScalarType: FftNum; + + /// The backend's precomputed 90 degree rotation, which the radix 4 butterfly needs. + type Rotation: Copy + Send + Sync; + + type Butterfly3: Send + Sync; + type Butterfly5: Send + Sync; + type Butterfly6: Send + Sync; + type Butterfly7: Send + Sync; + + unsafe fn load(data: &[Complex], index: usize) -> Self; + unsafe fn store(data: &mut [Complex], value: Self, index: usize); + + /// Pairwise multiply the complex numbers in `left` with the complex numbers in `right`. + unsafe fn mul_complex(left: Self, right: Self) -> Self; + + /// Generates a chunk of twiddle factors starting at (X,Y) and incrementing X + /// `COMPLEX_PER_VECTOR` times. + unsafe fn make_mixedradix_twiddle_chunk( + x: usize, + y: usize, + len: usize, + direction: FftDirection, + ) -> Self; + + unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation; + unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3; + unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5; + unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6; + unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7; + + /// Each of these interprets the input as rows of a `COMPLEX_PER_VECTOR`-by-N 2D array, and + /// computes parallel butterflies down the columns of the 2D array. + unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2]; + unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3]; + unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4]; + unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5]; + unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6]; + unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7]; + + /// The three `fft_helper_*` wrappers from the backend's `*_common.rs`, which run the whole + /// chunk loop with the backend's target feature enabled so that things like loading twiddle + /// factor registers can be lifted out of the loop. + unsafe fn fft_helper_immut( + input: &[E], + output: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&[E], &mut [E], &mut [E]), + ); + unsafe fn fft_helper_outofplace( + input: &mut [E], + output: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&mut [E], &mut [E], &mut [E]), + ); + unsafe fn fft_helper_inplace( + buffer: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&mut [E], &mut [E]), + ); +} + +/// The per-layer cross-FFT kernels, holding whatever precomputed state each radix needs. +enum InternalRadixFactor { + Factor2, + Factor3(V::Butterfly3), + Factor4(V::Rotation), + Factor5(V::Butterfly5), + Factor6(V::Butterfly6), + Factor7(V::Butterfly7), +} + +impl InternalRadixFactor { + fn radix(&self) -> usize { + match self { + InternalRadixFactor::Factor2 => 2, + InternalRadixFactor::Factor3(_) => 3, + InternalRadixFactor::Factor4(_) => 4, + InternalRadixFactor::Factor5(_) => 5, + InternalRadixFactor::Factor6(_) => 6, + InternalRadixFactor::Factor7(_) => 7, + } + } +} + +/// FFT algorithm for lengths that factor into small radixes, SIMD accelerated version. +/// This is designed to be used via a Planner, and not created directly. +pub struct SimdRadixN { + twiddles: Box<[V]>, + + base_fft: Arc>, + base_len: usize, + + factors: Box<[TransposeFactor]>, + butterflies: Box<[InternalRadixFactor]>, + + len: usize, + direction: FftDirection, + + inplace_scratch_len: usize, + outofplace_scratch_len: usize, + immut_scratch_len: usize, +} + +impl SimdRadixN { + /// Constructs a SimdRadixN which computes FFTs of length `factor_product * base_fft.len()`. + pub fn new(factors: &[RadixFactor], base_fft: Arc>) -> Self { + // Internal sanity check: Make sure that the vector's scalar type is T. + // This struct has two generic parameters V and T, but T must always be V's scalar type, + // and they are only kept separate to help work around the lack of specialization. + assert_eq!(TypeId::of::(), TypeId::of::()); + + let base_len = base_fft.len(); + let direction = base_fft.fft_direction(); + let complex_per_vector = V::COMPLEX_PER_VECTOR; + + // Every cross-FFT layer processes a whole vector of columns at a time. The column count + // starts at base_len and is only ever multiplied by a factor, so this one check covers + // every layer. + assert!( + factors.is_empty() || base_len % complex_per_vector == 0, + "SimdRadixN requires a base length divisible by {}, got {}", + complex_per_vector, + base_len + ); + + // set up our cross FFT butterfly instances. simultaneously, compute the number of twiddles + let mut butterflies = Vec::with_capacity(factors.len()); + let mut cross_fft_len = base_len; + let mut twiddle_count = 0; + + for factor in factors { + // twiddles are stored a vector at a time, so a layer needs one chunk per vector column + twiddle_count += (cross_fft_len / complex_per_vector) * (factor.radix() - 1); + + butterflies.push(unsafe { + match factor { + RadixFactor::Factor2 => InternalRadixFactor::Factor2, + RadixFactor::Factor3 => { + InternalRadixFactor::Factor3(V::make_butterfly3(direction)) + } + RadixFactor::Factor4 => { + InternalRadixFactor::Factor4(V::make_rotate90(direction)) + } + RadixFactor::Factor5 => { + InternalRadixFactor::Factor5(V::make_butterfly5(direction)) + } + RadixFactor::Factor6 => { + InternalRadixFactor::Factor6(V::make_butterfly6(direction)) + } + RadixFactor::Factor7 => { + InternalRadixFactor::Factor7(V::make_butterfly7(direction)) + } + } + }); + + cross_fft_len *= factor.radix(); + } + let len = cross_fft_len; + + // set up our list of transpose factors - it's the same list but reversed, and we want to + // collapse duplicates. Note that we are only de-duplicating adjacent factors: if we're + // passed 7 * 2 * 7, we can't collapse the sevens because the exact order matters. + let mut transpose_factors: Vec = Vec::with_capacity(factors.len()); + for f in factors.iter().rev() { + let mut push_new = true; + if let Some(last) = transpose_factors.last_mut() { + if last.factor == *f { + last.count += 1; + push_new = false; + } + } + if push_new { + transpose_factors.push(TransposeFactor { + factor: *f, + count: 1, + }); + } + } + + // Same packing as the scalar RadixN: all layers in one array, bottom layer first, and + // within a layer, (radix - 1) twiddles per column. The difference is that a "column" here + // is a whole vector of columns, so each entry is a twiddle chunk rather than one twiddle. + let mut twiddle_factors: Vec = Vec::with_capacity(twiddle_count); + let mut cross_fft_len = base_len; + for factor in factors { + let num_vector_columns = cross_fft_len / complex_per_vector; + cross_fft_len *= factor.radix(); + + for i in 0..num_vector_columns { + for k in 1..factor.radix() { + unsafe { + twiddle_factors.push(V::make_mixedradix_twiddle_chunk( + i * complex_per_vector, + k, + cross_fft_len, + direction, + )); + } + } + } + } + + // figure out how much scratch space we need to request from callers + let base_inplace_scratch = base_fft.get_inplace_scratch_len(); + let inplace_scratch_len = if base_inplace_scratch > len { + len + base_inplace_scratch + } else { + len + }; + let outofplace_scratch_len = if base_inplace_scratch > len { + base_inplace_scratch + } else { + 0 + }; + + Self { + twiddles: twiddle_factors.into_boxed_slice(), + + base_fft, + base_len, + + factors: transpose_factors.into_boxed_slice(), + butterflies: butterflies.into_boxed_slice(), + + len, + direction, + + inplace_scratch_len, + outofplace_scratch_len, + immut_scratch_len: base_inplace_scratch, + } + } + + /// The flat transpose that reorders the input down to base-sized chunks. + #[inline(always)] + fn transpose(&self, input: &[Complex], output: &mut [Complex]) { + if let Some(unroll_factor) = self.factors.first() { + // for performance, we really, really want to unroll the transpose, but we need to make + // sure the output length is divisible by the unroll amount. choosing the first factor + // seems to reliably perform well + match unroll_factor.factor { + RadixFactor::Factor2 => { + factor_transpose::, 2>(self.base_len, input, output, &self.factors) + } + RadixFactor::Factor3 => { + factor_transpose::, 3>(self.base_len, input, output, &self.factors) + } + RadixFactor::Factor4 => { + factor_transpose::, 4>(self.base_len, input, output, &self.factors) + } + RadixFactor::Factor5 => { + factor_transpose::, 5>(self.base_len, input, output, &self.factors) + } + RadixFactor::Factor6 => { + factor_transpose::, 6>(self.base_len, input, output, &self.factors) + } + RadixFactor::Factor7 => { + factor_transpose::, 7>(self.base_len, input, output, &self.factors) + } + } + } else { + // no factors, so just pass data straight to our base + output.copy_from_slice(input); + } + } + + /// The stack of in-place cross-FFT layers, run after the base FFTs. + unsafe fn cross_ffts(&self, output: &mut [Complex]) { + let out: &mut [Complex] = workaround_transmute_mut(output); + + let mut cross_fft_len = self.base_len; + let mut layer_twiddles: &[V] = &self.twiddles; + + for factor in self.butterflies.iter() { + let num_columns = cross_fft_len; + cross_fft_len *= factor.radix(); + + for data in out.chunks_exact_mut(cross_fft_len) { + match factor { + InternalRadixFactor::Factor2 => { + cross_layer::(data, layer_twiddles, num_columns, |v| { + V::column_butterfly2(v) + }) + } + InternalRadixFactor::Factor3(bf) => { + cross_layer::(data, layer_twiddles, num_columns, |v| { + V::column_butterfly3(bf, v) + }) + } + InternalRadixFactor::Factor4(rotation) => { + cross_layer::(data, layer_twiddles, num_columns, |v| { + V::column_butterfly4(v, *rotation) + }) + } + InternalRadixFactor::Factor5(bf) => { + cross_layer::(data, layer_twiddles, num_columns, |v| { + V::column_butterfly5(bf, v) + }) + } + InternalRadixFactor::Factor6(bf) => { + cross_layer::(data, layer_twiddles, num_columns, |v| { + V::column_butterfly6(bf, v) + }) + } + InternalRadixFactor::Factor7(bf) => { + cross_layer::(data, layer_twiddles, num_columns, |v| { + V::column_butterfly7(bf, v) + }) + } + } + } + + // skip past all the twiddle factors used in this layer + let twiddle_offset = (num_columns / V::COMPLEX_PER_VECTOR) * (factor.radix() - 1); + layer_twiddles = &layer_twiddles[twiddle_offset..]; + } + } + + unsafe fn perform_fft_immut( + &self, + input: &[Complex], + output: &mut [Complex], + scratch: &mut [Complex], + ) { + self.transpose(input, output); + self.base_fft.process_with_scratch(output, scratch); + self.cross_ffts(output); + } + + unsafe fn perform_fft_out_of_place( + &self, + input: &mut [Complex], + output: &mut [Complex], + scratch: &mut [Complex], + ) { + self.transpose(input, output); + + // the input is free once the transpose is done, so use it as base scratch when we weren't + // handed any of our own + let base_scratch = if !scratch.is_empty() { scratch } else { input }; + self.base_fft.process_with_scratch(output, base_scratch); + + self.cross_ffts(output); + } +} + +impl Fft for SimdRadixN { + fn process_immutable_with_scratch( + &self, + input: &[Complex], + output: &mut [Complex], + scratch: &mut [Complex], + ) { + unsafe { + V::fft_helper_immut( + input, + output, + scratch, + self.len(), + self.get_immutable_scratch_len(), + |in_chunk, out_chunk, scratch| self.perform_fft_immut(in_chunk, out_chunk, scratch), + ); + } + } + fn process_outofplace_with_scratch( + &self, + input: &mut [Complex], + output: &mut [Complex], + scratch: &mut [Complex], + ) { + unsafe { + V::fft_helper_outofplace( + input, + output, + scratch, + self.len(), + self.get_outofplace_scratch_len(), + |in_chunk, out_chunk, scratch| { + self.perform_fft_out_of_place(in_chunk, out_chunk, scratch) + }, + ); + } + } + fn process_with_scratch(&self, buffer: &mut [Complex], scratch: &mut [Complex]) { + unsafe { + V::fft_helper_inplace( + buffer, + scratch, + self.len(), + self.get_inplace_scratch_len(), + |chunk, scratch| { + let (self_scratch, inner_scratch) = scratch.split_at_mut(self.len()); + self.perform_fft_out_of_place(chunk, self_scratch, inner_scratch); + chunk.copy_from_slice(self_scratch); + }, + ) + } + } + #[inline(always)] + fn get_inplace_scratch_len(&self) -> usize { + self.inplace_scratch_len + } + #[inline(always)] + fn get_outofplace_scratch_len(&self) -> usize { + self.outofplace_scratch_len + } + #[inline(always)] + fn get_immutable_scratch_len(&self) -> usize { + self.immut_scratch_len + } +} +impl Length for SimdRadixN { + #[inline(always)] + fn len(&self) -> usize { + self.len + } +} +impl Direction for SimdRadixN { + #[inline(always)] + fn fft_direction(&self) -> FftDirection { + self.direction + } +} + +/// One cross-FFT layer: for each vector of columns, gather RADIX rows strided by `num_columns`, +/// apply the twiddles, run the column butterfly, scatter back. +/// +/// Unrolled two vectors at a time, which is what the SIMD Radix4's `butterfly_4` does and is what +/// gets the two independent dependency chains needed to keep the FMA pipeline busy. +#[inline(always)] +unsafe fn cross_layer( + data: &mut [Complex], + twiddles: &[V], + num_columns: usize, + butterfly: F, +) where + F: Fn([V; RADIX]) -> [V; RADIX], +{ + let complex_per_vector = V::COMPLEX_PER_VECTOR; + let num_vector_columns = num_columns / complex_per_vector; + let tw_stride = RADIX - 1; + + debug_assert!(twiddles.len() >= num_vector_columns * tw_stride); + + // The row-0 twiddle is always 1, so it's neither stored nor applied. + let gather = |data: &[Complex], idx: usize, tw_base: usize| -> [V; RADIX] { + // row 0 first, so the array is fully initialized without `array::from_fn`, which is + // newer than the crate MSRV + let mut rows = [V::load(data, idx); RADIX]; + for (r, row) in rows.iter_mut().enumerate().skip(1) { + let v = V::load(data, idx + r * num_columns); + *row = V::mul_complex(v, *twiddles.get_unchecked(tw_base + r - 1)); + } + rows + }; + + let mut vcol = 0; + while vcol + 2 <= num_vector_columns { + let idx = vcol * complex_per_vector; + + let a = gather(data, idx, vcol * tw_stride); + let b = gather(data, idx + complex_per_vector, (vcol + 1) * tw_stride); + + let a = butterfly(a); + let b = butterfly(b); + + for (r, (a_row, b_row)) in a.iter().zip(b.iter()).enumerate() { + V::store(data, *a_row, idx + r * num_columns); + V::store(data, *b_row, idx + complex_per_vector + r * num_columns); + } + + vcol += 2; + } + + // an odd vector column count leaves one behind + if vcol < num_vector_columns { + let idx = vcol * complex_per_vector; + let a = butterfly(gather(data, idx, vcol * tw_stride)); + for (r, a_row) in a.iter().enumerate() { + V::store(data, *a_row, idx + r * num_columns); + } + } +} + +/// The test bodies, shared the same way the algorithm is. Every backend runs all of them, from a +/// thin test function per body that names its own vector types. The bodies that exercise both +/// element types take both vector types, so `V32` is always the f32 vector and `V64` the f64 one. +#[cfg(test)] +pub mod test_bodies { + use super::*; + use crate::test_utils::{check_fft_algorithm, construct_base}; + use num_traits::Float; + use rand::distributions::uniform::SampleUniform; + + const FACTOR_LIST: &[RadixFactor] = &[ + RadixFactor::Factor2, + RadixFactor::Factor3, + RadixFactor::Factor4, + RadixFactor::Factor5, + RadixFactor::Factor6, + RadixFactor::Factor7, + ]; + + /// Every empty, one-factor and two-factor recipe over each of `bases`, both directions. + pub fn factor_pairs(bases: &[usize]) + where + V: RadixNVector, + V::ScalarType: Float + SampleUniform, + { + for base in bases { + let base_forward = construct_base(*base, FftDirection::Forward); + let base_inverse = construct_base(*base, FftDirection::Inverse); + + check::(&[], Arc::clone(&base_forward)); + check::(&[], Arc::clone(&base_inverse)); + + for factor_a in FACTOR_LIST { + check::(&[*factor_a], Arc::clone(&base_forward)); + check::(&[*factor_a], Arc::clone(&base_inverse)); + + for factor_b in FACTOR_LIST { + let factors = &[*factor_a, *factor_b]; + check::(factors, Arc::clone(&base_forward)); + check::(factors, Arc::clone(&base_inverse)); + } + } + } + } + + /// The base doesn't have to be a scratch-free butterfly. A composite base is a recursive + /// recipe that needs its own scratch, which is the case `design_radixn` hits whenever a length + /// has factors above 7 (for example 11 * 13 = 143). + pub fn composite_base() + where + V32: RadixNVector, + V64: RadixNVector, + { + let mut planner64 = crate::FftPlannerScalar::::new(); + let mut planner32 = crate::FftPlannerScalar::::new(); + + for direction in [FftDirection::Forward, FftDirection::Inverse] { + // odd base, f64 only + for base_len in [143, 55, 65] { + let base = planner64.plan_fft(base_len, direction); + assert!( + base.get_inplace_scratch_len() > 0, + "base {} was expected to need scratch", + base_len + ); + check::(&[RadixFactor::Factor6, RadixFactor::Factor4], base); + } + + // even base, usable by both element types + for base_len in [22, 26, 110] { + let base = planner32.plan_fft(base_len, direction); + assert!( + base.get_inplace_scratch_len() > 0, + "base {} was expected to need scratch", + base_len + ); + check::(&[RadixFactor::Factor3, RadixFactor::Factor4], base); + + let base = planner64.plan_fft(base_len, direction); + check::(&[RadixFactor::Factor3, RadixFactor::Factor4], base); + } + } + } + + /// The recipes the spike was benchmarked on, so the layer shapes that actually matter stay + /// covered. The benchmarked lengths used much bigger bases, but the base is just a butterfly + /// that the other tests already cover, so the smallest legal one is used here. That keeps the + /// naive `Dft` the result is checked against affordable. + pub fn large_recipes() + where + V32: RadixNVector, + V64: RadixNVector, + { + use RadixFactor::*; + // (factors, f64 base, f32 base). f32 needs an even base, so it gets its own. + let cases: [(&[RadixFactor], usize, usize); 5] = [ + (&[Factor6, Factor6, Factor6], 1, 2), // 216, 432 + (&[Factor6, Factor5, Factor5], 1, 2), // 150, 300 + (&[Factor6, Factor6, Factor4], 1, 2), // 144, 288 + (&[Factor6, Factor6, Factor3], 1, 2), // 108, 216 + (&[Factor6, Factor6, Factor6, Factor4], 1, 2), // 864, 1728 + ]; + recipes::(&cases); + } + + /// The deepest benchmarked recipe, six layers. The smallest legal base still leaves 14400 + /// (f64) and 28800 (f32) points, and the naive `Dft` they are checked against takes minutes + /// on a debug build, so this one is kept out of the normal run. Run it with + /// `cargo test --release -- --ignored radixn_six_layers`. + pub fn six_layers() + where + V32: RadixNVector, + V64: RadixNVector, + { + use RadixFactor::*; + let cases: [(&[RadixFactor], usize, usize); 1] = [( + &[Factor6, Factor6, Factor5, Factor5, Factor4, Factor4], + 1, + 2, + )]; // 14400, 28800 + recipes::(&cases); + } + + /// Runs each (factors, f64 base, f32 base) case in both directions. + fn recipes(cases: &[(&[RadixFactor], usize, usize)]) + where + V32: RadixNVector, + V64: RadixNVector, + { + for (factors, base64, base32) in cases { + for direction in [FftDirection::Forward, FftDirection::Inverse] { + check::(factors, construct_base(*base64, direction)); + check::(factors, construct_base(*base32, direction)); + } + } + } + + fn check(factors: &[RadixFactor], base_fft: Arc>) + where + V: RadixNVector, + V::ScalarType: Float + SampleUniform, + { + let len = base_fft.len() * factors.iter().map(|f| f.radix()).product::(); + let direction = base_fft.fft_direction(); + let fft: SimdRadixN = SimdRadixN::new(factors, base_fft); + + check_fft_algorithm::(&fft, len, direction); + } +} From 4c016905882519c5c1e2a5c93c84dfcf4b677458 Mon Sep 17 00:00:00 2001 From: Henrik Date: Wed, 9 Sep 2026 08:11:40 +0200 Subject: [PATCH 04/22] Port the RadixN work to SSE The shared layer in simd_radixn.rs and simd_planner.rs landed in a shape SSE can use unchanged, so this is additive apart from the two cfg gates in lib.rs, which grow an x86_64 arm. SSE is the most mechanical of the ports. SseVector already mirrors NeonVector method for method, the SseArray and SseArrayMut load and store traits match, and the butterfly structs for radix 3, 5, 6 and 7 have the same perform_fft_direct and perform_parallel_fft_direct shapes, so the two RadixNVector impls are the NEON ones with the names changed. sse_radixn.rs comes out at the same 224 lines as neon_radixn.rs. The planner gets the same treatment: a Recipe::RadixN variant, thin wrappers over simd_planner::design_radixn and design_butterfly_product, and the same dispatch chain order. Recipes audited over every length from 1 to 20000 for f32 and f64, and the change is exactly the one NEON saw: 26290 lengths gain a RadixN, and all 33 of the remaining differences are the butterfly-pair reorder at 320 and 576 propagating through nested designs. The resulting designs are identical to NEON's at all 40000, which is the check that the shared layer really is shared. Both reorder lengths measured faster on a Ryzen 7 250, forward, 10*len buffer, ns/iter, median of five runs: len dtype old new speedup 320 f32 9729 5642 1.72x 320 f64 11225 8123 1.38x 576 f32 18582 10532 1.76x 576 f64 31156 15097 2.06x That is a good deal more than the 1.13x to 1.20x the same plan change gave on an M1, so the size of the win is microarchitecture specific. Length 512, whose plan does not change, measures the same in both trees to within 0.5%, which rules out a build difference behind these numbers. A later run on the same machine put 576 f64 at 0.78x rather than 2.06x, so that one number is not settled and is being re-measured; the other three have been stable across runs. --- src/lib.rs | 10 +- src/sse/mod.rs | 1 + src/sse/sse_common.rs | 59 +++++++++++ src/sse/sse_planner.rs | 147 +++++++++++++++++++-------- src/sse/sse_radixn.rs | 224 +++++++++++++++++++++++++++++++++++++++++ 5 files changed, 399 insertions(+), 42 deletions(-) create mode 100644 src/sse/sse_radixn.rs diff --git a/src/lib.rs b/src/lib.rs index 7d0cb8ce..6a783341 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -131,11 +131,17 @@ mod plan; mod twiddles; // The SIMD backends all share one RadixN, generic over the vector type each of them provides -#[cfg(all(target_arch = "aarch64", feature = "neon"))] +#[cfg(any( + all(target_arch = "aarch64", feature = "neon"), + all(target_arch = "x86_64", feature = "sse"), +))] mod simd_radixn; // ...and the planner arithmetic that goes with it -#[cfg(all(target_arch = "aarch64", feature = "neon"))] +#[cfg(any( + all(target_arch = "aarch64", feature = "neon"), + all(target_arch = "x86_64", feature = "sse"), +))] mod simd_planner; use num_complex::Complex; diff --git a/src/sse/mod.rs b/src/sse/mod.rs index df50b59c..201375b3 100644 --- a/src/sse/mod.rs +++ b/src/sse/mod.rs @@ -7,6 +7,7 @@ mod sse_vector; pub mod sse_butterflies; pub mod sse_prime_butterflies; pub mod sse_radix4; +pub mod sse_radixn; mod sse_utils; diff --git a/src/sse/sse_common.rs b/src/sse/sse_common.rs index 78b26739..a7b5a1b7 100644 --- a/src/sse/sse_common.rs +++ b/src/sse/sse_common.rs @@ -144,6 +144,65 @@ macro_rules! boilerplate_fft_sse_oop { }; } +// The `RadixNVector::fft_helper_*` methods, which are the same forwarding calls for every SSE +// vector type: they hand the chunk loop to the target-feature-enabled wrappers below. +macro_rules! sse_radixn_fft_helpers { + () => { + #[inline(always)] + unsafe fn fft_helper_immut( + input: &[E], + output: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&[E], &mut [E], &mut [E]), + ) { + super::sse_common::sse_fft_helper_immut( + input, + output, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + #[inline(always)] + unsafe fn fft_helper_outofplace( + input: &mut [E], + output: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&mut [E], &mut [E], &mut [E]), + ) { + super::sse_common::sse_fft_helper_outofplace( + input, + output, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + #[inline(always)] + unsafe fn fft_helper_inplace( + buffer: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&mut [E], &mut [E]), + ) { + super::sse_common::sse_fft_helper_inplace( + buffer, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + }; +} + // A wrapper for the FFT helper functions that make sure the entire thing happens with the benefit of the SSE target feature, // so that things like loading twiddle factor registers etc can be lifted out of the loop #[target_feature(enable = "sse4.1")] diff --git a/src/sse/sse_planner.rs b/src/sse/sse_planner.rs index 6d2e790d..8df9dc6c 100644 --- a/src/sse/sse_planner.rs +++ b/src/sse/sse_planner.rs @@ -4,17 +4,24 @@ use std::collections::HashMap; use std::sync::Arc; -use crate::{common::FftNum, fft_cache::FftCache, FftDirection}; +use crate::{ + common::{FftNum, RadixFactor}, + fft_cache::FftCache, + FftDirection, +}; use crate::algorithm::*; use crate::sse::sse_butterflies::*; use crate::sse::sse_prime_butterflies; use crate::sse::sse_radix4::*; +use crate::sse::sse_radixn::*; use crate::Fft; use crate::math_utils::{PrimeFactor, PrimeFactors}; +use crate::simd_planner::{self, RadixNPlan}; const MIN_RADIX4_BITS: u32 = 6; // smallest size to consider radix 4 an option is 2^6 = 64 + const MAX_RADER_PRIME_FACTOR: usize = 23; // don't use Raders if the inner fft length has prime factor larger than this /// A Recipe is a structure that describes the design of a FFT, without actually creating it. @@ -50,6 +57,10 @@ pub enum Recipe { k: u32, base_fft: Arc, }, + RadixN { + factors: Box<[RadixFactor]>, + base_fft: Arc, + }, Butterfly1, Butterfly2, Butterfly3, @@ -74,6 +85,9 @@ impl Recipe { match self { Recipe::Dft(length) => *length, Recipe::Radix4 { k, base_fft } => base_fft.len() * (1 << (k * 2)), + Recipe::RadixN { factors, base_fft } => { + base_fft.len() * factors.iter().map(|f| f.radix()).product::() + } Recipe::Butterfly1 => 1, Recipe::Butterfly2 => 2, Recipe::Butterfly3 => 3, @@ -268,6 +282,16 @@ impl FftPlannerSse { panic!("Not f32 or f64"); } } + Recipe::RadixN { factors, base_fft } => { + let base_fft = self.build_fft(&base_fft, direction); + if id_t == id_f32 { + Arc::new(SseRadixN::::new(factors, base_fft)) as Arc> + } else if id_t == id_f64 { + Arc::new(SseRadixN::::new(factors, base_fft)) as Arc> + } else { + panic!("Not f32 or f64"); + } + } Recipe::Butterfly1 => { if id_t == id_f32 { Arc::new(SseF32Butterfly1::new(direction)) as Arc> @@ -446,49 +470,56 @@ impl FftPlannerSse { fft_instance } else if factors.is_prime() { self.design_prime(len) + } else if len.trailing_zeros() >= MIN_RADIX4_BITS + && factors.get_other_factors().is_empty() + && factors.get_power_of_three() < 2 + { + // pure powers of two, and 3 * 2^k, are Radix4's job. It's a specialised RadixN, and + // measurably faster than the generic driver on the shapes it covers. + self.design_radix4(factors) + } else if let Some(butterfly_product) = self.design_butterfly_product(len) { + butterfly_product + } else if let Some(radixn) = self.design_radixn(&factors) { + radixn } else if len.trailing_zeros() >= MIN_RADIX4_BITS { - if factors.get_other_factors().is_empty() && factors.get_power_of_three() < 2 { - self.design_radix4(factors) - } else { - let non_power_of_two = factors - .remove_factors(PrimeFactor { - value: 2, - count: len.trailing_zeros(), - }) - .unwrap(); - let power_of_two = PrimeFactors::compute(1 << len.trailing_zeros()); - self.design_mixed_radix(power_of_two, non_power_of_two) - } + // RadixN couldn't take this one, so fall back to peeling the power of two off the + // front and mixed-radixing the rest. + let non_power_of_two = factors + .remove_factors(PrimeFactor { + value: 2, + count: len.trailing_zeros(), + }) + .unwrap(); + let power_of_two = PrimeFactors::compute(1 << len.trailing_zeros()); + self.design_mixed_radix(power_of_two, non_power_of_two) } else { - // Can we do this as a mixed radix with just two butterflies? - // Loop through and find all combinations - // If more than one is found, keep the one where the factors are closer together. - // For example length 20 where 10x2 and 5x4 are possible, we use 5x4. - let mut bf_left = 0; - let mut bf_right = 0; - // If the length is below 14, or over 1024 we don't need to try this. - if len > 13 && len <= 1024 { - for (n, bf_l) in self.all_butterflies.iter().enumerate() { - if len % bf_l == 0 { - let bf_r = len / bf_l; - if self.all_butterflies.iter().skip(n).any(|&m| m == bf_r) { - bf_left = *bf_l; - bf_right = bf_r; - } - } - } - if bf_left > 0 { - let fact_l = PrimeFactors::compute(bf_left); - let fact_r = PrimeFactors::compute(bf_right); - return self.design_mixed_radix(fact_l, fact_r); - } - } - // Not possible with just butterflies, go with the general solution. let (left_factors, right_factors) = factors.partition_factors(); self.design_mixed_radix(left_factors, right_factors) } } + // Can we do this as a mixed radix with just two butterflies? + fn design_butterfly_product(&mut self, len: usize) -> Option> { + let (bf_left, bf_right) = + simd_planner::design_butterfly_product(len, &self.all_butterflies)?; + + let fact_l = PrimeFactors::compute(bf_left); + let fact_r = PrimeFactors::compute(bf_right); + Some(self.design_mixed_radix(fact_l, fact_r)) + } + + // Design a RadixN, or the Radix4 that some of its shapes are better served by. Returns None + // when RadixN can't cover this length, and the caller falls back to mixed radix. + fn design_radixn(&mut self, factors: &PrimeFactors) -> Option> { + let plan = simd_planner::design_radixn(factors, simd_planner::complex_per_vector::())?; + + let base_fft = self.design_fft_for_len(plan.base_len()); + Some(match plan { + RadixNPlan::Radix4 { k, .. } => Arc::new(Recipe::Radix4 { k, base_fft }), + RadixNPlan::RadixN { factors, .. } => Arc::new(Recipe::RadixN { factors, base_fft }), + }) + } + fn design_mixed_radix( &mut self, left_factors: PrimeFactors, @@ -644,6 +675,13 @@ mod unit_tests { } } + fn is_radixn(plan: &Recipe) -> bool { + match plan { + &Recipe::RadixN { .. } => true, + _ => false, + } + } + fn is_mixedradixsmall(plan: &Recipe) -> bool { match plan { &Recipe::MixedRadixSmall { .. } => true, @@ -722,7 +760,19 @@ mod unit_tests { #[test] fn test_plan_sse_mixedradix() { - // Products of several different primes should become MixedRadix + // Products of several primes that are all too big for a RadixN cross-FFT layer should + // become MixedRadix + let mut planner = FftPlannerSse::::new().unwrap(); + for len in [11 * 11 * 13, 11 * 13 * 17, 17 * 19 * 23, 11 * 13 * 17 * 19] { + let plan = planner.design_fft_for_len(len); + assert!(is_mixedradix(&plan), "Expected MixedRadix, got {:?}", plan); + assert_eq!(plan.len(), len, "Recipe reports wrong length"); + } + } + + #[test] + fn test_plan_sse_radixn() { + // Products of several small primes should become RadixN let mut planner = FftPlannerSse::::new().unwrap(); for pow2 in 2..5 { for pow3 in 2..5 { @@ -733,7 +783,7 @@ mod unit_tests { * 5usize.pow(pow5) * 7usize.pow(pow7); let plan = planner.design_fft_for_len(len); - assert!(is_mixedradix(&plan), "Expected MixedRadix, got {:?}", plan); + assert!(is_radixn(&plan), "Expected RadixN, got {:?}", plan); assert_eq!(plan.len(), len, "Recipe reports wrong length"); } } @@ -741,11 +791,28 @@ mod unit_tests { } } + #[test] + fn test_plan_sse_radixn_f32_needs_an_even_base() { + // An f32 vector holds two complex numbers, so RadixN needs an even column count and can + // never take an odd length. Those have to keep falling back to mixed radix. + let mut planner32 = FftPlannerSse::::new().unwrap(); + let mut planner64 = FftPlannerSse::::new().unwrap(); + for len in [1215, 10125, 3125] { + let plan32 = planner32.design_fft_for_len(len); + assert!(!is_radixn(&plan32), "Expected no RadixN, got {:?}", plan32); + assert_eq!(plan32.len(), len, "Recipe reports wrong length"); + + let plan64 = planner64.design_fft_for_len(len); + assert!(is_radixn(&plan64), "Expected RadixN, got {:?}", plan64); + assert_eq!(plan64.len(), len, "Recipe reports wrong length"); + } + } + #[test] fn test_plan_sse_mixedradixsmall() { // Products of two "small" lengths < 31 that have a common divisor >1, and isn't a power of 2 should be MixedRadixSmall let mut planner = FftPlannerSse::::new().unwrap(); - for len in [5 * 20, 5 * 25].iter() { + for len in [5 * 20, 6 * 9, 12 * 15, 10 * 15].iter() { let plan = planner.design_fft_for_len(*len); assert!( is_mixedradixsmall(&plan), diff --git a/src/sse/sse_radixn.rs b/src/sse/sse_radixn.rs new file mode 100644 index 00000000..45feaf2b --- /dev/null +++ b/src/sse/sse_radixn.rs @@ -0,0 +1,224 @@ +//! The SSE side of `SimdRadixN`. +//! +//! The algorithm itself lives in `src/simd_radixn.rs`, shared by every SIMD backend. All that is +//! left here is the `RadixNVector` impl for each vector type: the SSE loads, stores and vector +//! math, and the element-type-specific butterfly structs for radix 3, 5, 6 and 7. + +use std::arch::x86_64::{__m128, __m128d}; + +use num_complex::Complex; + +use crate::simd_radixn::{RadixNVector, SimdRadixN}; +use crate::FftDirection; + +use super::sse_butterflies::{ + SseF32Butterfly3, SseF32Butterfly5, SseF32Butterfly6, SseF64Butterfly3, SseF64Butterfly5, + SseF64Butterfly6, +}; +use super::sse_prime_butterflies::{SseF32Butterfly7, SseF64Butterfly7}; +use super::sse_vector::{Rotation90, SseArray, SseArrayMut, SseVector}; +use super::SseNum; + +/// FFT algorithm for lengths that factor into small radixes, SSE accelerated version. +/// This is designed to be used via a Planner, and not created directly. +pub type SseRadixN = SimdRadixN<::VectorType, T>; + +impl RadixNVector for __m128d { + const COMPLEX_PER_VECTOR: usize = 1; + + type ScalarType = f64; + type Rotation = Rotation90; + + type Butterfly3 = SseF64Butterfly3; + type Butterfly5 = SseF64Butterfly5; + type Butterfly6 = SseF64Butterfly6; + type Butterfly7 = SseF64Butterfly7; + + #[inline(always)] + unsafe fn load(data: &[Complex], index: usize) -> Self { + data.load_complex(index) + } + #[inline(always)] + unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { + data.store_complex(value, index) + } + + #[inline(always)] + unsafe fn mul_complex(left: Self, right: Self) -> Self { + SseVector::mul_complex(left, right) + } + #[inline(always)] + unsafe fn make_mixedradix_twiddle_chunk( + x: usize, + y: usize, + len: usize, + direction: FftDirection, + ) -> Self { + SseVector::make_mixedradix_twiddle_chunk(x, y, len, direction) + } + + #[inline(always)] + unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { + SseVector::make_rotate90(direction) + } + #[inline(always)] + unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { + SseF64Butterfly3::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { + SseF64Butterfly5::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { + SseF64Butterfly6::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { + SseF64Butterfly7::new(direction) + } + + #[inline(always)] + unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { + SseVector::column_butterfly2(rows) + } + #[inline(always)] + unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { + bf.perform_fft_direct(rows[0], rows[1], rows[2]) + } + #[inline(always)] + unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { + SseVector::column_butterfly4(rows, rotation) + } + #[inline(always)] + unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { + bf.perform_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4]) + } + #[inline(always)] + unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { + bf.perform_fft_direct(rows) + } + #[inline(always)] + unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { + bf.perform_fft_direct(rows) + } + + sse_radixn_fft_helpers!(); +} + +impl RadixNVector for __m128 { + const COMPLEX_PER_VECTOR: usize = 2; + + type ScalarType = f32; + type Rotation = Rotation90; + + type Butterfly3 = SseF32Butterfly3; + type Butterfly5 = SseF32Butterfly5; + type Butterfly6 = SseF32Butterfly6; + type Butterfly7 = SseF32Butterfly7; + + #[inline(always)] + unsafe fn load(data: &[Complex], index: usize) -> Self { + data.load_complex(index) + } + #[inline(always)] + unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { + data.store_complex(value, index) + } + + #[inline(always)] + unsafe fn mul_complex(left: Self, right: Self) -> Self { + SseVector::mul_complex(left, right) + } + #[inline(always)] + unsafe fn make_mixedradix_twiddle_chunk( + x: usize, + y: usize, + len: usize, + direction: FftDirection, + ) -> Self { + SseVector::make_mixedradix_twiddle_chunk(x, y, len, direction) + } + + #[inline(always)] + unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { + SseVector::make_rotate90(direction) + } + #[inline(always)] + unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { + SseF32Butterfly3::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { + SseF32Butterfly5::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { + SseF32Butterfly6::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { + SseF32Butterfly7::new(direction) + } + + #[inline(always)] + unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { + SseVector::column_butterfly2(rows) + } + #[inline(always)] + unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { + bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2]) + } + #[inline(always)] + unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { + SseVector::column_butterfly4(rows, rotation) + } + #[inline(always)] + unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { + bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4]) + } + #[inline(always)] + unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { + bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4], rows[5]) + } + #[inline(always)] + unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { + bf.perform_parallel_fft_direct(rows) + } + + sse_radixn_fft_helpers!(); +} + +#[cfg(test)] +mod unit_tests { + use super::*; + use crate::simd_radixn::test_bodies; + + #[test] + fn test_sse_radixn_f64() { + // f64 fits one complex per vector, so every base length is legal + test_bodies::factor_pairs::<__m128d>(&[1, 2, 3, 4, 5, 6]); + } + + #[test] + fn test_sse_radixn_f32() { + // f32 fits two complex per vector, so the base length has to be even + test_bodies::factor_pairs::<__m128>(&[2, 4, 6]); + } + + #[test] + fn test_sse_radixn_composite_base() { + test_bodies::composite_base::<__m128, __m128d>(); + } + + #[test] + fn test_sse_radixn_large_recipes() { + test_bodies::large_recipes::<__m128, __m128d>(); + } + + #[test] + #[ignore] + fn test_sse_radixn_six_layers() { + test_bodies::six_layers::<__m128, __m128d>(); + } +} From b5d0714c06cbad0883d050013d6d6ec40a71e3cb Mon Sep 17 00:00:00 2001 From: Henrik Date: Wed, 9 Sep 2026 08:19:42 +0200 Subject: [PATCH 05/22] Port the RadixN work to wasm_simd The third and last SIMD backend. Same algorithm and the same planner branch as the other two, which the backends being structurally identical makes largely mechanical. Two things differ. The butterflies for radix 3, 5 and 6 are written against raw v128, while WasmVector32 and WasmVector64 are newtypes over it, so the column butterflies unwrap and rewrap around them; butterfly 7 already speaks the wrapper types. And wasm_simd_planner globs its own module, so mod.rs re-exports the new one. The planner tests copied from sse also get their names fixed to test_plan_wasm_simd_* and test_wasm_simd_*, matching what the other two call the same test bodies. All five RadixN tests are #[wasm_bindgen_test], not #[test]. The wasm-bindgen harness only collects the former, so a plain #[test] here is silently dropped on the one backend where these are slowest. wasm-pack test --node lists 73 passing and 1 ignored, the ignored one being the six-layer case. Recipes audited over every length from 1 to 20000 for f32 and f64, under node via wasm32-wasip1: the same 26290 lengths gain a RadixN and the same 33 are the butterfly-pair reorder at 320 and 576, with nothing else moving. The designs are identical to NEON's and SSE's at all 40000, so all three backends now plan these lengths the same way. --- src/lib.rs | 2 + src/wasm_simd/mod.rs | 2 + src/wasm_simd/wasm_simd_common.rs | 59 +++++++ src/wasm_simd/wasm_simd_planner.rs | 160 +++++++++++++------ src/wasm_simd/wasm_simd_radixn.rs | 239 +++++++++++++++++++++++++++++ 5 files changed, 414 insertions(+), 48 deletions(-) create mode 100644 src/wasm_simd/wasm_simd_radixn.rs diff --git a/src/lib.rs b/src/lib.rs index 6a783341..97f78bd9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -134,6 +134,7 @@ mod twiddles; #[cfg(any( all(target_arch = "aarch64", feature = "neon"), all(target_arch = "x86_64", feature = "sse"), + all(target_arch = "wasm32", feature = "wasm_simd"), ))] mod simd_radixn; @@ -141,6 +142,7 @@ mod simd_radixn; #[cfg(any( all(target_arch = "aarch64", feature = "neon"), all(target_arch = "x86_64", feature = "sse"), + all(target_arch = "wasm32", feature = "wasm_simd"), ))] mod simd_planner; diff --git a/src/wasm_simd/mod.rs b/src/wasm_simd/mod.rs index af7325d5..b319afc4 100644 --- a/src/wasm_simd/mod.rs +++ b/src/wasm_simd/mod.rs @@ -7,6 +7,7 @@ mod wasm_simd_vector; pub mod wasm_simd_butterflies; pub mod wasm_simd_prime_butterflies; pub mod wasm_simd_radix4; +pub mod wasm_simd_radixn; mod wasm_simd_utils; @@ -17,6 +18,7 @@ use core::arch::wasm32::v128; pub use self::wasm_simd_butterflies::*; pub use self::wasm_simd_radix4::*; +pub use self::wasm_simd_radixn::*; use self::wasm_simd_vector::WasmVector; use self::wasm_simd_vector::WasmVector32; use self::wasm_simd_vector::WasmVector64; diff --git a/src/wasm_simd/wasm_simd_common.rs b/src/wasm_simd/wasm_simd_common.rs index af32692d..fc797339 100644 --- a/src/wasm_simd/wasm_simd_common.rs +++ b/src/wasm_simd/wasm_simd_common.rs @@ -145,6 +145,65 @@ macro_rules! boilerplate_fft_wasm_simd_oop { }; } +// The `RadixNVector::fft_helper_*` methods, which are the same forwarding calls for every WASM +// SIMD vector type: they hand the chunk loop to the target-feature-enabled wrappers below. +macro_rules! wasm_simd_radixn_fft_helpers { + () => { + #[inline(always)] + unsafe fn fft_helper_immut( + input: &[E], + output: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&[E], &mut [E], &mut [E]), + ) { + super::wasm_simd_common::wasm_simd_fft_helper_immut( + input, + output, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + #[inline(always)] + unsafe fn fft_helper_outofplace( + input: &mut [E], + output: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&mut [E], &mut [E], &mut [E]), + ) { + super::wasm_simd_common::wasm_simd_fft_helper_outofplace( + input, + output, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + #[inline(always)] + unsafe fn fft_helper_inplace( + buffer: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&mut [E], &mut [E]), + ) { + super::wasm_simd_common::wasm_simd_fft_helper_inplace( + buffer, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + }; +} + // A wrapper for the FFT helper functions that make sure the entire thing happens with the benefit of the Wasm SIMD target feature, // so that things like loading twiddle factor registers etc can be lifted out of the loop #[target_feature(enable = "simd128")] diff --git a/src/wasm_simd/wasm_simd_planner.rs b/src/wasm_simd/wasm_simd_planner.rs index 4b4fd0a6..6ccc3ae3 100644 --- a/src/wasm_simd/wasm_simd_planner.rs +++ b/src/wasm_simd/wasm_simd_planner.rs @@ -4,12 +4,15 @@ use crate::algorithm::{ BluesteinsAlgorithm, Dft, GoodThomasAlgorithm, GoodThomasAlgorithmSmall, MixedRadix, MixedRadixSmall, RadersAlgorithm, }; +use crate::common::RadixFactor; use crate::math_utils::PrimeFactor; +use crate::simd_planner::{self, RadixNPlan}; use crate::wasm_simd::*; use crate::{fft_cache::FftCache, math_utils::PrimeFactors, Fft, FftDirection, FftNum}; use std::{any::TypeId, collections::HashMap, sync::Arc}; const MIN_RADIX4_BITS: u32 = 6; // smallest size to consider radix 4 an option is 2^6 = 64 + const MAX_RADER_PRIME_FACTOR: usize = 23; // don't use Raders if the inner fft length has prime factor larger than this /// A Recipe is a structure that describes the design of a FFT, without actually creating it. @@ -45,6 +48,10 @@ pub enum Recipe { k: u32, base_fft: Arc, }, + RadixN { + factors: Box<[RadixFactor]>, + base_fft: Arc, + }, Butterfly1, Butterfly2, Butterfly3, @@ -69,6 +76,9 @@ impl Recipe { match self { Recipe::Dft(length) => *length, Recipe::Radix4 { k, base_fft } => base_fft.len() * (1 << (k * 2)), + Recipe::RadixN { factors, base_fft } => { + base_fft.len() * factors.iter().map(|f| f.radix()).product::() + } Recipe::Butterfly1 => 1, Recipe::Butterfly2 => 2, Recipe::Butterfly3 => 3, @@ -241,6 +251,16 @@ impl FftPlannerWasmSimd { panic!("Not f32 or f64"); } } + Recipe::RadixN { factors, base_fft } => { + let base_fft = self.build_fft(&base_fft, direction); + if id_t == id_f32 { + Arc::new(WasmSimdRadixN::::new(factors, base_fft)) as Arc> + } else if id_t == id_f64 { + Arc::new(WasmSimdRadixN::::new(factors, base_fft)) as Arc> + } else { + panic!("Not f32 or f64"); + } + } Recipe::Butterfly1 => { if id_t == id_f32 { Arc::new(WasmSimdF32Butterfly1::new(direction)) as Arc> @@ -418,48 +438,56 @@ impl FftPlannerWasmSimd { fft_instance } else if factors.is_prime() { self.design_prime(len) + } else if len.trailing_zeros() >= MIN_RADIX4_BITS + && factors.get_other_factors().is_empty() + && factors.get_power_of_three() < 2 + { + // pure powers of two, and 3 * 2^k, are Radix4's job. It's a specialised RadixN, and + // measurably faster than the generic driver on the shapes it covers. + self.design_radix4(factors) + } else if let Some(butterfly_product) = self.design_butterfly_product(len) { + butterfly_product + } else if let Some(radixn) = self.design_radixn(&factors) { + radixn } else if len.trailing_zeros() >= MIN_RADIX4_BITS { - if factors.get_other_factors().is_empty() && factors.get_power_of_three() < 2 { - self.design_radix4(factors) - } else { - let non_power_of_two = factors - .remove_factors(PrimeFactor { - value: 2, - count: len.trailing_zeros(), - }) - .unwrap(); - let power_of_two = PrimeFactors::compute(1 << len.trailing_zeros()); - self.design_mixed_radix(power_of_two, non_power_of_two) - } + // RadixN couldn't take this one, so fall back to peeling the power of two off the + // front and mixed-radixing the rest. + let non_power_of_two = factors + .remove_factors(PrimeFactor { + value: 2, + count: len.trailing_zeros(), + }) + .unwrap(); + let power_of_two = PrimeFactors::compute(1 << len.trailing_zeros()); + self.design_mixed_radix(power_of_two, non_power_of_two) } else { - // Can we do this as a mixed radix with just two butterflies? - // Loop through and find all combinations - // If more than one is found, keep the one where the factors are closer together. - // For example length 20 where 10x2 and 5x4 are possible, we use 5x4. - let mut bf_left = 0; - let mut bf_right = 0; - // If the length is below 14, or over 1024 we don't need to try this. - if len > 13 && len <= 1024 { - for (n, bf_l) in self.all_butterflies.iter().enumerate() { - if len % bf_l == 0 { - let bf_r = len / bf_l; - if self.all_butterflies.iter().skip(n).any(|&m| m == bf_r) { - bf_left = *bf_l; - bf_right = bf_r; - } - } - } - if bf_left > 0 { - let fact_l = PrimeFactors::compute(bf_left); - let fact_r = PrimeFactors::compute(bf_right); - return self.design_mixed_radix(fact_l, fact_r); - } - } - // Not possible with just butterflies, go with the general solution. let (left_factors, right_factors) = factors.partition_factors(); self.design_mixed_radix(left_factors, right_factors) } } + + // Can we do this as a mixed radix with just two butterflies? + fn design_butterfly_product(&mut self, len: usize) -> Option> { + let (bf_left, bf_right) = + simd_planner::design_butterfly_product(len, &self.all_butterflies)?; + + let fact_l = PrimeFactors::compute(bf_left); + let fact_r = PrimeFactors::compute(bf_right); + Some(self.design_mixed_radix(fact_l, fact_r)) + } + + // Design a RadixN, or the Radix4 that some of its shapes are better served by. Returns None + // when RadixN can't cover this length, and the caller falls back to mixed radix. + fn design_radixn(&mut self, factors: &PrimeFactors) -> Option> { + let plan = simd_planner::design_radixn(factors, simd_planner::complex_per_vector::())?; + + let base_fft = self.design_fft_for_len(plan.base_len()); + Some(match plan { + RadixNPlan::Radix4 { k, .. } => Arc::new(Recipe::Radix4 { k, base_fft }), + RadixNPlan::RadixN { factors, .. } => Arc::new(Recipe::RadixN { factors, base_fft }), + }) + } + fn design_mixed_radix( &mut self, left_factors: PrimeFactors, @@ -616,6 +644,13 @@ mod unit_tests { } } + fn is_radixn(plan: &Recipe) -> bool { + match plan { + &Recipe::RadixN { .. } => true, + _ => false, + } + } + fn is_mixedradixsmall(plan: &Recipe) -> bool { match plan { &Recipe::MixedRadixSmall { .. } => true, @@ -645,7 +680,7 @@ mod unit_tests { } #[wasm_bindgen_test] - fn test_plan_sse_trivial() { + fn test_plan_wasm_simd_trivial() { // Length 0 and 1 should use Dft let mut planner = FftPlannerWasmSimd::::new().unwrap(); for len in 0..1 { @@ -656,7 +691,7 @@ mod unit_tests { } #[wasm_bindgen_test] - fn test_plan_sse_largepoweroftwo() { + fn test_plan_wasm_simd_largepoweroftwo() { // Powers of 2 above 6 should use Radix4 let mut planner = FftPlannerWasmSimd::::new().unwrap(); for pow in 6..32 { @@ -668,7 +703,7 @@ mod unit_tests { } #[wasm_bindgen_test] - fn test_plan_sse_butterflies() { + fn test_plan_wasm_simd_butterflies() { // Check that all butterflies are used let mut planner = FftPlannerWasmSimd::::new().unwrap(); assert_eq!(*planner.design_fft_for_len(2), Recipe::Butterfly2); @@ -693,8 +728,20 @@ mod unit_tests { } #[wasm_bindgen_test] - fn test_plan_sse_mixedradix() { - // Products of several different primes should become MixedRadix + fn test_plan_wasm_simd_mixedradix() { + // Products of several primes that are all too big for a RadixN cross-FFT layer should + // become MixedRadix + let mut planner = FftPlannerWasmSimd::::new().unwrap(); + for len in [11 * 11 * 13, 11 * 13 * 17, 17 * 19 * 23, 11 * 13 * 17 * 19] { + let plan = planner.design_fft_for_len(len); + assert!(is_mixedradix(&plan), "Expected MixedRadix, got {:?}", plan); + assert_eq!(plan.len(), len, "Recipe reports wrong length"); + } + } + + #[wasm_bindgen_test] + fn test_plan_wasm_simd_radixn() { + // Products of several small primes should become RadixN let mut planner = FftPlannerWasmSimd::::new().unwrap(); for pow2 in 2..5 { for pow3 in 2..5 { @@ -705,7 +752,7 @@ mod unit_tests { * 5usize.pow(pow5) * 7usize.pow(pow7); let plan = planner.design_fft_for_len(len); - assert!(is_mixedradix(&plan), "Expected MixedRadix, got {:?}", plan); + assert!(is_radixn(&plan), "Expected RadixN, got {:?}", plan); assert_eq!(plan.len(), len, "Recipe reports wrong length"); } } @@ -714,10 +761,27 @@ mod unit_tests { } #[wasm_bindgen_test] - fn test_plan_sse_mixedradixsmall() { + fn test_plan_wasm_simd_radixn_f32_needs_an_even_base() { + // An f32 vector holds two complex numbers, so RadixN needs an even column count and can + // never take an odd length. Those have to keep falling back to mixed radix. + let mut planner32 = FftPlannerWasmSimd::::new().unwrap(); + let mut planner64 = FftPlannerWasmSimd::::new().unwrap(); + for len in [1215, 10125, 3125] { + let plan32 = planner32.design_fft_for_len(len); + assert!(!is_radixn(&plan32), "Expected no RadixN, got {:?}", plan32); + assert_eq!(plan32.len(), len, "Recipe reports wrong length"); + + let plan64 = planner64.design_fft_for_len(len); + assert!(is_radixn(&plan64), "Expected RadixN, got {:?}", plan64); + assert_eq!(plan64.len(), len, "Recipe reports wrong length"); + } + } + + #[wasm_bindgen_test] + fn test_plan_wasm_simd_mixedradixsmall() { // Products of two "small" lengths < 31 that have a common divisor >1, and isn't a power of 2 should be MixedRadixSmall let mut planner = FftPlannerWasmSimd::::new().unwrap(); - for len in [5 * 20, 5 * 25].iter() { + for len in [5 * 20, 6 * 9, 12 * 15, 10 * 15].iter() { let plan = planner.design_fft_for_len(*len); assert!( is_mixedradixsmall(&plan), @@ -729,7 +793,7 @@ mod unit_tests { } #[wasm_bindgen_test] - fn test_plan_sse_goodthomasbutterfly() { + fn test_plan_wasm_simd_goodthomasbutterfly() { let mut planner = FftPlannerWasmSimd::::new().unwrap(); for len in [3 * 7, 5 * 7, 11 * 13, 2 * 29].iter() { let plan = planner.design_fft_for_len(*len); @@ -743,7 +807,7 @@ mod unit_tests { } #[wasm_bindgen_test] - fn test_plan_sse_bluestein_vs_rader() { + fn test_plan_wasm_simd_bluestein_vs_rader() { let difficultprimes: [usize; 11] = [59, 83, 107, 149, 167, 173, 179, 359, 719, 1439, 2879]; let easyprimes: [usize; 24] = [ 53, 61, 67, 71, 73, 79, 89, 97, 101, 103, 109, 113, 127, 131, 137, 139, 151, 157, 163, @@ -768,7 +832,7 @@ mod unit_tests { } #[wasm_bindgen_test] - fn test_sse_fft_cache() { + fn test_wasm_simd_fft_cache() { { // Check that FFTs are reused if they're both forward let mut planner = FftPlannerWasmSimd::::new().unwrap(); @@ -796,7 +860,7 @@ mod unit_tests { } #[wasm_bindgen_test] - fn test_sse_recipe_cache() { + fn test_wasm_simd_recipe_cache() { // Check that all butterflies are used let mut planner = FftPlannerWasmSimd::::new().unwrap(); let fft_a = planner.design_fft_for_len(1234); diff --git a/src/wasm_simd/wasm_simd_radixn.rs b/src/wasm_simd/wasm_simd_radixn.rs new file mode 100644 index 00000000..6f43cd4d --- /dev/null +++ b/src/wasm_simd/wasm_simd_radixn.rs @@ -0,0 +1,239 @@ +//! The WASM SIMD side of `SimdRadixN`. +//! +//! The algorithm itself lives in `src/simd_radixn.rs`, shared by every SIMD backend. All that is +//! left here is the `RadixNVector` impl for each vector type: the WASM SIMD loads, stores and +//! vector math, and the element-type-specific butterfly structs for radix 3, 5, 6 and 7. + +use num_complex::Complex; + +use crate::simd_radixn::{RadixNVector, SimdRadixN}; +use crate::FftDirection; + +use super::wasm_simd_butterflies::{ + WasmSimdF32Butterfly3, WasmSimdF32Butterfly5, WasmSimdF32Butterfly6, WasmSimdF64Butterfly3, + WasmSimdF64Butterfly5, WasmSimdF64Butterfly6, +}; +use super::wasm_simd_prime_butterflies::{WasmSimdF32Butterfly7, WasmSimdF64Butterfly7}; +use super::wasm_simd_vector::{ + Rotation90, WasmSimdArray, WasmSimdArrayMut, WasmVector, WasmVector32, WasmVector64, +}; +use super::WasmNum; + +/// FFT algorithm for lengths that factor into small radixes, WASM SIMD accelerated version. +/// This is designed to be used via a Planner, and not created directly. +pub type WasmSimdRadixN = SimdRadixN<::VectorType, T>; + +impl RadixNVector for WasmVector64 { + const COMPLEX_PER_VECTOR: usize = 1; + + type ScalarType = f64; + type Rotation = Rotation90; + + type Butterfly3 = WasmSimdF64Butterfly3; + type Butterfly5 = WasmSimdF64Butterfly5; + type Butterfly6 = WasmSimdF64Butterfly6; + type Butterfly7 = WasmSimdF64Butterfly7; + + #[inline(always)] + unsafe fn load(data: &[Complex], index: usize) -> Self { + data.load_complex(index) + } + #[inline(always)] + unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { + data.store_complex(value, index) + } + + #[inline(always)] + unsafe fn mul_complex(left: Self, right: Self) -> Self { + WasmVector::mul_complex(left, right) + } + #[inline(always)] + unsafe fn make_mixedradix_twiddle_chunk( + x: usize, + y: usize, + len: usize, + direction: FftDirection, + ) -> Self { + WasmVector::make_mixedradix_twiddle_chunk(x, y, len, direction) + } + + #[inline(always)] + unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { + WasmVector::make_rotate90(direction) + } + #[inline(always)] + unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { + WasmSimdF64Butterfly3::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { + WasmSimdF64Butterfly5::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { + WasmSimdF64Butterfly6::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { + WasmSimdF64Butterfly7::new(direction) + } + + #[inline(always)] + unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { + WasmVector::column_butterfly2(rows) + } + // Butterflies 3, 5 and 6 are written against raw `v128` while `WasmVector64` is a newtype over + // it, so their results come back needing rewrapping. Butterfly 7 already speaks the wrapper + // types and needs none of it. + #[inline(always)] + unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { + bf.perform_fft_direct(rows[0].0, rows[1].0, rows[2].0) + .map(f64::wrap) + } + #[inline(always)] + unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { + WasmVector::column_butterfly4(rows, rotation) + } + #[inline(always)] + unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { + bf.perform_fft_direct(rows[0].0, rows[1].0, rows[2].0, rows[3].0, rows[4].0) + .map(f64::wrap) + } + #[inline(always)] + unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { + bf.perform_fft_direct([ + rows[0].0, rows[1].0, rows[2].0, rows[3].0, rows[4].0, rows[5].0, + ]) + .map(f64::wrap) + } + #[inline(always)] + unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { + bf.perform_fft_direct(rows) + } + + wasm_simd_radixn_fft_helpers!(); +} + +impl RadixNVector for WasmVector32 { + const COMPLEX_PER_VECTOR: usize = 2; + + type ScalarType = f32; + type Rotation = Rotation90; + + type Butterfly3 = WasmSimdF32Butterfly3; + type Butterfly5 = WasmSimdF32Butterfly5; + type Butterfly6 = WasmSimdF32Butterfly6; + type Butterfly7 = WasmSimdF32Butterfly7; + + #[inline(always)] + unsafe fn load(data: &[Complex], index: usize) -> Self { + data.load_complex(index) + } + #[inline(always)] + unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { + data.store_complex(value, index) + } + + #[inline(always)] + unsafe fn mul_complex(left: Self, right: Self) -> Self { + WasmVector::mul_complex(left, right) + } + #[inline(always)] + unsafe fn make_mixedradix_twiddle_chunk( + x: usize, + y: usize, + len: usize, + direction: FftDirection, + ) -> Self { + WasmVector::make_mixedradix_twiddle_chunk(x, y, len, direction) + } + + #[inline(always)] + unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { + WasmVector::make_rotate90(direction) + } + #[inline(always)] + unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { + WasmSimdF32Butterfly3::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { + WasmSimdF32Butterfly5::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { + WasmSimdF32Butterfly6::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { + WasmSimdF32Butterfly7::new(direction) + } + + #[inline(always)] + unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { + WasmVector::column_butterfly2(rows) + } + // See the f64 impl for why butterflies 3, 5 and 6 rewrap their results and 7 doesn't. + #[inline(always)] + unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { + bf.perform_parallel_fft_direct(rows[0].0, rows[1].0, rows[2].0) + .map(f32::wrap) + } + #[inline(always)] + unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { + WasmVector::column_butterfly4(rows, rotation) + } + #[inline(always)] + unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { + bf.perform_parallel_fft_direct(rows[0].0, rows[1].0, rows[2].0, rows[3].0, rows[4].0) + .map(f32::wrap) + } + #[inline(always)] + unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { + bf.perform_parallel_fft_direct( + rows[0].0, rows[1].0, rows[2].0, rows[3].0, rows[4].0, rows[5].0, + ) + .map(f32::wrap) + } + #[inline(always)] + unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { + bf.perform_parallel_fft_direct(rows) + } + + wasm_simd_radixn_fft_helpers!(); +} + +#[cfg(test)] +mod unit_tests { + use super::*; + use crate::simd_radixn::test_bodies; + use wasm_bindgen_test::wasm_bindgen_test; + + #[wasm_bindgen_test] + fn test_wasm_simd_radixn_f64() { + // f64 fits one complex per vector, so every base length is legal + test_bodies::factor_pairs::(&[1, 2, 3, 4, 5, 6]); + } + + #[wasm_bindgen_test] + fn test_wasm_simd_radixn_f32() { + // f32 fits two complex per vector, so the base length has to be even + test_bodies::factor_pairs::(&[2, 4, 6]); + } + + #[wasm_bindgen_test] + fn test_wasm_simd_radixn_composite_base() { + test_bodies::composite_base::(); + } + + #[wasm_bindgen_test] + fn test_wasm_simd_radixn_large_recipes() { + test_bodies::large_recipes::(); + } + + #[wasm_bindgen_test] + #[ignore] + fn test_wasm_simd_radixn_six_layers() { + test_bodies::six_layers::(); + } +} From 9ce0980496aaaa1f8b9b185bba7ee50493669e8e Mon Sep 17 00:00:00 2001 From: Henrik Date: Thu, 10 Sep 2026 22:45:18 +0200 Subject: [PATCH 06/22] Match the scalar RadixN dispatch shape Move the factor dispatch outside the chunk loop, the way algorithm/radixn.rs already does it, so each layer runs one monomorphized loop over its chunks. Measured perf-neutral on NEON, this is for consistency. --- src/simd_radixn.rs | 28 ++++++++++++++++++++-------- 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/src/simd_radixn.rs b/src/simd_radixn.rs index c800142d..ba4fb932 100644 --- a/src/simd_radixn.rs +++ b/src/simd_radixn.rs @@ -321,34 +321,46 @@ impl SimdRadixN { let num_columns = cross_fft_len; cross_fft_len *= factor.radix(); - for data in out.chunks_exact_mut(cross_fft_len) { - match factor { - InternalRadixFactor::Factor2 => { + // Dispatch once per layer rather than once per chunk, so each layer runs a single + // monomorphized loop over its chunks. Mirrors the scalar `RadixN`. + match factor { + InternalRadixFactor::Factor2 => { + for data in out.chunks_exact_mut(cross_fft_len) { cross_layer::(data, layer_twiddles, num_columns, |v| { V::column_butterfly2(v) }) } - InternalRadixFactor::Factor3(bf) => { + } + InternalRadixFactor::Factor3(bf) => { + for data in out.chunks_exact_mut(cross_fft_len) { cross_layer::(data, layer_twiddles, num_columns, |v| { V::column_butterfly3(bf, v) }) } - InternalRadixFactor::Factor4(rotation) => { + } + InternalRadixFactor::Factor4(rotation) => { + for data in out.chunks_exact_mut(cross_fft_len) { cross_layer::(data, layer_twiddles, num_columns, |v| { V::column_butterfly4(v, *rotation) }) } - InternalRadixFactor::Factor5(bf) => { + } + InternalRadixFactor::Factor5(bf) => { + for data in out.chunks_exact_mut(cross_fft_len) { cross_layer::(data, layer_twiddles, num_columns, |v| { V::column_butterfly5(bf, v) }) } - InternalRadixFactor::Factor6(bf) => { + } + InternalRadixFactor::Factor6(bf) => { + for data in out.chunks_exact_mut(cross_fft_len) { cross_layer::(data, layer_twiddles, num_columns, |v| { V::column_butterfly6(bf, v) }) } - InternalRadixFactor::Factor7(bf) => { + } + InternalRadixFactor::Factor7(bf) => { + for data in out.chunks_exact_mut(cross_fft_len) { cross_layer::(data, layer_twiddles, num_columns, |v| { V::column_butterfly7(bf, v) }) From 415a29fd41cd0f66292d9d3a611743a34791a963 Mon Sep 17 00:00:00 2001 From: Henrik Date: Wed, 16 Sep 2026 14:10:59 +0200 Subject: [PATCH 07/22] Precompute the RadixN transpose indices, and drop the per-layer divide factor_transpose recomputes every column's reversed index on each call, with an out-of-line reverse_remainders call per column and two hardware divides, and chunks_exact_mut adds one more divide per cross layer. None of that scales with the length, so it is a large share of a short FFT, and more so on x86 where a 64-bit divide takes tens of cycles. Compute the reversed columns once in new() and walk the layer chunks with split_at_mut. The per-element work is unchanged. factor_transpose itself is left alone, since the scalar RadixN still uses it. --- src/simd_radixn.rs | 150 +++++++++++++++++++++++++++++---------------- 1 file changed, 97 insertions(+), 53 deletions(-) diff --git a/src/simd_radixn.rs b/src/simd_radixn.rs index ba4fb932..29616446 100644 --- a/src/simd_radixn.rs +++ b/src/simd_radixn.rs @@ -19,7 +19,7 @@ use std::sync::Arc; use num_complex::Complex; -use crate::array_utils::{factor_transpose, workaround_transmute_mut, TransposeFactor}; +use crate::array_utils::{reverse_remainders, workaround_transmute_mut, TransposeFactor}; use crate::common::{FftNum, RadixFactor}; use crate::{Direction, Fft, FftDirection, Length}; @@ -136,7 +136,10 @@ pub struct SimdRadixN { base_fft: Arc>, base_len: usize, - factors: Box<[TransposeFactor]>, + // The factor the transpose is unrolled by, and the output column of each input column. None + // when there are no factors and the transpose is a plain copy. + unroll_factor: Option, + reversed_columns: Box<[usize]>, butterflies: Box<[InternalRadixFactor]>, len: usize, @@ -223,6 +226,17 @@ impl SimdRadixN { } } + // Precompute where each column lands. Working this out per call costs two hardware divides + // and an out-of-line `reverse_remainders` call per column, which is a large share of a + // short FFT, and much more so on x86 where a 64-bit divide takes tens of cycles. + let width = len / base_len; + let reversed_columns: Box<[usize]> = (0..width) + .map(|x| reverse_remainders(x, &transpose_factors)) + .collect(); + // `table_transpose` indexes unchecked, and relies on this. + assert!(reversed_columns.iter().all(|&r| r < width)); + let unroll_factor = transpose_factors.first().map(|f| f.factor); + // Same packing as the scalar RadixN: all layers in one array, bottom layer first, and // within a layer, (radix - 1) twiddles per column. The difference is that a "column" here // is a whole vector of columns, so each entry is a twiddle chunk rather than one twiddle. @@ -265,7 +279,8 @@ impl SimdRadixN { base_fft, base_len, - factors: transpose_factors.into_boxed_slice(), + unroll_factor, + reversed_columns, butterflies: butterflies.into_boxed_slice(), len, @@ -280,29 +295,18 @@ impl SimdRadixN { /// The flat transpose that reorders the input down to base-sized chunks. #[inline(always)] fn transpose(&self, input: &[Complex], output: &mut [Complex]) { - if let Some(unroll_factor) = self.factors.first() { + if let Some(unroll_factor) = self.unroll_factor { // for performance, we really, really want to unroll the transpose, but we need to make // sure the output length is divisible by the unroll amount. choosing the first factor // seems to reliably perform well - match unroll_factor.factor { - RadixFactor::Factor2 => { - factor_transpose::, 2>(self.base_len, input, output, &self.factors) - } - RadixFactor::Factor3 => { - factor_transpose::, 3>(self.base_len, input, output, &self.factors) - } - RadixFactor::Factor4 => { - factor_transpose::, 4>(self.base_len, input, output, &self.factors) - } - RadixFactor::Factor5 => { - factor_transpose::, 5>(self.base_len, input, output, &self.factors) - } - RadixFactor::Factor6 => { - factor_transpose::, 6>(self.base_len, input, output, &self.factors) - } - RadixFactor::Factor7 => { - factor_transpose::, 7>(self.base_len, input, output, &self.factors) - } + let (height, columns) = (self.base_len, &*self.reversed_columns); + match unroll_factor { + RadixFactor::Factor2 => table_transpose::<_, 2>(height, columns, input, output), + RadixFactor::Factor3 => table_transpose::<_, 3>(height, columns, input, output), + RadixFactor::Factor4 => table_transpose::<_, 4>(height, columns, input, output), + RadixFactor::Factor5 => table_transpose::<_, 5>(height, columns, input, output), + RadixFactor::Factor6 => table_transpose::<_, 6>(height, columns, input, output), + RadixFactor::Factor7 => table_transpose::<_, 7>(height, columns, input, output), } } else { // no factors, so just pass data straight to our base @@ -325,46 +329,34 @@ impl SimdRadixN { // monomorphized loop over its chunks. Mirrors the scalar `RadixN`. match factor { InternalRadixFactor::Factor2 => { - for data in out.chunks_exact_mut(cross_fft_len) { - cross_layer::(data, layer_twiddles, num_columns, |v| { - V::column_butterfly2(v) - }) - } + cross_layer_chunks::(out, layer_twiddles, num_columns, |v| { + V::column_butterfly2(v) + }) } InternalRadixFactor::Factor3(bf) => { - for data in out.chunks_exact_mut(cross_fft_len) { - cross_layer::(data, layer_twiddles, num_columns, |v| { - V::column_butterfly3(bf, v) - }) - } + cross_layer_chunks::(out, layer_twiddles, num_columns, |v| { + V::column_butterfly3(bf, v) + }) } InternalRadixFactor::Factor4(rotation) => { - for data in out.chunks_exact_mut(cross_fft_len) { - cross_layer::(data, layer_twiddles, num_columns, |v| { - V::column_butterfly4(v, *rotation) - }) - } + cross_layer_chunks::(out, layer_twiddles, num_columns, |v| { + V::column_butterfly4(v, *rotation) + }) } InternalRadixFactor::Factor5(bf) => { - for data in out.chunks_exact_mut(cross_fft_len) { - cross_layer::(data, layer_twiddles, num_columns, |v| { - V::column_butterfly5(bf, v) - }) - } + cross_layer_chunks::(out, layer_twiddles, num_columns, |v| { + V::column_butterfly5(bf, v) + }) } InternalRadixFactor::Factor6(bf) => { - for data in out.chunks_exact_mut(cross_fft_len) { - cross_layer::(data, layer_twiddles, num_columns, |v| { - V::column_butterfly6(bf, v) - }) - } + cross_layer_chunks::(out, layer_twiddles, num_columns, |v| { + V::column_butterfly6(bf, v) + }) } InternalRadixFactor::Factor7(bf) => { - for data in out.chunks_exact_mut(cross_fft_len) { - cross_layer::(data, layer_twiddles, num_columns, |v| { - V::column_butterfly7(bf, v) - }) - } + cross_layer_chunks::(out, layer_twiddles, num_columns, |v| { + V::column_butterfly7(bf, v) + }) } } @@ -480,6 +472,58 @@ impl Direction for SimdRadixN { } } +/// Run `cross_layer` over every chunk of `data`, each `num_columns * RADIX` long. +/// +/// This is `chunks_exact_mut` without the divide it does to find the chunk count. At short lengths +/// that one divide per layer is a measurable share of the whole FFT. +#[inline(always)] +unsafe fn cross_layer_chunks( + data: &mut [Complex], + twiddles: &[V], + num_columns: usize, + butterfly: F, +) where + F: Fn([V; RADIX]) -> [V; RADIX], +{ + let chunk_len = num_columns * RADIX; + let mut rest = data; + while rest.len() >= chunk_len { + let (chunk, tail) = rest.split_at_mut(chunk_len); + cross_layer::(chunk, twiddles, num_columns, &butterfly); + rest = tail; + } + debug_assert!(rest.is_empty()); +} + +/// `factor_transpose` with the reversed column indices looked up instead of recomputed. +/// +/// `reversed_columns[x]` is the output column of input column `x`, so its length is the width and +/// nothing here needs to divide. Every entry must be below the width, which `SimdRadixN::new` +/// asserts, and `D` must divide the width. +#[inline(always)] +fn table_transpose( + height: usize, + reversed_columns: &[usize], + input: &[T], + output: &mut [T], +) { + let width = reversed_columns.len(); + assert!(input.len() == width * height && output.len() == input.len()); + + for (group, rev) in reversed_columns.chunks_exact(D).enumerate() { + let x = group * D; + let rev: &[usize; D] = rev.try_into().unwrap(); + for y in 0..height { + let row = x + y * width; + for (i, &r) in rev.iter().enumerate() { + unsafe { + *output.get_unchecked_mut(y + r * height) = *input.get_unchecked(row + i); + } + } + } + } +} + /// One cross-FFT layer: for each vector of columns, gather RADIX rows strided by `num_columns`, /// apply the twiddles, run the column butterfly, scatter back. /// From c8c268e806448a7af7953f44c04ce6977aba71b1 Mon Sep 17 00:00:00 2001 From: Henrik Date: Wed, 16 Sep 2026 21:52:13 +0200 Subject: [PATCH 08/22] Move the shared SIMD code into src/simd, and rename RadixNVector to SimdVector The trait gets its own simd_vector.rs so other algorithms can be written against it. Each backend's impls move next to its own vector trait impls in *_vector.rs, together with the fft_helper forwarding macro, leaving *_radixn.rs with just the type alias and tests. --- src/lib.rs | 13 +- src/neon/neon_common.rs | 59 ------- src/neon/neon_planner.rs | 2 +- src/neon/neon_radixn.rs | 189 +--------------------- src/neon/neon_vector.rs | 235 +++++++++++++++++++++++++++ src/simd/mod.rs | 5 + src/{ => simd}/simd_planner.rs | 0 src/{ => simd}/simd_radixn.rs | 125 +++------------ src/simd/simd_vector.rs | 93 +++++++++++ src/sse/sse_common.rs | 59 ------- src/sse/sse_planner.rs | 2 +- src/sse/sse_radixn.rs | 189 +--------------------- src/sse/sse_vector.rs | 235 +++++++++++++++++++++++++++ src/wasm_simd/wasm_simd_common.rs | 59 ------- src/wasm_simd/wasm_simd_planner.rs | 2 +- src/wasm_simd/wasm_simd_radixn.rs | 203 +---------------------- src/wasm_simd/wasm_simd_vector.rs | 249 +++++++++++++++++++++++++++++ 17 files changed, 863 insertions(+), 856 deletions(-) create mode 100644 src/simd/mod.rs rename src/{ => simd}/simd_planner.rs (100%) rename src/{ => simd}/simd_radixn.rs (83%) create mode 100644 src/simd/simd_vector.rs diff --git a/src/lib.rs b/src/lib.rs index 97f78bd9..11ce6132 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -130,21 +130,14 @@ mod math_utils; mod plan; mod twiddles; -// The SIMD backends all share one RadixN, generic over the vector type each of them provides +// Code shared by the SIMD backends: the `SimdVector` trait, the algorithms written against it, and +// the planner arithmetic that goes with them #[cfg(any( all(target_arch = "aarch64", feature = "neon"), all(target_arch = "x86_64", feature = "sse"), all(target_arch = "wasm32", feature = "wasm_simd"), ))] -mod simd_radixn; - -// ...and the planner arithmetic that goes with it -#[cfg(any( - all(target_arch = "aarch64", feature = "neon"), - all(target_arch = "x86_64", feature = "sse"), - all(target_arch = "wasm32", feature = "wasm_simd"), -))] -mod simd_planner; +mod simd; use num_complex::Complex; use num_traits::Zero; diff --git a/src/neon/neon_common.rs b/src/neon/neon_common.rs index 384b64a8..df31e9fe 100644 --- a/src/neon/neon_common.rs +++ b/src/neon/neon_common.rs @@ -144,65 +144,6 @@ macro_rules! boilerplate_fft_neon_oop { }; } -// The `RadixNVector::fft_helper_*` methods, which are the same forwarding calls for every NEON -// vector type: they hand the chunk loop to the target-feature-enabled wrappers below. -macro_rules! neon_radixn_fft_helpers { - () => { - #[inline(always)] - unsafe fn fft_helper_immut( - input: &[E], - output: &mut [E], - scratch: &mut [E], - chunk_size: usize, - required_scratch: usize, - chunk_fn: impl FnMut(&[E], &mut [E], &mut [E]), - ) { - super::neon_common::neon_fft_helper_immut( - input, - output, - scratch, - chunk_size, - required_scratch, - chunk_fn, - ) - } - #[inline(always)] - unsafe fn fft_helper_outofplace( - input: &mut [E], - output: &mut [E], - scratch: &mut [E], - chunk_size: usize, - required_scratch: usize, - chunk_fn: impl FnMut(&mut [E], &mut [E], &mut [E]), - ) { - super::neon_common::neon_fft_helper_outofplace( - input, - output, - scratch, - chunk_size, - required_scratch, - chunk_fn, - ) - } - #[inline(always)] - unsafe fn fft_helper_inplace( - buffer: &mut [E], - scratch: &mut [E], - chunk_size: usize, - required_scratch: usize, - chunk_fn: impl FnMut(&mut [E], &mut [E]), - ) { - super::neon_common::neon_fft_helper_inplace( - buffer, - scratch, - chunk_size, - required_scratch, - chunk_fn, - ) - } - }; -} - // A wrapper for the FFT helper functions that make sure the entire thing happens with the benefit of the NEON target feature, // so that things like loading twiddle factor registers etc can be lifted out of the loop #[target_feature(enable = "neon")] diff --git a/src/neon/neon_planner.rs b/src/neon/neon_planner.rs index 56840a9b..2328805f 100644 --- a/src/neon/neon_planner.rs +++ b/src/neon/neon_planner.rs @@ -18,7 +18,7 @@ use crate::neon::neon_radixn::*; use crate::Fft; use crate::math_utils::{PrimeFactor, PrimeFactors}; -use crate::simd_planner::{self, RadixNPlan}; +use crate::simd::simd_planner::{self, RadixNPlan}; const MIN_RADIX4_BITS: u32 = 6; // smallest size to consider radix 4 an option is 2^6 = 64 diff --git a/src/neon/neon_radixn.rs b/src/neon/neon_radixn.rs index d714bdae..62f07f4d 100644 --- a/src/neon/neon_radixn.rs +++ b/src/neon/neon_radixn.rs @@ -1,198 +1,21 @@ //! The NEON side of `SimdRadixN`. //! -//! The algorithm itself lives in `src/simd_radixn.rs`, shared by every SIMD backend. All that is -//! left here is the `RadixNVector` impl for each vector type: the NEON loads, stores and vector -//! math, and the element-type-specific butterfly structs for radix 3, 5, 6 and 7. +//! The algorithm itself lives in `src/simd/simd_radixn.rs`, shared by every SIMD backend, and the +//! `SimdVector` impls it runs on are in `neon_vector.rs`. All that is left here is the type alias +//! and the tests. -use std::arch::aarch64::{float32x4_t, float64x2_t}; +use crate::simd::simd_radixn::SimdRadixN; -use num_complex::Complex; - -use crate::simd_radixn::{RadixNVector, SimdRadixN}; -use crate::FftDirection; - -use super::neon_butterflies::{ - NeonF32Butterfly3, NeonF32Butterfly5, NeonF32Butterfly6, NeonF64Butterfly3, NeonF64Butterfly5, - NeonF64Butterfly6, -}; -use super::neon_prime_butterflies::{NeonF32Butterfly7, NeonF64Butterfly7}; -use super::neon_vector::{NeonArray, NeonArrayMut, NeonVector, Rotation90}; use super::NeonNum; /// FFT algorithm for lengths that factor into small radixes, NEON accelerated version. /// This is designed to be used via a Planner, and not created directly. pub type NeonRadixN = SimdRadixN<::VectorType, T>; -impl RadixNVector for float64x2_t { - const COMPLEX_PER_VECTOR: usize = 1; - - type ScalarType = f64; - type Rotation = Rotation90; - - type Butterfly3 = NeonF64Butterfly3; - type Butterfly5 = NeonF64Butterfly5; - type Butterfly6 = NeonF64Butterfly6; - type Butterfly7 = NeonF64Butterfly7; - - #[inline(always)] - unsafe fn load(data: &[Complex], index: usize) -> Self { - data.load_complex(index) - } - #[inline(always)] - unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { - data.store_complex(value, index) - } - - #[inline(always)] - unsafe fn mul_complex(left: Self, right: Self) -> Self { - NeonVector::mul_complex(left, right) - } - #[inline(always)] - unsafe fn make_mixedradix_twiddle_chunk( - x: usize, - y: usize, - len: usize, - direction: FftDirection, - ) -> Self { - NeonVector::make_mixedradix_twiddle_chunk(x, y, len, direction) - } - - #[inline(always)] - unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { - NeonVector::make_rotate90(direction) - } - #[inline(always)] - unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { - NeonF64Butterfly3::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { - NeonF64Butterfly5::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { - NeonF64Butterfly6::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { - NeonF64Butterfly7::new(direction) - } - - #[inline(always)] - unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { - NeonVector::column_butterfly2(rows) - } - #[inline(always)] - unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { - bf.perform_fft_direct(rows[0], rows[1], rows[2]) - } - #[inline(always)] - unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { - NeonVector::column_butterfly4(rows, rotation) - } - #[inline(always)] - unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { - bf.perform_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4]) - } - #[inline(always)] - unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { - bf.perform_fft_direct(rows) - } - #[inline(always)] - unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { - bf.perform_fft_direct(rows) - } - - neon_radixn_fft_helpers!(); -} - -impl RadixNVector for float32x4_t { - const COMPLEX_PER_VECTOR: usize = 2; - - type ScalarType = f32; - type Rotation = Rotation90; - - type Butterfly3 = NeonF32Butterfly3; - type Butterfly5 = NeonF32Butterfly5; - type Butterfly6 = NeonF32Butterfly6; - type Butterfly7 = NeonF32Butterfly7; - - #[inline(always)] - unsafe fn load(data: &[Complex], index: usize) -> Self { - data.load_complex(index) - } - #[inline(always)] - unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { - data.store_complex(value, index) - } - - #[inline(always)] - unsafe fn mul_complex(left: Self, right: Self) -> Self { - NeonVector::mul_complex(left, right) - } - #[inline(always)] - unsafe fn make_mixedradix_twiddle_chunk( - x: usize, - y: usize, - len: usize, - direction: FftDirection, - ) -> Self { - NeonVector::make_mixedradix_twiddle_chunk(x, y, len, direction) - } - - #[inline(always)] - unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { - NeonVector::make_rotate90(direction) - } - #[inline(always)] - unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { - NeonF32Butterfly3::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { - NeonF32Butterfly5::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { - NeonF32Butterfly6::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { - NeonF32Butterfly7::new(direction) - } - - #[inline(always)] - unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { - NeonVector::column_butterfly2(rows) - } - #[inline(always)] - unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { - bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2]) - } - #[inline(always)] - unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { - NeonVector::column_butterfly4(rows, rotation) - } - #[inline(always)] - unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { - bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4]) - } - #[inline(always)] - unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { - bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4], rows[5]) - } - #[inline(always)] - unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { - bf.perform_parallel_fft_direct(rows) - } - - neon_radixn_fft_helpers!(); -} - #[cfg(test)] mod unit_tests { - use super::*; - use crate::simd_radixn::test_bodies; + use crate::simd::simd_radixn::test_bodies; + use std::arch::aarch64::{float32x4_t, float64x2_t}; #[test] fn test_neon_radixn_f64() { diff --git a/src/neon/neon_vector.rs b/src/neon/neon_vector.rs index 035cf4db..167f47c6 100644 --- a/src/neon/neon_vector.rs +++ b/src/neon/neon_vector.rs @@ -6,6 +6,11 @@ use std::ops::{Deref, DerefMut}; use crate::{array_utils::DoubleBuf, twiddles, FftDirection}; +use super::neon_butterflies::{ + NeonF32Butterfly3, NeonF32Butterfly5, NeonF32Butterfly6, NeonF64Butterfly3, NeonF64Butterfly5, + NeonF64Butterfly6, +}; +use super::neon_prime_butterflies::{NeonF32Butterfly7, NeonF64Butterfly7}; use super::NeonNum; // Read these indexes from an NeonArray and build an array of simd vectors. @@ -548,6 +553,236 @@ where } } +// The `SimdVector` impls, which let this backend use the algorithms in `src/simd`. The trait is +// named by path instead of imported, because importing it would make methods like +// `Self::column_butterfly2` ambiguous with the backend's own vector trait. + +// The `SimdVector::fft_helper_*` methods, which are the same forwarding calls for every NEON +// vector type: they hand the chunk loop to the target-feature-enabled wrappers in +// `neon_common.rs`. +macro_rules! neon_vector_fft_helpers { + () => { + #[inline(always)] + unsafe fn fft_helper_immut( + input: &[E], + output: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&[E], &mut [E], &mut [E]), + ) { + super::neon_common::neon_fft_helper_immut( + input, + output, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + #[inline(always)] + unsafe fn fft_helper_outofplace( + input: &mut [E], + output: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&mut [E], &mut [E], &mut [E]), + ) { + super::neon_common::neon_fft_helper_outofplace( + input, + output, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + #[inline(always)] + unsafe fn fft_helper_inplace( + buffer: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&mut [E], &mut [E]), + ) { + super::neon_common::neon_fft_helper_inplace( + buffer, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + }; +} + +impl crate::simd::simd_vector::SimdVector for float64x2_t { + const COMPLEX_PER_VECTOR: usize = 1; + + type ScalarType = f64; + type Rotation = Rotation90; + + type Butterfly3 = NeonF64Butterfly3; + type Butterfly5 = NeonF64Butterfly5; + type Butterfly6 = NeonF64Butterfly6; + type Butterfly7 = NeonF64Butterfly7; + + #[inline(always)] + unsafe fn load(data: &[Complex], index: usize) -> Self { + data.load_complex(index) + } + #[inline(always)] + unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { + data.store_complex(value, index) + } + + #[inline(always)] + unsafe fn mul_complex(left: Self, right: Self) -> Self { + NeonVector::mul_complex(left, right) + } + #[inline(always)] + unsafe fn make_mixedradix_twiddle_chunk( + x: usize, + y: usize, + len: usize, + direction: FftDirection, + ) -> Self { + NeonVector::make_mixedradix_twiddle_chunk(x, y, len, direction) + } + + #[inline(always)] + unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { + NeonVector::make_rotate90(direction) + } + #[inline(always)] + unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { + NeonF64Butterfly3::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { + NeonF64Butterfly5::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { + NeonF64Butterfly6::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { + NeonF64Butterfly7::new(direction) + } + + #[inline(always)] + unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { + NeonVector::column_butterfly2(rows) + } + #[inline(always)] + unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { + bf.perform_fft_direct(rows[0], rows[1], rows[2]) + } + #[inline(always)] + unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { + NeonVector::column_butterfly4(rows, rotation) + } + #[inline(always)] + unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { + bf.perform_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4]) + } + #[inline(always)] + unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { + bf.perform_fft_direct(rows) + } + #[inline(always)] + unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { + bf.perform_fft_direct(rows) + } + + neon_vector_fft_helpers!(); +} + +impl crate::simd::simd_vector::SimdVector for float32x4_t { + const COMPLEX_PER_VECTOR: usize = 2; + + type ScalarType = f32; + type Rotation = Rotation90; + + type Butterfly3 = NeonF32Butterfly3; + type Butterfly5 = NeonF32Butterfly5; + type Butterfly6 = NeonF32Butterfly6; + type Butterfly7 = NeonF32Butterfly7; + + #[inline(always)] + unsafe fn load(data: &[Complex], index: usize) -> Self { + data.load_complex(index) + } + #[inline(always)] + unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { + data.store_complex(value, index) + } + + #[inline(always)] + unsafe fn mul_complex(left: Self, right: Self) -> Self { + NeonVector::mul_complex(left, right) + } + #[inline(always)] + unsafe fn make_mixedradix_twiddle_chunk( + x: usize, + y: usize, + len: usize, + direction: FftDirection, + ) -> Self { + NeonVector::make_mixedradix_twiddle_chunk(x, y, len, direction) + } + + #[inline(always)] + unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { + NeonVector::make_rotate90(direction) + } + #[inline(always)] + unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { + NeonF32Butterfly3::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { + NeonF32Butterfly5::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { + NeonF32Butterfly6::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { + NeonF32Butterfly7::new(direction) + } + + #[inline(always)] + unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { + NeonVector::column_butterfly2(rows) + } + #[inline(always)] + unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { + bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2]) + } + #[inline(always)] + unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { + NeonVector::column_butterfly4(rows, rotation) + } + #[inline(always)] + unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { + bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4]) + } + #[inline(always)] + unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { + bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4], rows[5]) + } + #[inline(always)] + unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { + bf.perform_parallel_fft_direct(rows) + } + + neon_vector_fft_helpers!(); +} + #[cfg(test)] mod unit_tests { use super::*; diff --git a/src/simd/mod.rs b/src/simd/mod.rs new file mode 100644 index 00000000..22d7c175 --- /dev/null +++ b/src/simd/mod.rs @@ -0,0 +1,5 @@ +//! Code shared by every SIMD backend, written once against the `SimdVector` trait. + +pub mod simd_planner; +pub mod simd_radixn; +pub mod simd_vector; diff --git a/src/simd_planner.rs b/src/simd/simd_planner.rs similarity index 100% rename from src/simd_planner.rs rename to src/simd/simd_planner.rs diff --git a/src/simd_radixn.rs b/src/simd/simd_radixn.rs similarity index 83% rename from src/simd_radixn.rs rename to src/simd/simd_radixn.rs index 29616446..fd3f401c 100644 --- a/src/simd_radixn.rs +++ b/src/simd/simd_radixn.rs @@ -5,9 +5,8 @@ //! cross-FFT layers use SIMD column butterflies instead of the scalar ones, so a whole vector of //! columns is processed per butterfly call. //! -//! Everything here is generic over `RadixNVector`, which is the handful of vector operations the -//! algorithm needs. A backend implements that trait once per vector type and gets the algorithm, -//! so `SimdRadixN` is the only copy of it. +//! Everything here is generic over `SimdVector`, from `simd_vector.rs`. A backend implements that +//! trait once per vector type and gets the algorithm, so `SimdRadixN` is the only copy of it. //! //! Because a column butterfly consumes `COMPLEX_PER_VECTOR` columns at a time, the column count at //! every layer has to be a whole number of vectors. The column count starts at `base_len` and only @@ -23,90 +22,10 @@ use crate::array_utils::{reverse_remainders, workaround_transmute_mut, Transpose use crate::common::{FftNum, RadixFactor}; use crate::{Direction, Fft, FftDirection, Length}; -/// Everything `SimdRadixN` needs from a backend's vector type. -/// -/// Radix 2 and 4 are vector-generic in every backend, so they map straight onto the backend's -/// vector trait. Radix 3, 5, 6 and 7 only exist as element-type-specific structs holding -/// precomputed twiddles, so this trait pairs each vector type with its own set of them. For f64 a -/// vector is one complex number and the plain `perform_fft_direct` is already a single column; for -/// f32 a vector is two complex numbers and `perform_parallel_fft_direct` does two columns at once. -/// -/// Safety: every method here requires the current machine to support the backend's SIMD -/// instruction set. -pub trait RadixNVector: Copy + Send + Sync + Sized { - const COMPLEX_PER_VECTOR: usize; - - /// The scalar this vector holds. Always the same type as the `T` of the `SimdRadixN` using it. - type ScalarType: FftNum; - - /// The backend's precomputed 90 degree rotation, which the radix 4 butterfly needs. - type Rotation: Copy + Send + Sync; - - type Butterfly3: Send + Sync; - type Butterfly5: Send + Sync; - type Butterfly6: Send + Sync; - type Butterfly7: Send + Sync; - - unsafe fn load(data: &[Complex], index: usize) -> Self; - unsafe fn store(data: &mut [Complex], value: Self, index: usize); - - /// Pairwise multiply the complex numbers in `left` with the complex numbers in `right`. - unsafe fn mul_complex(left: Self, right: Self) -> Self; - - /// Generates a chunk of twiddle factors starting at (X,Y) and incrementing X - /// `COMPLEX_PER_VECTOR` times. - unsafe fn make_mixedradix_twiddle_chunk( - x: usize, - y: usize, - len: usize, - direction: FftDirection, - ) -> Self; - - unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation; - unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3; - unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5; - unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6; - unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7; - - /// Each of these interprets the input as rows of a `COMPLEX_PER_VECTOR`-by-N 2D array, and - /// computes parallel butterflies down the columns of the 2D array. - unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2]; - unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3]; - unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4]; - unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5]; - unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6]; - unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7]; - - /// The three `fft_helper_*` wrappers from the backend's `*_common.rs`, which run the whole - /// chunk loop with the backend's target feature enabled so that things like loading twiddle - /// factor registers can be lifted out of the loop. - unsafe fn fft_helper_immut( - input: &[E], - output: &mut [E], - scratch: &mut [E], - chunk_size: usize, - required_scratch: usize, - chunk_fn: impl FnMut(&[E], &mut [E], &mut [E]), - ); - unsafe fn fft_helper_outofplace( - input: &mut [E], - output: &mut [E], - scratch: &mut [E], - chunk_size: usize, - required_scratch: usize, - chunk_fn: impl FnMut(&mut [E], &mut [E], &mut [E]), - ); - unsafe fn fft_helper_inplace( - buffer: &mut [E], - scratch: &mut [E], - chunk_size: usize, - required_scratch: usize, - chunk_fn: impl FnMut(&mut [E], &mut [E]), - ); -} +use super::simd_vector::SimdVector; /// The per-layer cross-FFT kernels, holding whatever precomputed state each radix needs. -enum InternalRadixFactor { +enum InternalRadixFactor { Factor2, Factor3(V::Butterfly3), Factor4(V::Rotation), @@ -115,7 +34,7 @@ enum InternalRadixFactor { Factor7(V::Butterfly7), } -impl InternalRadixFactor { +impl InternalRadixFactor { fn radix(&self) -> usize { match self { InternalRadixFactor::Factor2 => 2, @@ -130,7 +49,7 @@ impl InternalRadixFactor { /// FFT algorithm for lengths that factor into small radixes, SIMD accelerated version. /// This is designed to be used via a Planner, and not created directly. -pub struct SimdRadixN { +pub struct SimdRadixN { twiddles: Box<[V]>, base_fft: Arc>, @@ -150,7 +69,7 @@ pub struct SimdRadixN { immut_scratch_len: usize, } -impl SimdRadixN { +impl SimdRadixN { /// Constructs a SimdRadixN which computes FFTs of length `factor_product * base_fft.len()`. pub fn new(factors: &[RadixFactor], base_fft: Arc>) -> Self { // Internal sanity check: Make sure that the vector's scalar type is T. @@ -394,7 +313,7 @@ impl SimdRadixN { } } -impl Fft for SimdRadixN { +impl Fft for SimdRadixN { fn process_immutable_with_scratch( &self, input: &[Complex], @@ -459,13 +378,13 @@ impl Fft for SimdRadixN { self.immut_scratch_len } } -impl Length for SimdRadixN { +impl Length for SimdRadixN { #[inline(always)] fn len(&self) -> usize { self.len } } -impl Direction for SimdRadixN { +impl Direction for SimdRadixN { #[inline(always)] fn fft_direction(&self) -> FftDirection { self.direction @@ -477,7 +396,7 @@ impl Direction for SimdRadixN { /// This is `chunks_exact_mut` without the divide it does to find the chunk count. At short lengths /// that one divide per layer is a measurable share of the whole FFT. #[inline(always)] -unsafe fn cross_layer_chunks( +unsafe fn cross_layer_chunks( data: &mut [Complex], twiddles: &[V], num_columns: usize, @@ -530,7 +449,7 @@ fn table_transpose( /// Unrolled two vectors at a time, which is what the SIMD Radix4's `butterfly_4` does and is what /// gets the two independent dependency chains needed to keep the FMA pipeline busy. #[inline(always)] -unsafe fn cross_layer( +unsafe fn cross_layer( data: &mut [Complex], twiddles: &[V], num_columns: usize, @@ -606,7 +525,7 @@ pub mod test_bodies { /// Every empty, one-factor and two-factor recipe over each of `bases`, both directions. pub fn factor_pairs(bases: &[usize]) where - V: RadixNVector, + V: SimdVector, V::ScalarType: Float + SampleUniform, { for base in bases { @@ -634,8 +553,8 @@ pub mod test_bodies { /// has factors above 7 (for example 11 * 13 = 143). pub fn composite_base() where - V32: RadixNVector, - V64: RadixNVector, + V32: SimdVector, + V64: SimdVector, { let mut planner64 = crate::FftPlannerScalar::::new(); let mut planner32 = crate::FftPlannerScalar::::new(); @@ -674,8 +593,8 @@ pub mod test_bodies { /// naive `Dft` the result is checked against affordable. pub fn large_recipes() where - V32: RadixNVector, - V64: RadixNVector, + V32: SimdVector, + V64: SimdVector, { use RadixFactor::*; // (factors, f64 base, f32 base). f32 needs an even base, so it gets its own. @@ -695,8 +614,8 @@ pub mod test_bodies { /// `cargo test --release -- --ignored radixn_six_layers`. pub fn six_layers() where - V32: RadixNVector, - V64: RadixNVector, + V32: SimdVector, + V64: SimdVector, { use RadixFactor::*; let cases: [(&[RadixFactor], usize, usize); 1] = [( @@ -710,8 +629,8 @@ pub mod test_bodies { /// Runs each (factors, f64 base, f32 base) case in both directions. fn recipes(cases: &[(&[RadixFactor], usize, usize)]) where - V32: RadixNVector, - V64: RadixNVector, + V32: SimdVector, + V64: SimdVector, { for (factors, base64, base32) in cases { for direction in [FftDirection::Forward, FftDirection::Inverse] { @@ -723,7 +642,7 @@ pub mod test_bodies { fn check(factors: &[RadixFactor], base_fft: Arc>) where - V: RadixNVector, + V: SimdVector, V::ScalarType: Float + SampleUniform, { let len = base_fft.len() * factors.iter().map(|f| f.radix()).product::(); diff --git a/src/simd/simd_vector.rs b/src/simd/simd_vector.rs new file mode 100644 index 00000000..6c8c572d --- /dev/null +++ b/src/simd/simd_vector.rs @@ -0,0 +1,93 @@ +//! The vector operations shared by every SIMD backend. +//! +//! Each backend has its own vector trait (`NeonVector`, `SseVector`, `WasmVector`), and those +//! don't share a supertrait. `SimdVector` is a separate, stripped down trait that each backend +//! implements once per vector type, next to its own vector trait impls. Algorithms written against +//! it, like `SimdRadixN`, then only exist once. + +use num_complex::Complex; + +use crate::common::FftNum; +use crate::FftDirection; + +/// The vector operations the shared SIMD algorithms need from a backend's vector type. +/// +/// Radix 2 and 4 are vector-generic in every backend, so they map straight onto the backend's +/// vector trait. Radix 3, 5, 6 and 7 only exist as element-type-specific structs holding +/// precomputed twiddles, so this trait pairs each vector type with its own set of them. For f64 a +/// vector is one complex number and the plain `perform_fft_direct` is already a single column; for +/// f32 a vector is two complex numbers and `perform_parallel_fft_direct` does two columns at once. +/// +/// Safety: every method here requires the current machine to support the backend's SIMD +/// instruction set. +pub trait SimdVector: Copy + Send + Sync + Sized { + const COMPLEX_PER_VECTOR: usize; + + /// The scalar this vector holds. Always the same type as the `T` of the algorithm using it. + type ScalarType: FftNum; + + /// The backend's precomputed 90 degree rotation, which the radix 4 butterfly needs. + type Rotation: Copy + Send + Sync; + + type Butterfly3: Send + Sync; + type Butterfly5: Send + Sync; + type Butterfly6: Send + Sync; + type Butterfly7: Send + Sync; + + unsafe fn load(data: &[Complex], index: usize) -> Self; + unsafe fn store(data: &mut [Complex], value: Self, index: usize); + + /// Pairwise multiply the complex numbers in `left` with the complex numbers in `right`. + unsafe fn mul_complex(left: Self, right: Self) -> Self; + + /// Generates a chunk of twiddle factors starting at (X,Y) and incrementing X + /// `COMPLEX_PER_VECTOR` times. + unsafe fn make_mixedradix_twiddle_chunk( + x: usize, + y: usize, + len: usize, + direction: FftDirection, + ) -> Self; + + unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation; + unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3; + unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5; + unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6; + unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7; + + /// Each of these interprets the input as rows of a `COMPLEX_PER_VECTOR`-by-N 2D array, and + /// computes parallel butterflies down the columns of the 2D array. + unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2]; + unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3]; + unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4]; + unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5]; + unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6]; + unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7]; + + /// The three `fft_helper_*` wrappers from the backend's `*_common.rs`, which run the whole + /// chunk loop with the backend's target feature enabled so that things like loading twiddle + /// factor registers can be lifted out of the loop. + unsafe fn fft_helper_immut( + input: &[E], + output: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&[E], &mut [E], &mut [E]), + ); + unsafe fn fft_helper_outofplace( + input: &mut [E], + output: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&mut [E], &mut [E], &mut [E]), + ); + unsafe fn fft_helper_inplace( + buffer: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&mut [E], &mut [E]), + ); +} diff --git a/src/sse/sse_common.rs b/src/sse/sse_common.rs index a7b5a1b7..78b26739 100644 --- a/src/sse/sse_common.rs +++ b/src/sse/sse_common.rs @@ -144,65 +144,6 @@ macro_rules! boilerplate_fft_sse_oop { }; } -// The `RadixNVector::fft_helper_*` methods, which are the same forwarding calls for every SSE -// vector type: they hand the chunk loop to the target-feature-enabled wrappers below. -macro_rules! sse_radixn_fft_helpers { - () => { - #[inline(always)] - unsafe fn fft_helper_immut( - input: &[E], - output: &mut [E], - scratch: &mut [E], - chunk_size: usize, - required_scratch: usize, - chunk_fn: impl FnMut(&[E], &mut [E], &mut [E]), - ) { - super::sse_common::sse_fft_helper_immut( - input, - output, - scratch, - chunk_size, - required_scratch, - chunk_fn, - ) - } - #[inline(always)] - unsafe fn fft_helper_outofplace( - input: &mut [E], - output: &mut [E], - scratch: &mut [E], - chunk_size: usize, - required_scratch: usize, - chunk_fn: impl FnMut(&mut [E], &mut [E], &mut [E]), - ) { - super::sse_common::sse_fft_helper_outofplace( - input, - output, - scratch, - chunk_size, - required_scratch, - chunk_fn, - ) - } - #[inline(always)] - unsafe fn fft_helper_inplace( - buffer: &mut [E], - scratch: &mut [E], - chunk_size: usize, - required_scratch: usize, - chunk_fn: impl FnMut(&mut [E], &mut [E]), - ) { - super::sse_common::sse_fft_helper_inplace( - buffer, - scratch, - chunk_size, - required_scratch, - chunk_fn, - ) - } - }; -} - // A wrapper for the FFT helper functions that make sure the entire thing happens with the benefit of the SSE target feature, // so that things like loading twiddle factor registers etc can be lifted out of the loop #[target_feature(enable = "sse4.1")] diff --git a/src/sse/sse_planner.rs b/src/sse/sse_planner.rs index 8df9dc6c..bfd3721a 100644 --- a/src/sse/sse_planner.rs +++ b/src/sse/sse_planner.rs @@ -18,7 +18,7 @@ use crate::sse::sse_radixn::*; use crate::Fft; use crate::math_utils::{PrimeFactor, PrimeFactors}; -use crate::simd_planner::{self, RadixNPlan}; +use crate::simd::simd_planner::{self, RadixNPlan}; const MIN_RADIX4_BITS: u32 = 6; // smallest size to consider radix 4 an option is 2^6 = 64 diff --git a/src/sse/sse_radixn.rs b/src/sse/sse_radixn.rs index 45feaf2b..c1eba04a 100644 --- a/src/sse/sse_radixn.rs +++ b/src/sse/sse_radixn.rs @@ -1,198 +1,21 @@ //! The SSE side of `SimdRadixN`. //! -//! The algorithm itself lives in `src/simd_radixn.rs`, shared by every SIMD backend. All that is -//! left here is the `RadixNVector` impl for each vector type: the SSE loads, stores and vector -//! math, and the element-type-specific butterfly structs for radix 3, 5, 6 and 7. +//! The algorithm itself lives in `src/simd/simd_radixn.rs`, shared by every SIMD backend, and the +//! `SimdVector` impls it runs on are in `sse_vector.rs`. All that is left here is the type alias +//! and the tests. -use std::arch::x86_64::{__m128, __m128d}; +use crate::simd::simd_radixn::SimdRadixN; -use num_complex::Complex; - -use crate::simd_radixn::{RadixNVector, SimdRadixN}; -use crate::FftDirection; - -use super::sse_butterflies::{ - SseF32Butterfly3, SseF32Butterfly5, SseF32Butterfly6, SseF64Butterfly3, SseF64Butterfly5, - SseF64Butterfly6, -}; -use super::sse_prime_butterflies::{SseF32Butterfly7, SseF64Butterfly7}; -use super::sse_vector::{Rotation90, SseArray, SseArrayMut, SseVector}; use super::SseNum; /// FFT algorithm for lengths that factor into small radixes, SSE accelerated version. /// This is designed to be used via a Planner, and not created directly. pub type SseRadixN = SimdRadixN<::VectorType, T>; -impl RadixNVector for __m128d { - const COMPLEX_PER_VECTOR: usize = 1; - - type ScalarType = f64; - type Rotation = Rotation90; - - type Butterfly3 = SseF64Butterfly3; - type Butterfly5 = SseF64Butterfly5; - type Butterfly6 = SseF64Butterfly6; - type Butterfly7 = SseF64Butterfly7; - - #[inline(always)] - unsafe fn load(data: &[Complex], index: usize) -> Self { - data.load_complex(index) - } - #[inline(always)] - unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { - data.store_complex(value, index) - } - - #[inline(always)] - unsafe fn mul_complex(left: Self, right: Self) -> Self { - SseVector::mul_complex(left, right) - } - #[inline(always)] - unsafe fn make_mixedradix_twiddle_chunk( - x: usize, - y: usize, - len: usize, - direction: FftDirection, - ) -> Self { - SseVector::make_mixedradix_twiddle_chunk(x, y, len, direction) - } - - #[inline(always)] - unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { - SseVector::make_rotate90(direction) - } - #[inline(always)] - unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { - SseF64Butterfly3::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { - SseF64Butterfly5::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { - SseF64Butterfly6::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { - SseF64Butterfly7::new(direction) - } - - #[inline(always)] - unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { - SseVector::column_butterfly2(rows) - } - #[inline(always)] - unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { - bf.perform_fft_direct(rows[0], rows[1], rows[2]) - } - #[inline(always)] - unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { - SseVector::column_butterfly4(rows, rotation) - } - #[inline(always)] - unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { - bf.perform_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4]) - } - #[inline(always)] - unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { - bf.perform_fft_direct(rows) - } - #[inline(always)] - unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { - bf.perform_fft_direct(rows) - } - - sse_radixn_fft_helpers!(); -} - -impl RadixNVector for __m128 { - const COMPLEX_PER_VECTOR: usize = 2; - - type ScalarType = f32; - type Rotation = Rotation90; - - type Butterfly3 = SseF32Butterfly3; - type Butterfly5 = SseF32Butterfly5; - type Butterfly6 = SseF32Butterfly6; - type Butterfly7 = SseF32Butterfly7; - - #[inline(always)] - unsafe fn load(data: &[Complex], index: usize) -> Self { - data.load_complex(index) - } - #[inline(always)] - unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { - data.store_complex(value, index) - } - - #[inline(always)] - unsafe fn mul_complex(left: Self, right: Self) -> Self { - SseVector::mul_complex(left, right) - } - #[inline(always)] - unsafe fn make_mixedradix_twiddle_chunk( - x: usize, - y: usize, - len: usize, - direction: FftDirection, - ) -> Self { - SseVector::make_mixedradix_twiddle_chunk(x, y, len, direction) - } - - #[inline(always)] - unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { - SseVector::make_rotate90(direction) - } - #[inline(always)] - unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { - SseF32Butterfly3::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { - SseF32Butterfly5::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { - SseF32Butterfly6::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { - SseF32Butterfly7::new(direction) - } - - #[inline(always)] - unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { - SseVector::column_butterfly2(rows) - } - #[inline(always)] - unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { - bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2]) - } - #[inline(always)] - unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { - SseVector::column_butterfly4(rows, rotation) - } - #[inline(always)] - unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { - bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4]) - } - #[inline(always)] - unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { - bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4], rows[5]) - } - #[inline(always)] - unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { - bf.perform_parallel_fft_direct(rows) - } - - sse_radixn_fft_helpers!(); -} - #[cfg(test)] mod unit_tests { - use super::*; - use crate::simd_radixn::test_bodies; + use crate::simd::simd_radixn::test_bodies; + use std::arch::x86_64::{__m128, __m128d}; #[test] fn test_sse_radixn_f64() { diff --git a/src/sse/sse_vector.rs b/src/sse/sse_vector.rs index 66daa4c1..85b16962 100644 --- a/src/sse/sse_vector.rs +++ b/src/sse/sse_vector.rs @@ -7,6 +7,11 @@ use std::ops::{Deref, DerefMut}; use crate::array_utils::DoubleBuf; use crate::{twiddles, FftDirection}; +use super::sse_butterflies::{ + SseF32Butterfly3, SseF32Butterfly5, SseF32Butterfly6, SseF64Butterfly3, SseF64Butterfly5, + SseF64Butterfly6, +}; +use super::sse_prime_butterflies::{SseF32Butterfly7, SseF64Butterfly7}; use super::SseNum; // Read these indexes from an SseArray and build an array of simd vectors. @@ -540,3 +545,233 @@ where self.output.store_partial_lo_complex(vector, index); } } + +// The `SimdVector` impls, which let this backend use the algorithms in `src/simd`. The trait is +// named by path instead of imported, because importing it would make methods like +// `Self::column_butterfly2` ambiguous with the backend's own vector trait. + +// The `SimdVector::fft_helper_*` methods, which are the same forwarding calls for every SSE +// vector type: they hand the chunk loop to the target-feature-enabled wrappers in +// `sse_common.rs`. +macro_rules! sse_vector_fft_helpers { + () => { + #[inline(always)] + unsafe fn fft_helper_immut( + input: &[E], + output: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&[E], &mut [E], &mut [E]), + ) { + super::sse_common::sse_fft_helper_immut( + input, + output, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + #[inline(always)] + unsafe fn fft_helper_outofplace( + input: &mut [E], + output: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&mut [E], &mut [E], &mut [E]), + ) { + super::sse_common::sse_fft_helper_outofplace( + input, + output, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + #[inline(always)] + unsafe fn fft_helper_inplace( + buffer: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&mut [E], &mut [E]), + ) { + super::sse_common::sse_fft_helper_inplace( + buffer, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + }; +} + +impl crate::simd::simd_vector::SimdVector for __m128d { + const COMPLEX_PER_VECTOR: usize = 1; + + type ScalarType = f64; + type Rotation = Rotation90; + + type Butterfly3 = SseF64Butterfly3; + type Butterfly5 = SseF64Butterfly5; + type Butterfly6 = SseF64Butterfly6; + type Butterfly7 = SseF64Butterfly7; + + #[inline(always)] + unsafe fn load(data: &[Complex], index: usize) -> Self { + data.load_complex(index) + } + #[inline(always)] + unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { + data.store_complex(value, index) + } + + #[inline(always)] + unsafe fn mul_complex(left: Self, right: Self) -> Self { + SseVector::mul_complex(left, right) + } + #[inline(always)] + unsafe fn make_mixedradix_twiddle_chunk( + x: usize, + y: usize, + len: usize, + direction: FftDirection, + ) -> Self { + SseVector::make_mixedradix_twiddle_chunk(x, y, len, direction) + } + + #[inline(always)] + unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { + SseVector::make_rotate90(direction) + } + #[inline(always)] + unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { + SseF64Butterfly3::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { + SseF64Butterfly5::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { + SseF64Butterfly6::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { + SseF64Butterfly7::new(direction) + } + + #[inline(always)] + unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { + SseVector::column_butterfly2(rows) + } + #[inline(always)] + unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { + bf.perform_fft_direct(rows[0], rows[1], rows[2]) + } + #[inline(always)] + unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { + SseVector::column_butterfly4(rows, rotation) + } + #[inline(always)] + unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { + bf.perform_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4]) + } + #[inline(always)] + unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { + bf.perform_fft_direct(rows) + } + #[inline(always)] + unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { + bf.perform_fft_direct(rows) + } + + sse_vector_fft_helpers!(); +} + +impl crate::simd::simd_vector::SimdVector for __m128 { + const COMPLEX_PER_VECTOR: usize = 2; + + type ScalarType = f32; + type Rotation = Rotation90; + + type Butterfly3 = SseF32Butterfly3; + type Butterfly5 = SseF32Butterfly5; + type Butterfly6 = SseF32Butterfly6; + type Butterfly7 = SseF32Butterfly7; + + #[inline(always)] + unsafe fn load(data: &[Complex], index: usize) -> Self { + data.load_complex(index) + } + #[inline(always)] + unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { + data.store_complex(value, index) + } + + #[inline(always)] + unsafe fn mul_complex(left: Self, right: Self) -> Self { + SseVector::mul_complex(left, right) + } + #[inline(always)] + unsafe fn make_mixedradix_twiddle_chunk( + x: usize, + y: usize, + len: usize, + direction: FftDirection, + ) -> Self { + SseVector::make_mixedradix_twiddle_chunk(x, y, len, direction) + } + + #[inline(always)] + unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { + SseVector::make_rotate90(direction) + } + #[inline(always)] + unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { + SseF32Butterfly3::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { + SseF32Butterfly5::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { + SseF32Butterfly6::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { + SseF32Butterfly7::new(direction) + } + + #[inline(always)] + unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { + SseVector::column_butterfly2(rows) + } + #[inline(always)] + unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { + bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2]) + } + #[inline(always)] + unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { + SseVector::column_butterfly4(rows, rotation) + } + #[inline(always)] + unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { + bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4]) + } + #[inline(always)] + unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { + bf.perform_parallel_fft_direct(rows[0], rows[1], rows[2], rows[3], rows[4], rows[5]) + } + #[inline(always)] + unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { + bf.perform_parallel_fft_direct(rows) + } + + sse_vector_fft_helpers!(); +} diff --git a/src/wasm_simd/wasm_simd_common.rs b/src/wasm_simd/wasm_simd_common.rs index fc797339..af32692d 100644 --- a/src/wasm_simd/wasm_simd_common.rs +++ b/src/wasm_simd/wasm_simd_common.rs @@ -145,65 +145,6 @@ macro_rules! boilerplate_fft_wasm_simd_oop { }; } -// The `RadixNVector::fft_helper_*` methods, which are the same forwarding calls for every WASM -// SIMD vector type: they hand the chunk loop to the target-feature-enabled wrappers below. -macro_rules! wasm_simd_radixn_fft_helpers { - () => { - #[inline(always)] - unsafe fn fft_helper_immut( - input: &[E], - output: &mut [E], - scratch: &mut [E], - chunk_size: usize, - required_scratch: usize, - chunk_fn: impl FnMut(&[E], &mut [E], &mut [E]), - ) { - super::wasm_simd_common::wasm_simd_fft_helper_immut( - input, - output, - scratch, - chunk_size, - required_scratch, - chunk_fn, - ) - } - #[inline(always)] - unsafe fn fft_helper_outofplace( - input: &mut [E], - output: &mut [E], - scratch: &mut [E], - chunk_size: usize, - required_scratch: usize, - chunk_fn: impl FnMut(&mut [E], &mut [E], &mut [E]), - ) { - super::wasm_simd_common::wasm_simd_fft_helper_outofplace( - input, - output, - scratch, - chunk_size, - required_scratch, - chunk_fn, - ) - } - #[inline(always)] - unsafe fn fft_helper_inplace( - buffer: &mut [E], - scratch: &mut [E], - chunk_size: usize, - required_scratch: usize, - chunk_fn: impl FnMut(&mut [E], &mut [E]), - ) { - super::wasm_simd_common::wasm_simd_fft_helper_inplace( - buffer, - scratch, - chunk_size, - required_scratch, - chunk_fn, - ) - } - }; -} - // A wrapper for the FFT helper functions that make sure the entire thing happens with the benefit of the Wasm SIMD target feature, // so that things like loading twiddle factor registers etc can be lifted out of the loop #[target_feature(enable = "simd128")] diff --git a/src/wasm_simd/wasm_simd_planner.rs b/src/wasm_simd/wasm_simd_planner.rs index 6ccc3ae3..e70a62ba 100644 --- a/src/wasm_simd/wasm_simd_planner.rs +++ b/src/wasm_simd/wasm_simd_planner.rs @@ -6,7 +6,7 @@ use crate::algorithm::{ }; use crate::common::RadixFactor; use crate::math_utils::PrimeFactor; -use crate::simd_planner::{self, RadixNPlan}; +use crate::simd::simd_planner::{self, RadixNPlan}; use crate::wasm_simd::*; use crate::{fft_cache::FftCache, math_utils::PrimeFactors, Fft, FftDirection, FftNum}; use std::{any::TypeId, collections::HashMap, sync::Arc}; diff --git a/src/wasm_simd/wasm_simd_radixn.rs b/src/wasm_simd/wasm_simd_radixn.rs index 6f43cd4d..24aeb8f0 100644 --- a/src/wasm_simd/wasm_simd_radixn.rs +++ b/src/wasm_simd/wasm_simd_radixn.rs @@ -1,212 +1,21 @@ //! The WASM SIMD side of `SimdRadixN`. //! -//! The algorithm itself lives in `src/simd_radixn.rs`, shared by every SIMD backend. All that is -//! left here is the `RadixNVector` impl for each vector type: the WASM SIMD loads, stores and -//! vector math, and the element-type-specific butterfly structs for radix 3, 5, 6 and 7. +//! The algorithm itself lives in `src/simd/simd_radixn.rs`, shared by every SIMD backend, and the +//! `SimdVector` impls it runs on are in `wasm_simd_vector.rs`. All that is left here is the type +//! alias and the tests. -use num_complex::Complex; +use crate::simd::simd_radixn::SimdRadixN; -use crate::simd_radixn::{RadixNVector, SimdRadixN}; -use crate::FftDirection; - -use super::wasm_simd_butterflies::{ - WasmSimdF32Butterfly3, WasmSimdF32Butterfly5, WasmSimdF32Butterfly6, WasmSimdF64Butterfly3, - WasmSimdF64Butterfly5, WasmSimdF64Butterfly6, -}; -use super::wasm_simd_prime_butterflies::{WasmSimdF32Butterfly7, WasmSimdF64Butterfly7}; -use super::wasm_simd_vector::{ - Rotation90, WasmSimdArray, WasmSimdArrayMut, WasmVector, WasmVector32, WasmVector64, -}; use super::WasmNum; /// FFT algorithm for lengths that factor into small radixes, WASM SIMD accelerated version. /// This is designed to be used via a Planner, and not created directly. pub type WasmSimdRadixN = SimdRadixN<::VectorType, T>; -impl RadixNVector for WasmVector64 { - const COMPLEX_PER_VECTOR: usize = 1; - - type ScalarType = f64; - type Rotation = Rotation90; - - type Butterfly3 = WasmSimdF64Butterfly3; - type Butterfly5 = WasmSimdF64Butterfly5; - type Butterfly6 = WasmSimdF64Butterfly6; - type Butterfly7 = WasmSimdF64Butterfly7; - - #[inline(always)] - unsafe fn load(data: &[Complex], index: usize) -> Self { - data.load_complex(index) - } - #[inline(always)] - unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { - data.store_complex(value, index) - } - - #[inline(always)] - unsafe fn mul_complex(left: Self, right: Self) -> Self { - WasmVector::mul_complex(left, right) - } - #[inline(always)] - unsafe fn make_mixedradix_twiddle_chunk( - x: usize, - y: usize, - len: usize, - direction: FftDirection, - ) -> Self { - WasmVector::make_mixedradix_twiddle_chunk(x, y, len, direction) - } - - #[inline(always)] - unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { - WasmVector::make_rotate90(direction) - } - #[inline(always)] - unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { - WasmSimdF64Butterfly3::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { - WasmSimdF64Butterfly5::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { - WasmSimdF64Butterfly6::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { - WasmSimdF64Butterfly7::new(direction) - } - - #[inline(always)] - unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { - WasmVector::column_butterfly2(rows) - } - // Butterflies 3, 5 and 6 are written against raw `v128` while `WasmVector64` is a newtype over - // it, so their results come back needing rewrapping. Butterfly 7 already speaks the wrapper - // types and needs none of it. - #[inline(always)] - unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { - bf.perform_fft_direct(rows[0].0, rows[1].0, rows[2].0) - .map(f64::wrap) - } - #[inline(always)] - unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { - WasmVector::column_butterfly4(rows, rotation) - } - #[inline(always)] - unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { - bf.perform_fft_direct(rows[0].0, rows[1].0, rows[2].0, rows[3].0, rows[4].0) - .map(f64::wrap) - } - #[inline(always)] - unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { - bf.perform_fft_direct([ - rows[0].0, rows[1].0, rows[2].0, rows[3].0, rows[4].0, rows[5].0, - ]) - .map(f64::wrap) - } - #[inline(always)] - unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { - bf.perform_fft_direct(rows) - } - - wasm_simd_radixn_fft_helpers!(); -} - -impl RadixNVector for WasmVector32 { - const COMPLEX_PER_VECTOR: usize = 2; - - type ScalarType = f32; - type Rotation = Rotation90; - - type Butterfly3 = WasmSimdF32Butterfly3; - type Butterfly5 = WasmSimdF32Butterfly5; - type Butterfly6 = WasmSimdF32Butterfly6; - type Butterfly7 = WasmSimdF32Butterfly7; - - #[inline(always)] - unsafe fn load(data: &[Complex], index: usize) -> Self { - data.load_complex(index) - } - #[inline(always)] - unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { - data.store_complex(value, index) - } - - #[inline(always)] - unsafe fn mul_complex(left: Self, right: Self) -> Self { - WasmVector::mul_complex(left, right) - } - #[inline(always)] - unsafe fn make_mixedradix_twiddle_chunk( - x: usize, - y: usize, - len: usize, - direction: FftDirection, - ) -> Self { - WasmVector::make_mixedradix_twiddle_chunk(x, y, len, direction) - } - - #[inline(always)] - unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { - WasmVector::make_rotate90(direction) - } - #[inline(always)] - unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { - WasmSimdF32Butterfly3::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { - WasmSimdF32Butterfly5::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { - WasmSimdF32Butterfly6::new(direction) - } - #[inline(always)] - unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { - WasmSimdF32Butterfly7::new(direction) - } - - #[inline(always)] - unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { - WasmVector::column_butterfly2(rows) - } - // See the f64 impl for why butterflies 3, 5 and 6 rewrap their results and 7 doesn't. - #[inline(always)] - unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { - bf.perform_parallel_fft_direct(rows[0].0, rows[1].0, rows[2].0) - .map(f32::wrap) - } - #[inline(always)] - unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { - WasmVector::column_butterfly4(rows, rotation) - } - #[inline(always)] - unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { - bf.perform_parallel_fft_direct(rows[0].0, rows[1].0, rows[2].0, rows[3].0, rows[4].0) - .map(f32::wrap) - } - #[inline(always)] - unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { - bf.perform_parallel_fft_direct( - rows[0].0, rows[1].0, rows[2].0, rows[3].0, rows[4].0, rows[5].0, - ) - .map(f32::wrap) - } - #[inline(always)] - unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { - bf.perform_parallel_fft_direct(rows) - } - - wasm_simd_radixn_fft_helpers!(); -} - #[cfg(test)] mod unit_tests { - use super::*; - use crate::simd_radixn::test_bodies; + use crate::simd::simd_radixn::test_bodies; + use crate::wasm_simd::wasm_simd_vector::{WasmVector32, WasmVector64}; use wasm_bindgen_test::wasm_bindgen_test; #[wasm_bindgen_test] diff --git a/src/wasm_simd/wasm_simd_vector.rs b/src/wasm_simd/wasm_simd_vector.rs index 11e6e41c..aa098bc4 100644 --- a/src/wasm_simd/wasm_simd_vector.rs +++ b/src/wasm_simd/wasm_simd_vector.rs @@ -6,6 +6,11 @@ use std::ops::{Deref, DerefMut}; use crate::{array_utils::DoubleBuf, twiddles, FftDirection}; +use super::wasm_simd_butterflies::{ + WasmSimdF32Butterfly3, WasmSimdF32Butterfly5, WasmSimdF32Butterfly6, WasmSimdF64Butterfly3, + WasmSimdF64Butterfly5, WasmSimdF64Butterfly6, +}; +use super::wasm_simd_prime_butterflies::{WasmSimdF32Butterfly7, WasmSimdF64Butterfly7}; use super::WasmNum; /// Read these indexes from an WasmSimdArray and build an array of simd vectors. @@ -746,6 +751,250 @@ where } } +// The `SimdVector` impls, which let this backend use the algorithms in `src/simd`. The trait is +// named by path instead of imported, because importing it would make methods like +// `Self::column_butterfly2` ambiguous with the backend's own vector trait. + +// The `SimdVector::fft_helper_*` methods, which are the same forwarding calls for every WASM +// SIMD vector type: they hand the chunk loop to the target-feature-enabled wrappers in +// `wasm_simd_common.rs`. +macro_rules! wasm_simd_vector_fft_helpers { + () => { + #[inline(always)] + unsafe fn fft_helper_immut( + input: &[E], + output: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&[E], &mut [E], &mut [E]), + ) { + super::wasm_simd_common::wasm_simd_fft_helper_immut( + input, + output, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + #[inline(always)] + unsafe fn fft_helper_outofplace( + input: &mut [E], + output: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&mut [E], &mut [E], &mut [E]), + ) { + super::wasm_simd_common::wasm_simd_fft_helper_outofplace( + input, + output, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + #[inline(always)] + unsafe fn fft_helper_inplace( + buffer: &mut [E], + scratch: &mut [E], + chunk_size: usize, + required_scratch: usize, + chunk_fn: impl FnMut(&mut [E], &mut [E]), + ) { + super::wasm_simd_common::wasm_simd_fft_helper_inplace( + buffer, + scratch, + chunk_size, + required_scratch, + chunk_fn, + ) + } + }; +} + +impl crate::simd::simd_vector::SimdVector for WasmVector64 { + const COMPLEX_PER_VECTOR: usize = 1; + + type ScalarType = f64; + type Rotation = Rotation90; + + type Butterfly3 = WasmSimdF64Butterfly3; + type Butterfly5 = WasmSimdF64Butterfly5; + type Butterfly6 = WasmSimdF64Butterfly6; + type Butterfly7 = WasmSimdF64Butterfly7; + + #[inline(always)] + unsafe fn load(data: &[Complex], index: usize) -> Self { + data.load_complex(index) + } + #[inline(always)] + unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { + data.store_complex(value, index) + } + + #[inline(always)] + unsafe fn mul_complex(left: Self, right: Self) -> Self { + WasmVector::mul_complex(left, right) + } + #[inline(always)] + unsafe fn make_mixedradix_twiddle_chunk( + x: usize, + y: usize, + len: usize, + direction: FftDirection, + ) -> Self { + WasmVector::make_mixedradix_twiddle_chunk(x, y, len, direction) + } + + #[inline(always)] + unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { + WasmVector::make_rotate90(direction) + } + #[inline(always)] + unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { + WasmSimdF64Butterfly3::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { + WasmSimdF64Butterfly5::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { + WasmSimdF64Butterfly6::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { + WasmSimdF64Butterfly7::new(direction) + } + + #[inline(always)] + unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { + WasmVector::column_butterfly2(rows) + } + // Butterflies 3, 5 and 6 are written against raw `v128` while `WasmVector64` is a newtype over + // it, so their results come back needing rewrapping. Butterfly 7 already speaks the wrapper + // types and needs none of it. + #[inline(always)] + unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { + bf.perform_fft_direct(rows[0].0, rows[1].0, rows[2].0) + .map(f64::wrap) + } + #[inline(always)] + unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { + WasmVector::column_butterfly4(rows, rotation) + } + #[inline(always)] + unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { + bf.perform_fft_direct(rows[0].0, rows[1].0, rows[2].0, rows[3].0, rows[4].0) + .map(f64::wrap) + } + #[inline(always)] + unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { + bf.perform_fft_direct([ + rows[0].0, rows[1].0, rows[2].0, rows[3].0, rows[4].0, rows[5].0, + ]) + .map(f64::wrap) + } + #[inline(always)] + unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { + bf.perform_fft_direct(rows) + } + + wasm_simd_vector_fft_helpers!(); +} + +impl crate::simd::simd_vector::SimdVector for WasmVector32 { + const COMPLEX_PER_VECTOR: usize = 2; + + type ScalarType = f32; + type Rotation = Rotation90; + + type Butterfly3 = WasmSimdF32Butterfly3; + type Butterfly5 = WasmSimdF32Butterfly5; + type Butterfly6 = WasmSimdF32Butterfly6; + type Butterfly7 = WasmSimdF32Butterfly7; + + #[inline(always)] + unsafe fn load(data: &[Complex], index: usize) -> Self { + data.load_complex(index) + } + #[inline(always)] + unsafe fn store(mut data: &mut [Complex], value: Self, index: usize) { + data.store_complex(value, index) + } + + #[inline(always)] + unsafe fn mul_complex(left: Self, right: Self) -> Self { + WasmVector::mul_complex(left, right) + } + #[inline(always)] + unsafe fn make_mixedradix_twiddle_chunk( + x: usize, + y: usize, + len: usize, + direction: FftDirection, + ) -> Self { + WasmVector::make_mixedradix_twiddle_chunk(x, y, len, direction) + } + + #[inline(always)] + unsafe fn make_rotate90(direction: FftDirection) -> Self::Rotation { + WasmVector::make_rotate90(direction) + } + #[inline(always)] + unsafe fn make_butterfly3(direction: FftDirection) -> Self::Butterfly3 { + WasmSimdF32Butterfly3::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly5(direction: FftDirection) -> Self::Butterfly5 { + WasmSimdF32Butterfly5::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly6(direction: FftDirection) -> Self::Butterfly6 { + WasmSimdF32Butterfly6::new(direction) + } + #[inline(always)] + unsafe fn make_butterfly7(direction: FftDirection) -> Self::Butterfly7 { + WasmSimdF32Butterfly7::new(direction) + } + + #[inline(always)] + unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] { + WasmVector::column_butterfly2(rows) + } + // See the f64 impl for why butterflies 3, 5 and 6 rewrap their results and 7 doesn't. + #[inline(always)] + unsafe fn column_butterfly3(bf: &Self::Butterfly3, rows: [Self; 3]) -> [Self; 3] { + bf.perform_parallel_fft_direct(rows[0].0, rows[1].0, rows[2].0) + .map(f32::wrap) + } + #[inline(always)] + unsafe fn column_butterfly4(rows: [Self; 4], rotation: Self::Rotation) -> [Self; 4] { + WasmVector::column_butterfly4(rows, rotation) + } + #[inline(always)] + unsafe fn column_butterfly5(bf: &Self::Butterfly5, rows: [Self; 5]) -> [Self; 5] { + bf.perform_parallel_fft_direct(rows[0].0, rows[1].0, rows[2].0, rows[3].0, rows[4].0) + .map(f32::wrap) + } + #[inline(always)] + unsafe fn column_butterfly6(bf: &Self::Butterfly6, rows: [Self; 6]) -> [Self; 6] { + bf.perform_parallel_fft_direct( + rows[0].0, rows[1].0, rows[2].0, rows[3].0, rows[4].0, rows[5].0, + ) + .map(f32::wrap) + } + #[inline(always)] + unsafe fn column_butterfly7(bf: &Self::Butterfly7, rows: [Self; 7]) -> [Self; 7] { + bf.perform_parallel_fft_direct(rows) + } + + wasm_simd_vector_fft_helpers!(); +} + #[cfg(test)] mod unit_tests { use super::*; From 453cd7f5b325cc9a746cf62e2a045283bf3ba3e3 Mon Sep 17 00:00:00 2001 From: Henrik Date: Wed, 16 Sep 2026 21:52:48 +0200 Subject: [PATCH 09/22] Move split_cross_len to math_utils It is factoring arithmetic, so it belongs with PrimeFactors rather than on RadixFactor in common.rs. --- src/common.rs | 38 -------------------------------------- src/math_utils.rs | 40 ++++++++++++++++++++++++++++++++++++++++ src/plan.rs | 4 ++-- src/simd/simd_planner.rs | 4 ++-- 4 files changed, 44 insertions(+), 42 deletions(-) diff --git a/src/common.rs b/src/common.rs index d59dc6ea..2dd1db12 100644 --- a/src/common.rs +++ b/src/common.rs @@ -281,42 +281,4 @@ impl RadixFactor { RadixFactor::Factor7 => 7, } } - - /// Split a RadixN cross-FFT length into the layers that make it up, or None if it has a - /// factor no layer can handle. - /// - /// Every planner picks its base by its own rules, but once the base is divided out the rest - /// is the same arithmetic for all of them, so this is deliberately free of any policy. - pub fn split_cross_len(mut cross_len: usize) -> Option> { - let mut factors = Vec::new(); - while cross_len % 7 == 0 { - cross_len /= 7; - factors.push(RadixFactor::Factor7); - } - while cross_len % 6 == 0 { - cross_len /= 6; - factors.push(RadixFactor::Factor6); - } - while cross_len % 5 == 0 { - cross_len /= 5; - factors.push(RadixFactor::Factor5); - } - while cross_len % 3 == 0 { - cross_len /= 3; - factors.push(RadixFactor::Factor3); - } - if !cross_len.is_power_of_two() { - return None; - } - - // benchmarking suggests that we want to add the 4s *last*, i suspect because 4 is a - // better-than-usual value for the transpose - let cross_bits = cross_len.trailing_zeros(); - if cross_bits % 2 == 1 { - factors.push(RadixFactor::Factor2); - } - factors.extend(std::iter::repeat(RadixFactor::Factor4).take(cross_bits as usize / 2)); - - Some(factors.into_boxed_slice()) - } } diff --git a/src/math_utils.rs b/src/math_utils.rs index 164b72f5..92dc22cb 100644 --- a/src/math_utils.rs +++ b/src/math_utils.rs @@ -1,5 +1,7 @@ use num_traits::{One, PrimInt, Zero}; +use crate::common::RadixFactor; + pub fn primitive_root(prime: u64) -> Option { let test_exponents: Vec = distinct_prime_factors(prime - 1) .iter() @@ -507,6 +509,44 @@ impl PartialFactors { } } +/// Split a RadixN cross-FFT length into the layers that make it up, or None if it has a +/// factor no layer can handle. +/// +/// Every planner picks its base by its own rules, but once the base is divided out the rest +/// is the same arithmetic for all of them, so this is deliberately free of any policy. +pub fn split_cross_len(mut cross_len: usize) -> Option> { + let mut factors = Vec::new(); + while cross_len % 7 == 0 { + cross_len /= 7; + factors.push(RadixFactor::Factor7); + } + while cross_len % 6 == 0 { + cross_len /= 6; + factors.push(RadixFactor::Factor6); + } + while cross_len % 5 == 0 { + cross_len /= 5; + factors.push(RadixFactor::Factor5); + } + while cross_len % 3 == 0 { + cross_len /= 3; + factors.push(RadixFactor::Factor3); + } + if !cross_len.is_power_of_two() { + return None; + } + + // benchmarking suggests that we want to add the 4s *last*, i suspect because 4 is a + // better-than-usual value for the transpose + let cross_bits = cross_len.trailing_zeros(); + if cross_bits % 2 == 1 { + factors.push(RadixFactor::Factor2); + } + factors.extend(std::iter::repeat(RadixFactor::Factor4).take(cross_bits as usize / 2)); + + Some(factors.into_boxed_slice()) +} + #[cfg(test)] mod unit_tests { use super::*; diff --git a/src/plan.rs b/src/plan.rs index 56a790d2..e9c9b192 100644 --- a/src/plan.rs +++ b/src/plan.rs @@ -14,7 +14,7 @@ use crate::FftPlannerAvx; use crate::FftPlannerNeon; use crate::FftPlannerSse; -use crate::math_utils::PrimeFactors; +use crate::math_utils::{split_cross_len, PrimeFactors}; enum ChosenFftPlanner { Scalar(FftPlannerScalar), @@ -566,7 +566,7 @@ impl FftPlannerScalar { // we weren't able to use radix4, so fall back to RadixN // theoretically we could do this with the p2, p3, p5 etc values above, but our choice of base knocked them out of sync - let factors = RadixFactor::split_cross_len(cross_len) + let factors = split_cross_len(cross_len) .expect("Every factor RadixN can't handle should have gone into the base"); Arc::new(Recipe::RadixN { factors, base_fft }) diff --git a/src/simd/simd_planner.rs b/src/simd/simd_planner.rs index f404520a..445bf68f 100644 --- a/src/simd/simd_planner.rs +++ b/src/simd/simd_planner.rs @@ -13,7 +13,7 @@ //! style, so folding it in would be a planner change to measure, not a deduplication. use crate::common::RadixFactor; -use crate::math_utils::PrimeFactors; +use crate::math_utils::{split_cross_len, PrimeFactors}; use crate::FftNum; use std::any::TypeId; @@ -171,7 +171,7 @@ pub fn design_radixn(factors: &PrimeFactors, complex_per_vector: usize) -> Optio // Split what's left into cross-FFT layers. Every factor too big for a layer went into the // base above, so the split can't fail, and the same expect guards it in `src/plan.rs`. Some(RadixNPlan::RadixN { - factors: RadixFactor::split_cross_len(cross_len) + factors: split_cross_len(cross_len) .expect("Every factor RadixN can't handle should have gone into the base"), base_len, }) From 4649d7b6803ac1bb4c577644114376a068c4a302 Mon Sep 17 00:00:00 2001 From: Henrik Date: Wed, 16 Sep 2026 21:54:22 +0200 Subject: [PATCH 10/22] Inline the SimdRadixN perform methods, and count the unrolled columns up front The perform methods were one line forwards from the Fft closures. The column loop now computes how many pairs and whether a column is left over before it starts, instead of testing vcol + 2 <= num_vector_columns. --- src/simd/simd_radixn.rs | 66 +++++++++++++++++------------------------ 1 file changed, 28 insertions(+), 38 deletions(-) diff --git a/src/simd/simd_radixn.rs b/src/simd/simd_radixn.rs index fd3f401c..ddfd6e75 100644 --- a/src/simd/simd_radixn.rs +++ b/src/simd/simd_radixn.rs @@ -284,33 +284,6 @@ impl SimdRadixN { layer_twiddles = &layer_twiddles[twiddle_offset..]; } } - - unsafe fn perform_fft_immut( - &self, - input: &[Complex], - output: &mut [Complex], - scratch: &mut [Complex], - ) { - self.transpose(input, output); - self.base_fft.process_with_scratch(output, scratch); - self.cross_ffts(output); - } - - unsafe fn perform_fft_out_of_place( - &self, - input: &mut [Complex], - output: &mut [Complex], - scratch: &mut [Complex], - ) { - self.transpose(input, output); - - // the input is free once the transpose is done, so use it as base scratch when we weren't - // handed any of our own - let base_scratch = if !scratch.is_empty() { scratch } else { input }; - self.base_fft.process_with_scratch(output, base_scratch); - - self.cross_ffts(output); - } } impl Fft for SimdRadixN { @@ -327,7 +300,11 @@ impl Fft for SimdRadixN { scratch, self.len(), self.get_immutable_scratch_len(), - |in_chunk, out_chunk, scratch| self.perform_fft_immut(in_chunk, out_chunk, scratch), + |input, output, scratch| { + self.transpose(input, output); + self.base_fft.process_with_scratch(output, scratch); + self.cross_ffts(output); + }, ); } } @@ -344,8 +321,13 @@ impl Fft for SimdRadixN { scratch, self.len(), self.get_outofplace_scratch_len(), - |in_chunk, out_chunk, scratch| { - self.perform_fft_out_of_place(in_chunk, out_chunk, scratch) + |input, output, scratch| { + self.transpose(input, output); + // the input is free once the transpose is done, so use it as base scratch + // when we weren't handed any of our own + let base_scratch = if !scratch.is_empty() { scratch } else { input }; + self.base_fft.process_with_scratch(output, base_scratch); + self.cross_ffts(output); }, ); } @@ -358,9 +340,17 @@ impl Fft for SimdRadixN { self.len(), self.get_inplace_scratch_len(), |chunk, scratch| { - let (self_scratch, inner_scratch) = scratch.split_at_mut(self.len()); - self.perform_fft_out_of_place(chunk, self_scratch, inner_scratch); - chunk.copy_from_slice(self_scratch); + let (output, inner_scratch) = scratch.split_at_mut(self.len()); + self.transpose(chunk, output); + // same as out of place: the chunk is free once the transpose is done + let base_scratch = if !inner_scratch.is_empty() { + inner_scratch + } else { + &mut *chunk + }; + self.base_fft.process_with_scratch(output, base_scratch); + self.cross_ffts(output); + chunk.copy_from_slice(output); }, ) } @@ -475,8 +465,9 @@ unsafe fn cross_layer( rows }; - let mut vcol = 0; - while vcol + 2 <= num_vector_columns { + let (unroll_count, unroll_remainder) = (num_vector_columns / 2, num_vector_columns % 2); + for i in 0..unroll_count { + let vcol = i * 2; let idx = vcol * complex_per_vector; let a = gather(data, idx, vcol * tw_stride); @@ -489,12 +480,11 @@ unsafe fn cross_layer( V::store(data, *a_row, idx + r * num_columns); V::store(data, *b_row, idx + complex_per_vector + r * num_columns); } - - vcol += 2; } // an odd vector column count leaves one behind - if vcol < num_vector_columns { + if unroll_remainder > 0 { + let vcol = unroll_count * 2; let idx = vcol * complex_per_vector; let a = butterfly(gather(data, idx, vcol * tw_stride)); for (r, a_row) in a.iter().enumerate() { From cf625651574d0db380eff5bc7b44aab349d38f1b Mon Sep 17 00:00:00 2001 From: Henrik Date: Wed, 16 Sep 2026 21:57:51 +0200 Subject: [PATCH 11/22] Gather the cross layer rows with array::from_fn The comment said from_fn was newer than the MSRV, but it has been stable since 1.63 and the MSRV is 1.77. Perf-neutral on NEON: 0.995x to 1.002x over 10 lengths from 120 to 100800, f32 and f64. --- src/simd/simd_radixn.rs | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/simd/simd_radixn.rs b/src/simd/simd_radixn.rs index ddfd6e75..f4e4da85 100644 --- a/src/simd/simd_radixn.rs +++ b/src/simd/simd_radixn.rs @@ -455,14 +455,14 @@ unsafe fn cross_layer( // The row-0 twiddle is always 1, so it's neither stored nor applied. let gather = |data: &[Complex], idx: usize, tw_base: usize| -> [V; RADIX] { - // row 0 first, so the array is fully initialized without `array::from_fn`, which is - // newer than the crate MSRV - let mut rows = [V::load(data, idx); RADIX]; - for (r, row) in rows.iter_mut().enumerate().skip(1) { + std::array::from_fn(|r| { let v = V::load(data, idx + r * num_columns); - *row = V::mul_complex(v, *twiddles.get_unchecked(tw_base + r - 1)); - } - rows + if r == 0 { + v + } else { + V::mul_complex(v, *twiddles.get_unchecked(tw_base + r - 1)) + } + }) }; let (unroll_count, unroll_remainder) = (num_vector_columns / 2, num_vector_columns % 2); From ba8d179bd8a947dc4e4f9ad631f9b3b95b19650d Mon Sep 17 00:00:00 2001 From: Henrik Date: Wed, 16 Sep 2026 22:38:32 +0200 Subject: [PATCH 12/22] Bring over the planner tuning harness from the counted cost spike The harness, the tuning feature and the op count derivation. The lab notes, the weight grid scripts and the one-off analysis scripts stay on counted_cost_spike. --- Cargo.toml | 1 + src/lib.rs | 5 + src/neon/neon_planner.rs | 4 +- src/plan.rs | 4 +- src/sse/sse_planner.rs | 4 +- src/tuning/adapters.rs | 473 +++++++++ src/tuning/mod.rs | 636 +++++++++++ src/wasm_simd/wasm_simd_planner.rs | 4 +- tools/planner_tuning/.gitignore | 8 + tools/planner_tuning/COST-MODEL.md | 375 +++++++ tools/planner_tuning/Cargo.toml | 26 + tools/planner_tuning/OP-COUNTS.md | 165 +++ tools/planner_tuning/README.md | 148 +++ tools/planner_tuning/plot_sweep.py | 152 +++ tools/planner_tuning/run_wasm.mjs | 14 + tools/planner_tuning/src/counted.rs | 564 ++++++++++ tools/planner_tuning/src/emit.rs | 77 ++ tools/planner_tuning/src/main.rs | 1509 +++++++++++++++++++++++++++ tools/planner_tuning/src/model.rs | 164 +++ 19 files changed, 4325 insertions(+), 8 deletions(-) create mode 100644 src/tuning/adapters.rs create mode 100644 src/tuning/mod.rs create mode 100644 tools/planner_tuning/.gitignore create mode 100644 tools/planner_tuning/COST-MODEL.md create mode 100644 tools/planner_tuning/Cargo.toml create mode 100644 tools/planner_tuning/OP-COUNTS.md create mode 100644 tools/planner_tuning/README.md create mode 100755 tools/planner_tuning/plot_sweep.py create mode 100644 tools/planner_tuning/run_wasm.mjs create mode 100644 tools/planner_tuning/src/counted.rs create mode 100644 tools/planner_tuning/src/emit.rs create mode 100644 tools/planner_tuning/src/main.rs create mode 100644 tools/planner_tuning/src/model.rs diff --git a/Cargo.toml b/Cargo.toml index d0132078..39e3b246 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,6 +31,7 @@ avx = [] sse = [] neon = [] wasm_simd = [] +tuning = [] [dependencies] diff --git a/src/lib.rs b/src/lib.rs index 10bc9cec..719b4c2c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -611,5 +611,10 @@ mod wasm_simd { pub use self::wasm_simd::wasm_simd_planner::FftPlannerWasmSimd; +// Internal support for the planner-tuning tools. Not part of the public API. +#[cfg(feature = "tuning")] +#[doc(hidden)] +pub mod tuning; + #[cfg(test)] mod test_utils; diff --git a/src/neon/neon_planner.rs b/src/neon/neon_planner.rs index 2328805f..bea1e657 100644 --- a/src/neon/neon_planner.rs +++ b/src/neon/neon_planner.rs @@ -239,7 +239,7 @@ impl FftPlannerNeon { } // Make a recipe for a length - fn design_fft_for_len(&mut self, len: usize) -> Arc { + pub(crate) fn design_fft_for_len(&mut self, len: usize) -> Arc { if len < 1 { Arc::new(Recipe::Dft(len)) } else if let Some(recipe) = self.recipe_cache.get(&len) { @@ -253,7 +253,7 @@ impl FftPlannerNeon { } // Create the fft from a recipe, take from cache if possible - fn build_fft(&mut self, recipe: &Recipe, direction: FftDirection) -> Arc> { + pub(crate) fn build_fft(&mut self, recipe: &Recipe, direction: FftDirection) -> Arc> { let len = recipe.len(); if let Some(instance) = self.algorithm_cache.get(len, direction) { instance diff --git a/src/plan.rs b/src/plan.rs index e9c9b192..7055060b 100644 --- a/src/plan.rs +++ b/src/plan.rs @@ -309,7 +309,7 @@ impl FftPlannerScalar { } // Make a recipe for a length - fn design_fft_for_len(&mut self, len: usize) -> Arc { + pub(crate) fn design_fft_for_len(&mut self, len: usize) -> Arc { if len < 2 { Arc::new(Recipe::Dft(len)) } else if let Some(recipe) = self.recipe_cache.get(&len) { @@ -323,7 +323,7 @@ impl FftPlannerScalar { } // Create the fft from a recipe, take from cache if possible - fn build_fft(&mut self, recipe: &Recipe, direction: FftDirection) -> Arc> { + pub(crate) fn build_fft(&mut self, recipe: &Recipe, direction: FftDirection) -> Arc> { let len = recipe.len(); if let Some(instance) = self.algorithm_cache.get(len, direction) { instance diff --git a/src/sse/sse_planner.rs b/src/sse/sse_planner.rs index bfd3721a..34e821cd 100644 --- a/src/sse/sse_planner.rs +++ b/src/sse/sse_planner.rs @@ -239,7 +239,7 @@ impl FftPlannerSse { } // Make a recipe for a length - fn design_fft_for_len(&mut self, len: usize) -> Arc { + pub(crate) fn design_fft_for_len(&mut self, len: usize) -> Arc { if len < 1 { Arc::new(Recipe::Dft(len)) } else if let Some(recipe) = self.recipe_cache.get(&len) { @@ -253,7 +253,7 @@ impl FftPlannerSse { } // Create the fft from a recipe, take from cache if possible - fn build_fft(&mut self, recipe: &Recipe, direction: FftDirection) -> Arc> { + pub(crate) fn build_fft(&mut self, recipe: &Recipe, direction: FftDirection) -> Arc> { let len = recipe.len(); if let Some(instance) = self.algorithm_cache.get(len, direction) { instance diff --git a/src/tuning/adapters.rs b/src/tuning/adapters.rs new file mode 100644 index 00000000..d9e01aba --- /dev/null +++ b/src/tuning/adapters.rs @@ -0,0 +1,473 @@ +//! Per-planner adapters between [`Spec`] and each planner's own `Recipe`. +//! +//! The three SIMD planners have structurally identical `Recipe` types, so their adapters are +//! generated by a macro. The scalar one is written out because it differs: it has `RadixN`, and +//! a different set of butterflies with no separate prime-butterfly kind. + +use std::any::TypeId; +use std::sync::Arc; + +use super::{Spec, TunablePlanner}; +use crate::{Fft, FftDirection, FftNum}; + +/// Panic message used when a spec asks for something a planner cannot express. +fn unsupported(planner: &str, what: &str) -> ! { + panic!("the {} planner cannot express {}", planner, what) +} + +// --------------------------------------------------------------------------- +// Scalar +// --------------------------------------------------------------------------- + +/// Butterfly lengths the scalar planner has kernels for. +const SCALAR_BUTTERFLIES: [usize; 20] = [ + 2, 3, 4, 5, 6, 7, 8, 9, 11, 12, 13, 16, 17, 19, 23, 24, 27, 29, 31, 32, +]; + +/// Bases worth trying under Radix4 and RadixN in the scalar planner. These are the sizes its own +/// `design_radixn` chooses between, which are the ones that have been benchmarked upstream. +const SCALAR_RADIX_BASES: [usize; 9] = [5, 6, 7, 8, 9, 12, 16, 24, 27]; + +pub struct ScalarTuner { + planner: crate::FftPlannerScalar, +} + +impl ScalarTuner { + fn to_spec(recipe: &crate::plan::Recipe) -> Arc { + use crate::plan::Recipe; + Arc::new(match recipe { + Recipe::Dft(len) => Spec::Dft(*len), + Recipe::Radix4 { k, base_fft } => Spec::Radix4 { + k: *k, + base: Self::to_spec(base_fft), + }, + Recipe::RadixN { factors, base_fft } => Spec::RadixN { + radixes: factors.iter().map(|f| f.radix()).collect(), + base: Self::to_spec(base_fft), + }, + Recipe::MixedRadix { + left_fft, + right_fft, + } => Spec::MixedRadix { + left: Self::to_spec(left_fft), + right: Self::to_spec(right_fft), + small: false, + }, + Recipe::MixedRadixSmall { + left_fft, + right_fft, + } => Spec::MixedRadix { + left: Self::to_spec(left_fft), + right: Self::to_spec(right_fft), + small: true, + }, + Recipe::GoodThomasAlgorithm { + left_fft, + right_fft, + } => Spec::GoodThomas { + left: Self::to_spec(left_fft), + right: Self::to_spec(right_fft), + small: false, + }, + Recipe::GoodThomasAlgorithmSmall { + left_fft, + right_fft, + } => Spec::GoodThomas { + left: Self::to_spec(left_fft), + right: Self::to_spec(right_fft), + small: true, + }, + Recipe::RadersAlgorithm { inner_fft } => Spec::Raders { + inner: Self::to_spec(inner_fft), + }, + Recipe::BluesteinsAlgorithm { len, inner_fft } => Spec::Bluesteins { + len: *len, + inner: Self::to_spec(inner_fft), + }, + butterfly => Spec::Butterfly(butterfly.len()), + }) + } + + fn from_spec(spec: &Spec) -> Arc { + use crate::common::RadixFactor; + use crate::plan::Recipe; + Arc::new(match spec { + Spec::Dft(len) => Recipe::Dft(*len), + Spec::Radix4 { k, base } => Recipe::Radix4 { + k: *k, + base_fft: Self::from_spec(base), + }, + Spec::RadixN { radixes, base } => Recipe::RadixN { + factors: radixes + .iter() + .map(|r| match r { + 2 => RadixFactor::Factor2, + 3 => RadixFactor::Factor3, + 4 => RadixFactor::Factor4, + 5 => RadixFactor::Factor5, + 6 => RadixFactor::Factor6, + 7 => RadixFactor::Factor7, + other => unsupported("scalar", &format!("a radix of {}", other)), + }) + .collect::>() + .into_boxed_slice(), + base_fft: Self::from_spec(base), + }, + Spec::MixedRadix { left, right, small } => { + let (left_fft, right_fft) = (Self::from_spec(left), Self::from_spec(right)); + if *small { + Recipe::MixedRadixSmall { + left_fft, + right_fft, + } + } else { + Recipe::MixedRadix { + left_fft, + right_fft, + } + } + } + Spec::GoodThomas { left, right, small } => { + let (left_fft, right_fft) = (Self::from_spec(left), Self::from_spec(right)); + if *small { + Recipe::GoodThomasAlgorithmSmall { + left_fft, + right_fft, + } + } else { + Recipe::GoodThomasAlgorithm { + left_fft, + right_fft, + } + } + } + Spec::Raders { inner } => Recipe::RadersAlgorithm { + inner_fft: Self::from_spec(inner), + }, + Spec::Bluesteins { len, inner } => Recipe::BluesteinsAlgorithm { + len: *len, + inner_fft: Self::from_spec(inner), + }, + Spec::Butterfly(len) => match len { + 2 => Recipe::Butterfly2, + 3 => Recipe::Butterfly3, + 4 => Recipe::Butterfly4, + 5 => Recipe::Butterfly5, + 6 => Recipe::Butterfly6, + 7 => Recipe::Butterfly7, + 8 => Recipe::Butterfly8, + 9 => Recipe::Butterfly9, + 11 => Recipe::Butterfly11, + 12 => Recipe::Butterfly12, + 13 => Recipe::Butterfly13, + 16 => Recipe::Butterfly16, + 17 => Recipe::Butterfly17, + 19 => Recipe::Butterfly19, + 23 => Recipe::Butterfly23, + 24 => Recipe::Butterfly24, + 27 => Recipe::Butterfly27, + 29 => Recipe::Butterfly29, + 31 => Recipe::Butterfly31, + 32 => Recipe::Butterfly32, + other => unsupported("scalar", &format!("a butterfly of {}", other)), + }, + }) + } +} + +impl TunablePlanner for ScalarTuner { + fn label() -> &'static str { + "scalar" + } + + fn new() -> Self { + Self { + planner: crate::FftPlannerScalar::new(), + } + } + + fn plan(&mut self, len: usize) -> Arc { + let recipe = self.planner.design_fft_for_len(len); + Self::to_spec(&recipe) + } + + fn build(&mut self, spec: &Spec, direction: FftDirection) -> Arc> { + super::check_unambiguous(spec, &mut Default::default()).expect("ambiguous recipe"); + let recipe = Self::from_spec(spec); + let mut planner = crate::FftPlannerScalar::::new(); + planner.build_fft(&recipe, direction) + } + + fn butterfly_lens() -> Vec { + SCALAR_BUTTERFLIES.to_vec() + } + + fn radix4_bases() -> Vec { + SCALAR_RADIX_BASES.to_vec() + } + + fn has_radixn() -> bool { + true + } +} + +// --------------------------------------------------------------------------- +// SIMD planners +// --------------------------------------------------------------------------- + +/// Butterfly lengths the SIMD planners have dedicated kernels for, excluding their prime +/// butterflies, which are listed separately by each module. +const SIMD_BUTTERFLIES: [usize; 14] = [1, 2, 3, 4, 5, 6, 8, 9, 10, 12, 15, 16, 24, 32]; + +/// Bases a SIMD `Radix4` can be built on. The smallest are excluded per element type at runtime, +/// because the kernels need a whole number of vector pairs. +const SIMD_RADIX4_BASES: [usize; 10] = [1, 2, 4, 8, 16, 32, 3, 6, 12, 24]; + +/// Bases a SIMD `RadixN` can be built on. Wider than the `Radix4` set because `SimdRadixN` needs +/// only a whole number of vectors in the base, not a whole number of vector pairs, so odd bases +/// are legal for f64. These are the sizes `simd_planner::design_radixn` chooses between, plus the +/// neighbouring butterfly lengths, so the sweep can see one factor either side of its pick. +const SIMD_RADIXN_BASES: [usize; 12] = [4, 5, 6, 7, 8, 9, 10, 12, 15, 16, 24, 32]; + +/// Generates a `TunablePlanner` adapter for one of the SIMD planners. +/// +/// They are identical apart from which module their `Recipe`, planner and prime butterflies live +/// in, so this exists to keep three copies of the same conversion from drifting apart. +macro_rules! simd_adapter { + ($tuner:ident, $label:literal, $planner:path, $recipe:path, $prime_lens:path) => { + pub struct $tuner { + planner: $planner, + } + + impl $tuner { + fn to_spec(recipe: &$recipe) -> Arc { + use $recipe as R; + Arc::new(match recipe { + R::Dft(len) => Spec::Dft(*len), + R::Radix4 { k, base_fft } => Spec::Radix4 { + k: *k, + base: Self::to_spec(base_fft), + }, + R::RadixN { factors, base_fft } => Spec::RadixN { + radixes: factors.iter().map(|f| f.radix()).collect(), + base: Self::to_spec(base_fft), + }, + R::MixedRadix { + left_fft, + right_fft, + } => Spec::MixedRadix { + left: Self::to_spec(left_fft), + right: Self::to_spec(right_fft), + small: false, + }, + R::MixedRadixSmall { + left_fft, + right_fft, + } => Spec::MixedRadix { + left: Self::to_spec(left_fft), + right: Self::to_spec(right_fft), + small: true, + }, + R::GoodThomasAlgorithm { + left_fft, + right_fft, + } => Spec::GoodThomas { + left: Self::to_spec(left_fft), + right: Self::to_spec(right_fft), + small: false, + }, + R::GoodThomasAlgorithmSmall { + left_fft, + right_fft, + } => Spec::GoodThomas { + left: Self::to_spec(left_fft), + right: Self::to_spec(right_fft), + small: true, + }, + R::RadersAlgorithm { inner_fft } => Spec::Raders { + inner: Self::to_spec(inner_fft), + }, + R::BluesteinsAlgorithm { len, inner_fft } => Spec::Bluesteins { + len: *len, + inner: Self::to_spec(inner_fft), + }, + butterfly => Spec::Butterfly(butterfly.len()), + }) + } + + fn from_spec(spec: &Spec) -> Arc<$recipe> { + use $recipe as R; + Arc::new(match spec { + Spec::Dft(len) => R::Dft(*len), + Spec::Radix4 { k, base } => R::Radix4 { + k: *k, + base_fft: Self::from_spec(base), + }, + Spec::RadixN { radixes, base } => R::RadixN { + factors: radixes + .iter() + .map(|r| match r { + 2 => crate::common::RadixFactor::Factor2, + 3 => crate::common::RadixFactor::Factor3, + 4 => crate::common::RadixFactor::Factor4, + 5 => crate::common::RadixFactor::Factor5, + 6 => crate::common::RadixFactor::Factor6, + 7 => crate::common::RadixFactor::Factor7, + other => { + unsupported($label, &format!("a radix of {}", other)) + } + }) + .collect::>() + .into_boxed_slice(), + base_fft: Self::from_spec(base), + }, + Spec::MixedRadix { left, right, small } => { + let (left_fft, right_fft) = (Self::from_spec(left), Self::from_spec(right)); + if *small { + R::MixedRadixSmall { + left_fft, + right_fft, + } + } else { + R::MixedRadix { + left_fft, + right_fft, + } + } + } + Spec::GoodThomas { left, right, small } => { + let (left_fft, right_fft) = (Self::from_spec(left), Self::from_spec(right)); + if *small { + R::GoodThomasAlgorithmSmall { + left_fft, + right_fft, + } + } else { + R::GoodThomasAlgorithm { + left_fft, + right_fft, + } + } + } + Spec::Raders { inner } => R::RadersAlgorithm { + inner_fft: Self::from_spec(inner), + }, + Spec::Bluesteins { len, inner } => R::BluesteinsAlgorithm { + len: *len, + inner_fft: Self::from_spec(inner), + }, + Spec::Butterfly(len) => match len { + 1 => R::Butterfly1, + 2 => R::Butterfly2, + 3 => R::Butterfly3, + 4 => R::Butterfly4, + 5 => R::Butterfly5, + 6 => R::Butterfly6, + 8 => R::Butterfly8, + 9 => R::Butterfly9, + 10 => R::Butterfly10, + 12 => R::Butterfly12, + 15 => R::Butterfly15, + 16 => R::Butterfly16, + 24 => R::Butterfly24, + 32 => R::Butterfly32, + len if $prime_lens().contains(len) => R::PrimeButterfly { len: *len }, + other => unsupported($label, &format!("a butterfly of {}", other)), + }, + }) + } + } + + impl TunablePlanner for $tuner { + fn label() -> &'static str { + $label + } + + fn new() -> Self { + Self { + planner: <$planner>::new() + .expect(concat!("this machine does not support ", $label)), + } + } + + fn plan(&mut self, len: usize) -> Arc { + let recipe = self.planner.design_fft_for_len(len); + Self::to_spec(&recipe) + } + + fn build(&mut self, spec: &Spec, direction: FftDirection) -> Arc> { + super::check_unambiguous(spec, &mut Default::default()) + .expect("ambiguous recipe"); + let recipe = Self::from_spec(spec); + let mut planner = <$planner>::new() + .expect(concat!("this machine does not support ", $label)); + planner.build_fft(&recipe, direction) + } + + fn butterfly_lens() -> Vec { + let mut lens: Vec = SIMD_BUTTERFLIES.to_vec(); + lens.extend_from_slice($prime_lens()); + lens.sort_unstable(); + lens.dedup(); + lens + } + + fn radix4_bases() -> Vec { + // The kernels require the base length to be a whole number of vector pairs, + // which rules out the smallest bases, and rules out more of them for f32. + let multiple = if TypeId::of::() == TypeId::of::() { + 4 + } else { + 2 + }; + SIMD_RADIX4_BASES + .iter() + .copied() + .filter(|base| base % multiple == 0) + .collect() + } + + fn radixn_bases() -> Vec { + // `SimdRadixN` only asserts a whole number of vectors in the base, so for f64 + // every base is legal and for f32 the odd ones are not. + let multiple = crate::simd::simd_planner::complex_per_vector::(); + SIMD_RADIXN_BASES + .iter() + .copied() + .filter(|base| base % multiple == 0) + .collect() + } + + fn has_radixn() -> bool { + true + } + } + }; +} + +#[cfg(all(target_arch = "aarch64", feature = "neon"))] +simd_adapter!( + NeonTuner, + "neon", + crate::FftPlannerNeon, + crate::neon::neon_planner::Recipe, + crate::neon::neon_prime_butterflies::prime_butterfly_lens +); + +#[cfg(all(target_arch = "x86_64", feature = "sse"))] +simd_adapter!( + SseTuner, + "sse", + crate::FftPlannerSse, + crate::sse::sse_planner::Recipe, + crate::sse::sse_prime_butterflies::prime_butterfly_lens +); + +#[cfg(all(target_arch = "wasm32", feature = "wasm_simd"))] +simd_adapter!( + WasmSimdTuner, + "wasm_simd", + crate::FftPlannerWasmSimd, + crate::wasm_simd::wasm_simd_planner::Recipe, + crate::wasm_simd::wasm_simd_prime_butterflies::prime_butterfly_lens +); diff --git a/src/tuning/mod.rs b/src/tuning/mod.rs new file mode 100644 index 00000000..4acc1164 --- /dev/null +++ b/src/tuning/mod.rs @@ -0,0 +1,636 @@ +//! Support code for tuning the FFT planners. +//! +//! Compiled only when the non-default `tuning` feature is enabled, and not part of the public +//! API. It exists so that measurement tools can build, name, and enumerate the recipes a planner +//! works with, rather than reimplementing construction and risking measuring something no +//! planner would ever build. +//! +//! Every planner has its own `Recipe` type, and the three SIMD ones are structurally identical +//! while the scalar one additionally has `RadixN`. Rather than duplicate the tooling per +//! planner, everything here is written against [`Spec`], a planner-independent description, and +//! each planner supplies a thin adapter implementing [`TunablePlanner`]. + +use std::collections::HashMap; +use std::sync::Arc; + +use num_integer::gcd; + +use crate::{Fft, FftDirection, FftNum}; + +mod adapters; +pub use adapters::*; + +// --------------------------------------------------------------------------- +// Planner-independent recipe description +// --------------------------------------------------------------------------- + +/// A recipe, described independently of which planner will build it. +#[derive(Debug, Clone, PartialEq)] +pub enum Spec { + Dft(usize), + /// A dedicated kernel for this length. Which concrete butterfly that is, including whether + /// it is one of the SIMD planners' prime butterflies, is the adapter's business. + Butterfly(usize), + Radix4 { + k: u32, + base: Arc, + }, + /// Only the scalar planner has this. + RadixN { + radixes: Vec, + base: Arc, + }, + MixedRadix { + left: Arc, + right: Arc, + small: bool, + }, + GoodThomas { + left: Arc, + right: Arc, + small: bool, + }, + Raders { + inner: Arc, + }, + Bluesteins { + len: usize, + inner: Arc, + }, +} + +impl Spec { + pub fn len(&self) -> usize { + match self { + Spec::Dft(len) | Spec::Butterfly(len) => *len, + Spec::Radix4 { k, base } => base.len() << (2 * k), + Spec::RadixN { radixes, base } => base.len() * radixes.iter().product::(), + Spec::MixedRadix { left, right, .. } | Spec::GoodThomas { left, right, .. } => { + left.len() * right.len() + } + Spec::Raders { inner } => inner.len() + 1, + Spec::Bluesteins { len, .. } => *len, + } + } + + /// A short tag naming the algorithm, used to attribute measured overhead. + pub fn kind(&self) -> &'static str { + match self { + Spec::Dft(_) => "dft", + Spec::Butterfly(_) => "butterfly", + Spec::Radix4 { .. } => "r4", + Spec::RadixN { .. } => "rn", + Spec::MixedRadix { small: false, .. } => "mr", + Spec::MixedRadix { small: true, .. } => "mrs", + Spec::GoodThomas { small: false, .. } => "gt", + Spec::GoodThomas { small: true, .. } => "gts", + Spec::Raders { .. } => "rad", + Spec::Bluesteins { .. } => "bs", + } + } + + pub fn children(&self) -> Vec<&Arc> { + match self { + Spec::Radix4 { base, .. } | Spec::RadixN { base, .. } => vec![base], + Spec::Raders { inner } | Spec::Bluesteins { inner, .. } => vec![inner], + Spec::MixedRadix { left, right, .. } | Spec::GoodThomas { left, right, .. } => { + vec![left, right] + } + _ => Vec::new(), + } + } +} + +// --------------------------------------------------------------------------- +// Naming +// --------------------------------------------------------------------------- + +/// Render a spec in the syntax accepted by [`parse`]. +pub fn to_spec_string(spec: &Spec) -> String { + match spec { + Spec::Dft(len) => format!("dft({})", len), + Spec::Butterfly(len) => format!("b{}", len), + Spec::Radix4 { k, base } => format!("r4({},{})", k, to_spec_string(base)), + Spec::RadixN { radixes, base } => format!( + "rn({},{})", + radixes + .iter() + .map(|r| r.to_string()) + .collect::>() + .join("."), + to_spec_string(base) + ), + Spec::MixedRadix { left, right, small } => format!( + "{}({},{})", + if *small { "mrs" } else { "mr" }, + to_spec_string(left), + to_spec_string(right) + ), + Spec::GoodThomas { left, right, small } => format!( + "{}({},{})", + if *small { "gts" } else { "gt" }, + to_spec_string(left), + to_spec_string(right) + ), + Spec::Raders { inner } => format!("rad({})", to_spec_string(inner)), + Spec::Bluesteins { len, inner } => format!("bs({},{})", len, to_spec_string(inner)), + } +} + +struct Parser<'a> { + s: &'a str, + pos: usize, +} + +impl<'a> Parser<'a> { + fn peek(&self) -> Option { + self.s[self.pos..].chars().next() + } + + fn eat(&mut self, c: char) -> Result<(), String> { + match self.peek() { + Some(got) if got == c => { + self.pos += got.len_utf8(); + Ok(()) + } + other => Err(format!( + "expected '{}' at offset {}, found {:?}", + c, self.pos, other + )), + } + } + + fn ident(&mut self) -> String { + let start = self.pos; + while matches!(self.peek(), Some(c) if c.is_ascii_alphanumeric()) { + self.pos += 1; + } + self.s[start..self.pos].to_string() + } + + fn number(&mut self) -> Result { + let start = self.pos; + while matches!(self.peek(), Some(c) if c.is_ascii_digit()) { + self.pos += 1; + } + self.s[start..self.pos] + .parse() + .map_err(|_| format!("expected a number at offset {}", start)) + } + + fn spec(&mut self) -> Result, String> { + let name = self.ident(); + if name.is_empty() { + return Err(format!("expected a recipe at offset {}", self.pos)); + } + if let Some(size) = name.strip_prefix('b') { + if let Ok(size) = size.parse::() { + return Ok(Arc::new(Spec::Butterfly(size))); + } + } + + self.eat('(')?; + let result = match name.as_str() { + "dft" => Arc::new(Spec::Dft(self.number()?)), + "r4" => { + let k = self.number()? as u32; + self.eat(',')?; + let base = self.spec()?; + Arc::new(Spec::Radix4 { k, base }) + } + "rn" => { + let mut radixes = vec![self.number()?]; + while self.peek() == Some('.') { + self.eat('.')?; + radixes.push(self.number()?); + } + self.eat(',')?; + let base = self.spec()?; + Arc::new(Spec::RadixN { radixes, base }) + } + "rad" => Arc::new(Spec::Raders { + inner: self.spec()?, + }), + "bs" => { + let len = self.number()?; + self.eat(',')?; + let inner = self.spec()?; + Arc::new(Spec::Bluesteins { len, inner }) + } + "mr" | "mrs" | "gt" | "gts" => { + let left = self.spec()?; + self.eat(',')?; + let right = self.spec()?; + let small = name.ends_with('s'); + Arc::new(if name.starts_with("mr") { + Spec::MixedRadix { left, right, small } + } else { + Spec::GoodThomas { left, right, small } + }) + } + other => return Err(format!("unknown recipe '{}'", other)), + }; + self.eat(')')?; + Ok(result) + } +} + +/// Parse a spec string, eg `mr(r4(2,b16),b12)`. +pub fn parse(spec: &str) -> Result, String> { + let cleaned: String = spec.chars().filter(|c| !c.is_whitespace()).collect(); + let mut parser = Parser { + s: &cleaned, + pos: 0, + }; + let parsed = parser.spec()?; + if parser.pos != cleaned.len() { + return Err(format!("trailing junk at offset {}", parser.pos)); + } + Ok(parsed) +} + +// --------------------------------------------------------------------------- +// The planner adapter +// --------------------------------------------------------------------------- + +/// A planner that measurement tools can drive. +pub trait TunablePlanner: Sized { + /// Name used to select this planner on the command line. + fn label() -> &'static str; + + fn new() -> Self; + + /// The recipe this planner picks for `len`, as a [`Spec`]. + fn plan(&mut self, len: usize) -> Arc; + + /// Build an FFT for an arbitrary spec. + /// + /// Each call uses a planner of its own, so that nothing is shared between the recipes being + /// compared. Panics if the spec is not something this planner can express. + fn build(&mut self, spec: &Spec, direction: FftDirection) -> Arc>; + + /// Lengths this planner has a dedicated kernel for. + fn butterfly_lens() -> Vec; + + /// Base lengths a `Radix4` can be built on for element type `T`. + fn radix4_bases() -> Vec; + + /// Base lengths a `RadixN` can be built on for element type `T`. + /// + /// Separate from [`radix4_bases`] because the constraint is looser: a `Radix4` kernel needs a + /// whole number of vector *pairs* in the base, while `SimdRadixN` needs only a whole number of + /// vectors. Defaults to the `Radix4` set for planners where the two coincide. + fn radixn_bases() -> Vec { + Self::radix4_bases() + } + + /// Whether this planner can express `RadixN`. + fn has_radixn() -> bool { + false + } +} + +// --------------------------------------------------------------------------- +// Candidate enumeration +// --------------------------------------------------------------------------- + +/// Walks a spec and fails if one length appears with two different recipes. +/// +/// Planner algorithm caches are keyed on length alone, so a tree containing two different +/// recipes of the same length would silently build the same FFT twice and quietly invalidate +/// whatever was measured. Every spec goes through this before it is built. +pub fn check_unambiguous(spec: &Spec, seen: &mut HashMap) -> Result<(), String> { + let rendered = to_spec_string(spec); + if let Some(previous) = seen.insert(spec.len(), rendered.clone()) { + if previous != rendered { + return Err(format!( + "length {} appears as both '{}' and '{}'; the algorithm cache cannot hold both", + spec.len(), + previous, + rendered + )); + } + } + for child in spec.children() { + check_unambiguous(child, seen)?; + } + Ok(()) +} + +/// Plausible alternatives to the planner's choice for `len`, the planner's own pick first. +/// +/// Deliberately broader than what any planner would consider: the point is to find what the best +/// available recipe actually is, so a planner's pick can be scored against it. +pub fn candidates>(planner: &mut P, len: usize) -> Vec> { + candidates_inner(planner, len, false) +} + +/// `candidates`, optionally skipping the wider-of-the-two-first ordering of each split. +/// +/// The prune has to happen here rather than as a filter afterwards. Generating a candidate costs +/// far more than pricing one: it renders the spec to a string, scans the seen-list linearly, and +/// walks the tree to check it is unambiguous. Filtering after the fact at length 1200 cut the +/// candidate count from 48 to 32 but plan time only from 211us to 181us; skipping the work up +/// front is what actually saves it. +fn candidates_inner>( + planner: &mut P, + len: usize, + planning: bool, +) -> Vec> { + let mut out: Vec> = vec![planner.plan(len)]; + let mut seen: Vec = vec![to_spec_string(&out[0])]; + + let mut push = |spec: Arc, out: &mut Vec>, seen: &mut Vec| { + if spec.len() != len { + return; + } + let rendered = to_spec_string(&spec); + if !seen.contains(&rendered) && check_unambiguous(&spec, &mut HashMap::new()).is_ok() { + seen.push(rendered); + out.push(spec); + } + }; + + let butterflies = P::butterfly_lens(); + + // Every two-way split, in both orders, as each algorithm that can express it. + for left_len in 2..=(len / 2) { + if len % left_len != 0 { + continue; + } + let right_len = len / left_len; + if planning && left_len > right_len { + continue; + } + let left = planner.plan(left_len); + let right = planner.plan(right_len); + let coprime = gcd(left_len, right_len) == 1; + let small = left_len < 33 && right_len < 33; + + // Both orderings, unless pruning. The loop above already restricts `left_len` to the + // smaller half when pruning, so the surviving order is the smaller-width-first one. + let orders: Vec<(Arc, Arc)> = if planning { + vec![(Arc::clone(&left), Arc::clone(&right))] + } else { + vec![ + (Arc::clone(&left), Arc::clone(&right)), + (Arc::clone(&right), Arc::clone(&left)), + ] + }; + for (left, right) in orders { + for small_flag in if small { + vec![false, true] + } else { + vec![false] + } { + push( + Arc::new(Spec::MixedRadix { + left: Arc::clone(&left), + right: Arc::clone(&right), + small: small_flag, + }), + &mut out, + &mut seen, + ); + if coprime { + push( + Arc::new(Spec::GoodThomas { + left: Arc::clone(&left), + right: Arc::clone(&right), + small: small_flag, + }), + &mut out, + &mut seen, + ); + } + } + } + } + + // Radix4 on every base that divides out to a power of four. + for base in P::radix4_bases() { + if base == 0 || len % base != 0 { + continue; + } + let cross = len / base; + if !cross.is_power_of_two() || cross.trailing_zeros() % 2 != 0 { + continue; + } + push( + Arc::new(Spec::Radix4 { + k: cross.trailing_zeros() / 2, + base: planner.plan(base), + }), + &mut out, + &mut seen, + ); + } + + // RadixN, where available, over the factorisations the algorithm supports. + if P::has_radixn() { + for base in P::radixn_bases() { + if base == 0 || len % base != 0 { + continue; + } + let mut cross = len / base; + if cross <= 1 { + continue; + } + let mut radixes = Vec::new(); + for radix in [7usize, 6, 5, 4, 3, 2] { + while cross % radix == 0 { + cross /= radix; + radixes.push(radix); + } + } + if cross == 1 && !radixes.is_empty() { + // Benchmarking upstream suggests the 4s want to go last. + radixes.sort_by_key(|r| if *r == 4 { 1 } else { 0 }); + push( + Arc::new(Spec::RadixN { + radixes, + base: planner.plan(base), + }), + &mut out, + &mut seen, + ); + } + } + } + + // Prime lengths: Rader's. It needs len - 1 to factor, so it is genuinely prime-only. + if len > 3 && crate::math_utils::PrimeFactors::compute(len).is_prime() { + push( + Arc::new(Spec::Raders { + inner: planner.plan(len - 1), + }), + &mut out, + &mut seen, + ); + } + + // Bluestein's over a range of inner lengths, at **every** length and not only at primes. + // + // Nothing about Bluestein's needs a prime: the only requirement is an inner FFT of at least + // 2*len - 1. The shipping planner reaches it solely from `design_prime`, so a composite can + // never be given it, and that is a real loss rather than a theoretical one. Length 671 is + // 11 x 61; its only candidates were a MixedRadix or GoodThomas wrapped around a Rader's for + // 61, and it measured slower than both of its prime neighbours 673 and 677, which do get + // Bluestein's. Built by hand, bs(671, r4(3,b24)) is 1.46x faster than the planner's pick. + // + // Over 1..1000 the lengths whose pick is more than 5% slower than Bluestein's costs nearby + // number 374 on NEON f32 at a geometric mean of 1.405x, and 357 on SSE f32 at 1.306x. The + // pattern is a composite whose factorisation forces Rader's onto a large prime factor: + // Rader's permutation work is scalar, so at f32 it does not shrink while everything around + // it does. + // When planning, offer Bluestein's only where the direct route can actually be bad: some + // prime factor with no butterfly of its own, which is what forces Rader's or an awkward + // split. If every prime factor has a butterfly the decomposition is all butterflies and + // Bluestein's, which needs an inner FFT of at least 2*len - 1, cannot compete. Across the + // four 1..1000 sweeps the model picked Bluestein's at 1315 lengths and **not one** of them + // had all its prime factors covered, so this costs nothing and skips the enumeration at + // every smooth length. `candidates` still offers it everywhere, which is how that was + // checked and how it would be re-checked. + let bluesteins_worth_it = !planning || { + let butterflies = P::butterfly_lens(); + let mut n = len; + let mut uncovered = false; + let mut d = 2; + while d * d <= n { + while n % d == 0 { + uncovered |= !butterflies.contains(&d); + n /= d; + } + d += 1; + } + if n > 1 { + uncovered |= !butterflies.contains(&n); + } + uncovered + }; + + if len > 3 && bluesteins_worth_it { + let min_inner = 2 * len - 1; + let mut inner_lens: Vec = Vec::new(); + for multiplier in [1usize, 3, 5, 7, 9, 15] { + let mut candidate = multiplier; + while candidate < min_inner { + candidate *= 2; + } + inner_lens.push(candidate); + } + inner_lens.sort_unstable(); + inner_lens.dedup(); + for inner_len in inner_lens { + push( + Arc::new(Spec::Bluesteins { + len, + inner: planner.plan(inner_len), + }), + &mut out, + &mut seen, + ); + } + } + + // A bare butterfly, when one exists at this size. + if butterflies.contains(&len) { + push(Arc::new(Spec::Butterfly(len)), &mut out, &mut seen); + } + + out +} + +/// [`candidates`], trimmed to at most `cap` entries. +/// +/// A length like 100800 has hundreds of two-way splits. Everything structural is kept (the +/// planner's pick, Radix4 and RadixN shapes, Rader's and Bluestein's) and only splits are +/// dropped, most lopsided first, on the grounds that a split with a tiny side is mostly just its +/// large side plus a transpose. +/// What an estimating planner would actually enumerate at `len`. +/// +/// Identical to `candidates_capped`, except that it returns the fixed planner's pick immediately +/// at lengths where there is nothing to decide. Enumeration is pure overhead there, and it is the +/// overhead that matters most: plan-time is a large fraction of plan-plus-build at small lengths +/// and a negligible one at large lengths, so the cheapest lengths are exactly where an estimating +/// planner can least afford to enumerate. +/// +/// Two classes need no decision, and both were checked against measurement rather than assumed: +/// +/// - **A length with its own butterfly.** A single hand-written kernel beats any decomposition. +/// Over lengths 8..128 on NEON the bare butterfly is fastest at all fifteen such lengths, and +/// is never beaten by a split. +/// - **A power of two.** Radix4 wins, and the fixed planner already picks the right base, which +/// is the part that is not obvious: at 1024 `r4(3,b16)` beats `r4(4,b4)` by 1.15x. Across four +/// datasets, at every power of two from 64 up the fixed planner's pick is exactly the fastest +/// measured candidate, regret 1.000. +/// +/// It also drops the wider-of-the-two-first ordering of every split. Each two-way split is +/// otherwise enumerated twice, which roughly doubles the candidate count at a highly composite +/// length for almost no information: the two orderings differ only in how `transpose_small` walks +/// the rectangle and in which inner FFT runs first. The smaller-width ordering is the better one +/// in 90 to 97% of measured pairs for the Small variants, and for the general variants the two +/// are usually indistinguishable, which makes dropping one free rather than merely cheap. +/// +/// Measured over four datasets, that keeps 58 to 63% of candidates for a geometric mean regret of +/// 1.0016 or better against the full set. The worst single case is 1.129x at length 62 on NEON +/// f32, where `gts(b31,b2)` beats `gts(b2,b31)`; both known exceptions involve b31 or b32, where +/// the parallel-pair f32 butterflies make the chunk count matter in a way none of this models. +/// The planner's own pick is always element zero, so pruning can never leave an estimating +/// planner worse than the fixed one. +/// +/// `candidates` and `candidates_capped` stay exhaustive, because scoring a planner's pick needs +/// the alternatives even where a planner would not look at them. That is how every claim above +/// was established, and re-establishing them after a kernel change needs the same breadth. +pub fn plan_candidates>( + planner: &mut P, + len: usize, + cap: usize, +) -> Vec> { + if len.is_power_of_two() || P::butterfly_lens().contains(&len) { + return vec![planner.plan(len)]; + } + cap_list(candidates_inner(planner, len, true), cap) +} + +pub fn candidates_capped>( + planner: &mut P, + len: usize, + cap: usize, +) -> Vec> { + cap_list(candidates(planner, len), cap) +} + +/// Trim a candidate list to `cap`, keeping the planner's pick and the most balanced splits. +fn cap_list(all: Vec>, cap: usize) -> Vec> { + if all.len() <= cap { + return all; + } + + let imbalance = |spec: &Spec| -> Option { + match spec { + Spec::MixedRadix { left, right, .. } | Spec::GoodThomas { left, right, .. } => { + Some(((left.len() as f64).ln() - (right.len() as f64).ln()).abs()) + } + _ => None, + } + }; + + let mut kept: Vec> = Vec::with_capacity(cap); + let mut splits: Vec> = Vec::new(); + for (index, spec) in all.into_iter().enumerate() { + if index == 0 || imbalance(&spec).is_none() { + kept.push(spec); + } else { + splits.push(spec); + } + } + splits.sort_by(|a, b| { + imbalance(a) + .unwrap() + .partial_cmp(&imbalance(b).unwrap()) + .unwrap_or(std::cmp::Ordering::Equal) + }); + kept.extend(splits.into_iter().take(cap.saturating_sub(kept.len()))); + kept +} diff --git a/src/wasm_simd/wasm_simd_planner.rs b/src/wasm_simd/wasm_simd_planner.rs index e70a62ba..d2a02e3d 100644 --- a/src/wasm_simd/wasm_simd_planner.rs +++ b/src/wasm_simd/wasm_simd_planner.rs @@ -210,7 +210,7 @@ impl FftPlannerWasmSimd { } impl FftPlannerWasmSimd { - fn design_fft_for_len(&mut self, len: usize) -> Arc { + pub(crate) fn design_fft_for_len(&mut self, len: usize) -> Arc { if len < 1 { Arc::new(Recipe::Dft(len)) } else if let Some(recipe) = self.recipe_cache.get(&len) { @@ -223,7 +223,7 @@ impl FftPlannerWasmSimd { } } - fn build_fft(&mut self, recipe: &Recipe, direction: FftDirection) -> Arc> { + pub(crate) fn build_fft(&mut self, recipe: &Recipe, direction: FftDirection) -> Arc> { let len = recipe.len(); if let Some(instance) = self.algorithm_cache.get(len, direction) { instance diff --git a/tools/planner_tuning/.gitignore b/tools/planner_tuning/.gitignore new file mode 100644 index 00000000..5910fe01 --- /dev/null +++ b/tools/planner_tuning/.gitignore @@ -0,0 +1,8 @@ +target +Cargo.lock +# measurement logs +*.txt +*.tsv +*.log +sweep*.txt +.venv diff --git a/tools/planner_tuning/COST-MODEL.md b/tools/planner_tuning/COST-MODEL.md new file mode 100644 index 00000000..070d3a94 --- /dev/null +++ b/tools/planner_tuning/COST-MODEL.md @@ -0,0 +1,375 @@ +# How the estimating planner estimates + +This is the "how it works" document. `OP-COUNTS.md` is where the instruction counts come from, +`RESULTS.md` is the evidence that the thing works, `NEXT-STEPS.md` is the live plan, and `README.md` +is how to run the tools. This file explains the mechanism that sits under all four: what a cost is, +how one is computed, which numbers are read off the source, which are fitted against measurements, +and what has to be redone when the kernels change. + +**Status.** There is no estimating planner in the library. `src/tuning/` is a feature-gated +description of what the planners can build, and the cost model lives in the measurement tool at +`tools/planner_tuning/src/counted.rs`. What exists today is a full working prototype driven from the +command line, and `NEXT-STEPS.md` holds the open question of whether it ships whole or only as the +one scoped decision it is most clearly right about. + +## 1. The loop + +An estimating planner does three things, and only the middle one is new: + +1. **Enumerate.** Build every recipe the planner could plausibly use at this length, as a tree of + [`Spec`](../../src/tuning/mod.rs) nodes. +2. **Price.** Give each candidate a number. +3. **Pick the minimum.** + +Step 1 is [`plan_candidates`](../../src/tuning/mod.rs#L585). Step 2 is +[`CountedModel::cost`](src/counted.rs#L409). Step 3 is a `min_by`, visible in +[`cmd_sweep`](src/main.rs#L452). That is the whole planner. Everything else in this document is +about step 2. + +The enumeration deliberately answers immediately at a length with its own butterfly and at a power +of two, because measurement showed there is nothing to decide at either, and it emits only the +smaller-width-first ordering of each two-way split. The reasoning and the measured cost of each of +those shortcuts is in the doc comment on `plan_candidates`. Element zero of the list is always the +current fixed planner's pick, so pruning can never leave the estimating planner behind the planner +it replaces. + +The separate exhaustive set (`candidates`, `candidates_capped`) exists because *scoring* a pick +needs alternatives that no planner would ever propose. Keep it exhaustive; it is how each shortcut +above was justified. + +## 2. What a cost is + +**One unit is one issued arithmetic instruction.** A cost is not nanoseconds and does not try to be: +nothing converts it to time, and nothing needs to, because only the ranking within one length is +ever used. That is why a single weight set travels across machines of different clock speeds, and it +is also the reason the absolute numbers below look large and mean nothing on their own. + +Every cost is the sum of two halves: + +``` +cost = counted arithmetic instructions (read off the source, see OP-COUNTS.md) + + a memory term (accesses x pattern x cache level) +``` + +A pure operation count, which is the `FFTW_ESTIMATE` analogue, scores worse than the shipping +planner: mean regret 1.359 against 1.093. Adding the memory term takes the same op counts to 1.003. +The memory term is not a refinement, it is the entire result. + +### The memory term + +[`mem()`](src/counted.rs#L382) is four lines and deliberately crude: + +``` +accesses -> divided by complexes-per-vector, except under a permutation +level -> L1 if the working set fits in l1_elems, else L2 if it fits l2_elems, else DRAM +cost = accesses * seq[level] * (1.0 sequential | strided_mult | permuted_mult) +``` + +Three things about it are load-bearing and worth stating plainly: + +1. **Sequential versus jumpy carries all of it.** Pricing every access the same degrades worst-case + regret from 1.043 to 1.246. The finer distinction between strided and permuted does much less. +2. **The cache level cannot reorder candidates, as the model is built.** Collapsing L1/L2/DRAM to one + flat cost produces byte-identical picks at every length measured, in cache and out. That is + structural rather than a finding about memory: the level is chosen from the whole transform's + working set (next section), which is the same for every candidate at one length, so it is a + common factor. **It is not true of the hardware.** At 100003 and 100049 on the Pi 5, Rader's is + about 3x faster than every Bluestein's candidate, because Bluestein's inner FFT is 262144 points + (4 MB of complex f64), inside the M1's 12 MB L2 and far outside the A76's 512 KB L2 and 2 MB L3. + No setting of the cache sizes or level weights can express that. See blind spot 4. +3. **A permuted pass is charged per complex number, not per vector.** A gather or scatter computes + an address per element and cannot fill a vector. This is invisible at f64, where the factor is 1, + and was worth a factor of 2 at f32; finding it is what took SSE f32 from 51 to 117 of 216 weight + settings clearing the 20% bar. `--permuted-vector` restores the old behaviour. + +### The working set is threaded down unchanged + +[`cost_ws`](src/counted.rs#L416) passes the *whole transform's* length to every nested algorithm, +not the nested algorithm's own length. Every pass of every inner FFT walks the top-level buffer, so +that is what decides where the traffic is served from. Pricing an inner FFT as if it ran standalone +is the specific mistake that sank the 2021 attempt. + +The exception is Bluestein's, whose inner FFT runs on a buffer of its own that is at least twice the +outer length. Threading the outer length down understates that working set, which is exactly the +Pi 5 defect above. Charging a Bluestein's node at its inner length would fix it and still keep a +node's cost a function of its own subtree, which is what length-keyed recipe memoisation needs. Not +tried yet. + +One caveat on reading `explain` output: it calls `cost()` per node, so each row is priced at its +own length as working set. Since the cache level cannot reorder candidates this almost never +changes a number, but the child rows of a very large transform are informational rather than exact +contributions. + +## 3. What each node costs + +All of this is [`cost_ws`](src/counted.rs#L416), one match arm per `Spec` variant. `len` is the +node's own length and `ws` the whole transform's. + +| node | arithmetic | memory | +|---|---|---| +| `Butterfly(len)` | counted table lookup | `2*len` sequential | +| `Radix4 { k, base }` | `reps` x base, plus per layer `len/4` butterfly4 and 3 twiddles each | one `2*len` permuted digit reversal, plus `2*len` strided per layer | +| `RadixN { radixes, base }` | `reps` x base, plus per layer `len/r` x (butterfly `r` + `r-1` twiddles) | same shape as Radix4, plus `radixn_extra` per element per layer | +| `MixedRadix` | `h` x left + `w` x right, plus `len` twiddle multiplies | three transposes, plus one sequential pass | +| `GoodThomas` | `h` x left + `w` x right, no twiddles at all | two permuted CRT passes plus one transpose | +| `Raders { inner }` | 2 x inner, `len` twiddles, `len * rader_index` | two permuted passes plus one sequential | +| `Bluesteins { len, inner }` | 2 x inner, `inner.len()` pointwise multiplies | sequential over the inner length and twice over the outer | +| `Dft(n)` | `100 * n^2` | none; a quadratic that only has to sort last | + +Two of these encode a decision the model would otherwise be unable to make: + +- **`small` versus general.** `MixedRadixSmall` and `GoodThomasAlgorithmSmall` call + `array_utils::transpose_small`, a naive strided double loop, so their transposes are priced + `Permuted`. The general variants hand the job to the `transpose` crate, which tiles the rectangle + and gets the reuse back, and pay `general_row` per row for the privilege. That makes the general + form cheaper per element and dearer per row, which is a crossover, and the measurements are the + shape of one: `MixedRadix` crosses under its Small variant near length 200. Without the per-row + term the model has only per-element costs, so it would pick the general form at every length. +- **Which ordering of a split.** `small_row * max(width - height, 0)` is the only thing that + separates `mrs(A,B)` from `mrs(B,A)`. `transpose_small`'s outer loop runs `width` times, so both + Small variants change by exactly `width - height` when the pair is reversed, which is why one + weight covers both. Charging the *difference* rather than the raw count keeps a square pair free + and leaves the better ordering at exactly the cost it had before the term existed. + +### A worked example, which is also a self-check + +``` +$ ./target/release/planner_tuning explain 'mrs(b8,rad(rn(3.3,b9)))' +recipe len times total cost own cost +mrs(b8,rad(rn(3.3,b9))) 656 x1 127628 13776 + b8 8 x82 4428 4428 + rad(rn(3.3,b9)) 82 x8 109424 69536 + rn(3.3,b9) 81 x16 39888 28080 + b9 9 x144 11808 11808 +``` + +The `b8` row by hand: 38 counted instructions from the NEON f64 table, plus `2*8 = 16` sequential +accesses at `seq[0] = 1.0` and a sequential multiplier of 1.0, is 54. It runs 82 times, giving 4428. + +The root's own cost by hand: three transposes at `3 * 2*656 * 2.5` permuted is 9840, no `small_row` +charge because the width is the smaller side already, plus `656` twiddle multiplies at 4 +instructions each is 2624, plus a `2*656` sequential pass is 1312. Total 13776. + +If a change to `counted.rs` is supposed to leave some family of recipes alone, `explain` on one of +them is the fastest way to confirm it did. + +## 4. Where every number comes from + +This is the distinction the whole approach rests on. Most of the model is read off the source and +must be maintained when the source changes; a small set of weights is fitted against measurements +and must be refitted when the machine changes. + +### Read off the source, never fitted + +| quantity | where it is read | notes | +|---|---|---| +| butterfly instruction counts | `src/neon/*.rs`, `src/sse/*.rs`, by hand | tables in `butterfly_compute`, derived in `OP-COUNTS.md` | +| generated prime butterflies | the generator's loop structure | closed form `(h-1)(2h+5)`, cannot drift when sizes are added | +| `mul_complex`, `column_butterfly4` | the vector trait impls | 4 on NEON f64, 6 on SSE and on NEON f32 | +| `registers` | the architecture | 32 `v` on aarch64, 16 `xmm` on x86-64 | +| access counts per pass | the algorithm source | loads plus stores, counted per pass | +| access *pattern* per pass | the algorithm source | which of sequential, strided, permuted applies | +| `l1_elems`, `l2_elems` | the machine's cache sizes | inert as the model is built: no pick depends on them. Defaults are the M1's | + +### Derived from the code, then confirmed against measurement + +| quantity | derivation | +|---|---| +| `rader_index` = 30 (f64), 45 (f32) | `raders_algorithm.rs` recomputes `index * root % len` per element, a loop-carried `mul -> umulh -> mul -> sub` chain of about 10 cycles. That is worth tens of instruction slots, not the 7 instructions counted, but how many depends on the core's instructions per cycle, so the optimum is per machine: about 40 on the M1, 25 to 30 on a Cortex-A76. The defaults are the values acceptable on all three measured machines (M1, Pi 5, ThinkCentre) by the 1..1000 sweep, not the optimum on any one. See `NEXT-STEPS.md`. | +| `small_row` = 10 | One outer iteration of `transpose_small`, which measures 1.48 ns on the i3 and 0.7 to 1.0 ns on the M1, or 9 to 13 instruction-equivalents at either machine's scale. Scores are byte-identical for anything from 2 to 24 on every dataset, because the term only ever separates two orderings that are otherwise exactly equal. It is a tie-break with a derivation, not a weight. | + +### Fitted by comparing against measured times + +Six numbers, and only these six. They convert a memory access into arithmetic-instruction +equivalents, which is the one thing that cannot be read off RustFFT's source because it is a +property of the machine. + +| parameter | what it prices | fitted over | +|---|---|---| +| `seq[0]`, `seq[1]`, `seq[2]` | one access at L1, L2, DRAM | pinned at 1.0; the other two are inert, see below | +| `strided_mult` | a fixed-stride pass against a sequential one | 1.0, 1.5, 2.5, 4.0 | +| `permuted_mult` | a gather or scatter against a sequential pass | 1.5, 2.5, 4.0, 6.0 | +| `radixn_extra` | the generic `SimdRadixN` driver against the hand-written `Radix4` kernel, per element per layer | 0, 1, 2, 3, 5, 8; defaults 0 on NEON, 6 on SSE f64, 1 on SSE f32, chosen by the 1..1000 sweep | +| `general_row` | per-row setup of the blocked transpose and the on-the-fly CRT mapping | the twelve measured general-over-small ratios | + +`spill` and `mul_complex` are diagnostic overrides rather than fitted weights. Both are off by +default. They exist so that a specific hypothesis can be tested (register pressure in the RadixN +cross layer, and whether SSE's shuffle-heavy complex multiply costs more than its instruction count) +without editing the model. + +### The fitted values, per dataset + +``` +dataset seq[1] strided permuted radixn_extra mean worst +NEON f64 any 1.5 1.5 0 1.003 1.043 +NEON f32 3.0 2.5 2.5 0 1.032 1.121 +SSE f64 2.0 2.5 1.5 5 1.052 1.152 +SSE f32 1.5 2.5 2.5 3 1.034 1.121 +``` + +**Each of those four is one (machine, backend, element type) cell.** Every NEON number comes from +the M1 and every SSE number from the ThinkCentre, so on those two alone "the NEON weights" and "the +M1 weights" are the same column. Which of the two it really is decides whether this ships: + +- **If the split is per backend**, it costs nothing. The backend is chosen at compile time plus a + feature check, and the element type is a generic parameter, so each combination can carry its own + constants exactly as it already carries its own op counts. Four constant sets, all clearing the + bar. +- **If the split is per machine**, no compiled-in constant is right for anybody, and the whole + approach needs a runtime calibration step that nothing here has designed. + +**A second ARM machine, the Pi 5 (Cortex-A76), says mostly per backend.** Its 1..1000 sweeps with +the M1's NEON weights give the same picks at every length, so only the timing differs. Grouped by +which algorithm each planner chose, every class of decision lands within 1.5% of the M1 except +Rader's versus Bluestein's. That includes MixedRadix versus GoodThomas, which `RESULTS.md` had as +the one preference following the machine (wasm on the M1 within 0.004 of NEON on the M1, SSE on the +ThinkCentre 0.088 away): on the Pi those calls score 1.009 in f64 and 1.018 in f32, against 1.008 +and 1.030 on the M1. + +The exception is genuinely per machine. `rader_index` is a latency converted into instruction slots, +so it depends on the core: about 40 on the M1, 25 to 30 on the A76. It is handled without runtime +calibration, by choosing the default that is acceptable on all three machines rather than optimal +on one (`NEXT-STEPS.md` has the three-machine table). That is the rule for any future per-machine weight too: sweep it +everywhere and take the value whose worst machine looks best. + +### The floor, if one set had to serve everything + +Worth knowing because it bounds the damage. Gridding a single memory weight set jointly against all +four dumps, allowing only `radixn_extra` to differ per backend since that one is a register-file +property, the best is `strided 1.5, permuted 2.5, radixn_extra 0 on NEON and 5 on SSE`, with the +cache levels irrelevant as always: + +``` + one shared set own weights fixed planner +NEON f64 1.003 / 1.043 1.003 / 1.043 1.093 / 1.495 +NEON f32 1.032 / 1.225 1.032 / 1.121 1.171 / 1.969 +SSE f64 1.085 / 1.323 1.052 / 1.152 1.246 / 1.734 +SSE f32 1.052 / 1.250 1.034 / 1.121 1.207 / 1.927 +``` + +(mean / worst.) Three of the four then miss the 20% target, so this is not the proposal. But it +still beats the fixed planner on both statistics on all four datasets, and NEON f64 loses nothing at +all. The weights are worth getting right; getting them wrong degrades the result rather than +inverting it. + +`radixn_extra` is the one fitted weight that is not just a fudge: it is exactly 0 on NEON and +positive on SSE, which is what a register-count argument predicts, since `cross_layer` holds 2R rows +live and 2R fits 32 `v` registers at every supported radix and does not fit 16 `xmm`. It also shrinks +from 5 to 2-3 when the element type halves, because a spilled register then covers twice the +elements. Predicted 2.5, observed 2 to 3. + +**That was before the RadixN transpose fix** (415a29f), which removed per-call divides the weight had +been partly absorbing. Refitted by the 1..1000 sweep on the ThinkCentre, the defaults are now 6 for +SSE f64 and 1 for SSE f32, so the halving prediction no longer holds. Zero on NEON still does. The +table above predates the fix; `NEXT-STEPS.md` has the refit. + +Note that `Params::default()` is the NEON working set with `permuted_mult` at 2.5 rather than the +fitted 1.5. It makes no difference to that dataset, but a run that means to reproduce a table above +should pass the weights explicitly rather than trust the defaults. + +## 5. How the weights are fitted + +The procedure is a grid search against a frozen measurement, with a held-out half. + +```sh +# 1. measure once. every candidate at every length, times written to a TSV. +./target/release/planner_tuning dump --planner neon --rounds 7 --cap 48 \ + --out dump_neon_f64.tsv + +# 2. split into halves. even-indexed lengths train, odd-indexed test, on sorted length, +# so each half spans the whole size range. +python3 split.py dump_neon_f64.tsv train.tsv test.tsv + +# 3. grid the weights against the training half only. +./grid.sh train.tsv # 216 points +./sweep.sh train.tsv # the older 108-point grid + +# 4. score the winner on the test half, once. +./target/release/planner_tuning score --seq-l2 1.5 --strided 1.5 --permuted 1.5 \ + --rader-index 30 test.tsv +``` + +Steps 2 to 4 are **pure replay**. No machine is involved, no planner is built, and a full grid takes +seconds. That is the property that makes this maintainable where the old measured-table model was +not: one measurement run per machine, then unlimited model iteration anywhere. + +**Fit on the training half only.** Gridding on the full dump and then quoting a held-out number is +not a held-out number. The honest SSE f32 held-out worst case is 1.163, not the 1.105 that a +fit-on-everything run reports. + +The metric is **regret**: measured time of the chosen recipe divided by measured time of the best +enumerated candidate, so 1.000 is optimal. It is a lower bound on the distance from optimal, because +the candidate set is finite and the inner recipes inside each candidate come from the same planner. + +Two sanity properties of the fit are worth knowing before touching a weight: + +- **96 of 108 settings clear the 20% bar** on NEON f64. The result does not depend on hitting the + weights precisely, which is the main reason to think they travel. +- **The DRAM weight is inert.** 3.0, 6.0, 10.0 and 16.0 give byte-identical results, for the + structural reason in section 2. +- **A flat weight on the tuning set proves nothing about the sweep.** The Rader's weight is flat + from 15 to 120 on the original 33 lengths, because none of them is a small prime. The 1..1000 + sweep moves geometric mean by up to 4% between 30 and 45. + +## 6. Updating the model when the code changes + +This is the cost of the approach. The model is accurate *because* it tracks the source, and that +means the source moving invalidates it. The upside is that every such update is a re-count, which is +mechanical and needs no machine, rather than a re-measurement campaign. + +| what changed | what to redo | +|---|---| +| a butterfly's body | re-count it into `OP-COUNTS.md`, update the table in `butterfly_compute` | +| a butterfly length added or removed | nothing for the generated primes, they are a closed form; a table entry for a hand-written one | +| a vector primitive, for instance FMA or `vcmlaq` arriving | update `Backend::mul_complex` or `column_butterfly4`, then re-count every butterfly built from it. This is what the fcma work would trigger | +| an algorithm's pass structure, for instance `raders_precompute` landing | re-derive that node's arm in `cost_ws`. A precomputed Rader's permutation drops the per-element cost by roughly 4x and `rader_index` stops being a latency term at all | +| a transpose or index computation swapped | recheck the `Pattern` on that pass, and whether `general_row` still describes the same per-row work | +| a new algorithm in a planner | a `Spec` variant, an adapter arm, and a `cost_ws` arm | +| a new backend | a `Backend` variant, its counts, and its register file size | +| a new machine | nothing counted changes. Run `sweep 1..1000` in f64 and f32 and compare it class by class against a machine on the same backend. Where a weight's optimum moves, choose the value acceptable on every machine, never refit to the new one alone | + +After any of them, in this order: + +1. `verify` at a spread of lengths, first. It checks every enumerated candidate against an f64 + reference DFT, which catches an illegal spec such as a Bluestein's inner shorter than `2n - 1`. + Never trust a timing taken before `verify` is clean. +2. `explain` on a recipe whose cost you can predict, as in the worked example above. +3. `score` against an existing dump, which is free, to see whether the ranking moved at all. +4. `sweep --planner

1..1000` if it did. Both cost-model defects found so far were invisible to + the 33- and 44-length tuning sets and showed up only in the full sweep. +5. A fresh `dump` and refit only if the counts changed enough to move the fitted weights, which is + unusual: the weights price memory, and a kernel change usually moves arithmetic. + +## 7. Known blind spots + +Short version; `RESULTS.md` has the numbers and `NEXT-STEPS.md` has what to do about each. + +1. **Width and height are tied.** The cost function gives `mr(A,B)` and `mr(B,A)` the same cost at + all 415 reversed pairs, apart from the `small_row` tie-break, yet 133 to 179 of them measure more + than 2% apart. This is the clearest unexploited improvement and it is derivable from the code. +2. **MixedRadix versus GoodThomas is a fixed preference, not a decision.** With identical inners the + cost difference is a constant multiple of `len`, so its sign cannot vary with length. It picks + GoodThomas at 142 of 142 pairs, which is right 121 times on the M1 and 38 times on the i3. Keep + any correction small: a stride-aware rewrite aimed at this regressed both backends and was + reverted. +3. **Three machines, two per backend on ARM only.** The Pi 5 separated machine from backend on + ARM (section 4), but x86 still has only the ThinkCentre. A Zen 5 box is the instrument for that. +4. **Bluestein's working set is not priced.** The cache level is taken from the outer transform + (section 2), so a Bluestein's inner FFT two to four times the outer length is charged as if it + fitted where the outer one does. Invisible on the M1 and in any 1..1000 sweep. On the Pi 5 it + picks Bluestein's at 100003 and 100049 where Rader's is about 3x faster. +5. **Plan time.** Enumerate-and-price is 20x to 1265x the fixed planner's plan time, which is the + wrong denominator: what a caller pays is plan plus build, and building is 113x planning overall. + Against plan-plus-build, medians on the M1 at `--cap 48`: + + ``` + len fixed plan plan+price build extra on plan+build in FFT executions + 1260 0.6 us 62.5 us 8.0 us +720% ~12 + 10080 0.5 us 138 us 63 us +219% ~3 + 100800 0.5 us 284 us 537 us +53% ~0.5 + ``` + + So the cost is about twelve executions of the transform being planned at the worst measured + length, under one above 100k, and zero at butterfly lengths and powers of two where enumeration + short-circuits. Memoising inner recipes across candidates is the obvious optimisation and has + deliberately not been done. diff --git a/tools/planner_tuning/Cargo.toml b/tools/planner_tuning/Cargo.toml new file mode 100644 index 00000000..6837611e --- /dev/null +++ b/tools/planner_tuning/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "planner_tuning" +version = "0.1.0" +edition = "2021" +publish = false + +# The SIMD planner available depends on the architecture, so the feature does too. Both builds +# always include the scalar planner, which needs no feature of its own. +[target.'cfg(target_arch = "aarch64")'.dependencies] +rustfft = { path = "../..", default-features = false, features = ["neon", "tuning"] } + +[target.'cfg(target_arch = "x86_64")'.dependencies] +rustfft = { path = "../..", default-features = false, features = ["sse", "tuning"] } + +# Wasm needs a runtime with a working monotonic clock to time anything, so build for +# wasm32-wasip1 and run under wasmtime rather than wasm32-unknown-unknown. +[target.'cfg(target_arch = "wasm32")'.dependencies] +rustfft = { path = "../..", default-features = false, features = ["wasm_simd", "tuning"] } + +[target.'cfg(not(any(target_arch = "aarch64", target_arch = "x86_64", target_arch = "wasm32")))'.dependencies] +rustfft = { path = "../..", default-features = false, features = ["tuning"] } + +[profile.release] +opt-level = 3 +lto = true +codegen-units = 1 diff --git a/tools/planner_tuning/OP-COUNTS.md b/tools/planner_tuning/OP-COUNTS.md new file mode 100644 index 00000000..c1996ad5 --- /dev/null +++ b/tools/planner_tuning/OP-COUNTS.md @@ -0,0 +1,165 @@ +# NEON and SSE f64 operation counts, and how they were derived + +Every number here was obtained by **reading the source**, not by benchmarking. The unit is one +issued NEON instruction. The counts feed the analytic cost model as leaf costs, replacing the +measured butterfly table. + +Counting was done by hand, using `grep`/`python` only as a calculator to histogram intrinsic names +inside a function body. Nothing infers cost from a timing. + +## Primitive costs + +From `src/neon/neon_vector.rs` (the `impl NeonVector for float64x2_t` block) and +`src/neon/neon_utils.rs` (`impl Rotate90F64`). + +| primitive | instr | derivation | +|---|---|---| +| `vaddq_f64`, `vsubq_f64`, `vmulq_f64`, `vfmaq_f64`, `vnegq_f64`, `vneg_f64`, `veorq_u64`, `vcombine_f64`, `vmulq_laneq_f64`, `vfmaq_laneq_f64` | 1 | one instruction each | +| `vreinterpretq_*`, `vget_low_f64`, `vget_high_f64` | 0 | pure bit-pattern views, or fold into the consuming `vcombine` | +| `vmovq_n_f64` of a literal | 0 | loop-invariant constant, hoisted | +| `vld1q_f64` / `vst1q_f64` (`load_complex` / `store_complex`) | 1 | one load or store | +| `solo_fft2_f64` | **2** | `vaddq_f64` + `vsubq_f64` | +| `NeonVector::column_butterfly2` | **2** | same, one add and one sub | +| `Rotate90F64::rotate`, `NeonVector::apply_rotate90` | **2** | `vcombine_f64` + `veorq_u64`; the `vget_*` fold in | +| `Rotate90F64::rotate_45` / `_135` / `_225` | **4** | `rotate` (2) + one `vaddq`/`vsubq` (1) + one `vmulq` (1) | +| `NeonVector::mul_complex` | **4** | `vcombine_f64` + `vneg_f64` + `vmulq_laneq_f64` + `vfmaq_laneq_f64` | +| `NeonVector::column_butterfly4` | **10** | 4 x `column_butterfly2` (8) + 1 x `apply_rotate90` (2) | + +`mul_complex` costing 4 is the single most load-bearing primitive, since it is what every twiddle +multiply in RadixN, Radix4 and MixedRadix costs. Note the source comment: ARMv8.2-A `vcmlaq_f64` +would collapse this to 1-2, which is what the `fcma_backend` branch is about. + +## Load and store + +Every butterfly reads each of its `len` complex inputs exactly once and writes each of its `len` +outputs exactly once, so **loads = stores = len**, giving `2 * len` memory instructions per call. +This was confirmed by inspection rather than grepped, because the larger butterflies (16, 24, 32) +issue their loads from inside a `let load = |i| {...}` closure invoked once per column, which a +naive text count gets wrong. + +## Hand-written butterflies + +From `src/neon/neon_butterflies.rs`, counting the body of `perform_fft_direct` (or of +`perform_fft_contiguous` for 16, 24 and 32, which have no separate `direct`), with `new()` excluded +because it is construction, not per-call work. Sub-butterfly calls are resolved bottom-up. + +| len | composition | compute instr | +|---|---|---| +| 1 | nothing | 0 | +| 2 | 1 x `solo_fft2` | **2** | +| 3 | 2 `vaddq` + 1 `vsubq` + 3 `vfmaq` + 1 `rotate` | **8** | +| 4 | 4 x `solo_fft2` + 1 `rotate` | **10** | +| 5 | 11 `vaddq` + 5 `vsubq` + 8 `vmulq` + 2 `rotate` | **28** | +| 6 | 2 x bf3 + 3 x `solo_fft2` | **22** | +| 8 | 2 x bf4 + 3 `rotate` + 1 `vaddq` + 1 `vsubq` + 2 `vmulq` + 4 x `solo_fft2` | **38** | +| 9 | 3x3: 6 x bf3 + 4 x `mul_complex` | **64** | +| 10 | 5x2 Good-Thomas: 2 x bf5 + 5 x bf2, no twiddles | **66** | +| 12 | 3 x bf4 + 4 x bf3 | **62** | +| 15 | 3 x bf5 + 5 x bf3 | **124** | +| 16 | 4x4: 8 x bf4 + 4 x `mul_complex` + 1 `rotate` + 2 `rotate_45` + 2 `rotate_135` | **114** | +| 24 | 6 x bf4 + 4 x bf6 + 8 x `mul_complex` + 1 `neg` + 2 `rotate` + 2 `rotate_45` + 1 `rotate_135` + 1 `rotate_225` | **201** | +| 32 | 8x4: 8 x bf4 + 4 x bf8 + 16 x `mul_complex` + 1 `rotate` + 2 `rotate_45` + 2 `rotate_135` | **314** | + +Two traps worth recording, both of which produced wrong counts on the first pass: + +1. Calls are often split across lines (`self\n .bf4\n .perform_fft_direct(...)`), so a + line-oriented match misses them. Flatten whitespace first. +2. Butterfly 32 calls `self.bf8.bf4.perform_fft_direct` for its eight size-4 column FFTs, a + three-level path. Missing those made it look cheaper per element than butterfly 16, which was + the signal that the count was wrong. + +## Generated prime butterflies + +From `src/neon/neon_prime_butterflies.rs`, which is emitted by +`tools/gen_simd_butterflies/src/main.rs:252-310`. Because the generator is a pair of loops, the +count is a **closed form** rather than a table, so it cannot drift when sizes are added or removed. + +With `h = (len + 1) / 2`, the generator emits, for each of the `h - 1` conjugate pairs: + +- one `column_butterfly2` (2), one `apply_rotate90` (2) and one `add` (1) in the input stage; +- an a-chain of `h - 1` `fmadd` (1 each); +- a b-chain of one `mul` plus `h - 2` `fmadd`/`nmadd` (1 each); +- one closing `column_butterfly2` (2). + +That is `(h-1) * [1 + 1 + (2h-3) + 2 + 4] = (h-1)(2h+5)` instructions. + +| len | 7 | 11 | 13 | 17 | 19 | 23 | 29 | 31 | +|---|---|---|---|---|---|---|---|---| +| h | 4 | 6 | 7 | 9 | 10 | 12 | 15 | 16 | +| compute | **39** | **85** | **114** | **184** | **225** | **319** | **490** | **555** | + +Verified against a direct histogram of the generated source for all eight lengths: exact match. + +## Resulting per-element cost + +`compute / len`, the number that decides whether a bigger RadixN base pays for itself: + +| len | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 15 | 16 | 17 | 19 | 23 | 24 | 29 | 31 | 32 | +|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---| +| instr/elem | 1.0 | 2.7 | 2.5 | 5.6 | 3.7 | 5.6 | 4.8 | 7.1 | 6.6 | 7.7 | 5.2 | 8.8 | 8.3 | 7.1 | 10.8 | 11.8 | 13.9 | 8.4 | 16.9 | 17.9 | 9.8 | + +The shape is the expected one: powers of two are cheapest per element, the prime butterflies grow +roughly linearly in `len` because they are O(n^2) kernels, and the composites sit in between. + +## SSE f64 + +Counted the same way, from `src/sse/sse_vector.rs`, `src/sse/sse_utils.rs` and +`src/sse/sse_butterflies.rs`. The decompositions are the same as NEON's; the primitive costs are +not, because SSE4.1 has no FMA. + +| primitive | NEON | SSE | SSE derivation | +|---|---|---|---| +| `fmadd` / `nmadd` | 1 | **2** | `_mm_mul_pd` plus `_mm_add_pd` / `_mm_sub_pd` | +| `mul_complex` | 4 | **6** | `_mm_unpacklo_pd`, `_mm_unpackhi_pd`, two `_mm_mul_pd`, `_mm_shuffle_pd`, `_mm_addsub_pd` | +| `apply_rotate90`, `Rotate90F64::rotate` | 2 | 2 | `_mm_shuffle_pd` + `_mm_xor_pd` | +| `rotate_45` / `_135` / `_225` | 4 | 4 | rotate, then add or sub, then mul | +| `column_butterfly2`, `solo_fft2_f64` | 2 | 2 | add + sub | +| `column_butterfly4` | 10 | 10 | four `column_butterfly2` plus one `apply_rotate90` | +| `add`, `mul`, `neg`, load, store | 1 | 1 | one instruction each | + +Prime butterflies re-derive to `(h-1)(4h+2)`, since only the `fmadd` chain changes weight. Verified +exactly against a histogram of `src/sse/sse_prime_butterflies.rs` for all eight lengths: 7 -> 54, +11 -> 130, 13 -> 180, 17 -> 304, 19 -> 378, 23 -> 550, 29 -> 868, 31 -> 990. + +Hand-written butterflies. **Two of these are not the NEON figure re-weighted**, which is why the +SSE source was counted rather than scaled: + +| len | compute | note | +|---|---|---| +| 1 | 0 | | +| 2 | 2 | | +| 3 | **10** | written without FMA as 4 add + 2 mul + 2 sub + 1 rotate, not NEON's 8 re-weighted to 11 | +| 4 | 10 | | +| 5 | 28 | | +| 6 | 26 | 2 x bf3 + 3 x solo_fft2 | +| 8 | **38** | uses `rotate_45` and `rotate_135` where NEON uses explicit multiplies; lands on the same total by coincidence | +| 9 | 84 | 6 x bf3 + 4 x mul_complex | +| 10 | 66 | 2 x bf5 + 5 x bf2, Good-Thomas, no twiddles | +| 12 | 70 | 3 x bf4 + 4 x bf3 | +| 15 | 134 | 3 x bf5 + 5 x bf3 | +| 16 | 122 | 8 x bf4 + 4 x mul_complex + rotations | +| 24 | 233 | 6 x bf4 + 4 x bf6 + 8 x mul_complex + neg + rotations | +| 32 | 346 | 8 x bf4 + 4 x bf8 + 16 x mul_complex + rotations | + +As on NEON, `src/sse/sse_radixn.rs` wires its cross-FFT layers to these same +`SseF64ButterflyN::perform_fft_direct` functions, so the table applies to RadixN directly. + +## NEON f32 + +Counted from `perform_parallel_fft_direct`, which is what `src/neon/neon_radixn.rs` wires the +cross-FFT layers to for f32 and what the `Fft` boilerplate uses for chunk pairs. That function +computes **two** FFTs at once, so the figures below are the raw count halved, giving a per-FFT cost +directly comparable with the f64 table. + +Primitive differences from f64, read from `impl NeonVector for float32x4_t`: `nmadd` costs 2 rather +than 1 (`vfmaq_f32` plus `vnegq_f32`) and `mul_complex` costs 6 rather than 4 (`vtrn1q`, `vtrn2q`, +`vnegq`, `vmulq`, `vrev64q`, `vfmaq`). The f32 rotate helpers in `neon_utils.rs` are `rotate_both` +at 2, `rotate_hi` at 3, and `rotate_both_45/135/225` at 4. + +| len | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 15 | 16 | 17 | 19 | 23 | 24 | 29 | 31 | 32 | +|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---| +| per FFT | 2 | 4 | 5 | 16 | 11 | 19.5 | 19 | 36 | 37 | 44.5 | 31 | 61.5 | 68 | 66 | 100 | 125.5 | 180 | 113.5 | 279.5 | 321 | 178 | + +Per element these are 0.50x to 0.58x the f64 figures, except length 2 at 1.00x, where the packing +needed by `parallel_fft2_contiguous_f32` eats the whole gain. That near-uniformity is why the table +turns out not to matter; see the f32 section of `RESULTS.md`. diff --git a/tools/planner_tuning/README.md b/tools/planner_tuning/README.md new file mode 100644 index 00000000..c4260269 --- /dev/null +++ b/tools/planner_tuning/README.md @@ -0,0 +1,148 @@ +# planner_tuning + +A measurement harness for RustFFT's planners. Not part of the library and not shipped: it exists +to answer "which recipe is actually fastest at this length, and does the planner pick it?" + +Everything here works against a `TunablePlanner`, so the same commands run on the scalar, NEON, +SSE and wasm planners. Recipes are built through the planner's own internals, so what gets timed +is exactly what the planner would construct. + +## Vocabulary + +A **recipe** is a `Spec` tree, printed compactly: `mr(b11,rad(mrs(b6,b10)))` is a MixedRadix of +butterfly 11 against Rader's over a MixedRadixSmall of 6 and 10. The prefixes are `b` butterfly, +`r4` Radix4, `rn` RadixN, `mr`/`mrs` MixedRadix and its Small variant, `gt`/`gts` Good-Thomas and +its Small variant, `rad` Rader's, `bs` Bluestein's. + +**Regret** is a pick divided by the best measured candidate at that length: 1.000 is optimal, +1.05 is five percent off. It is a *lower bound* on the distance from optimal, because the +candidate set is finite. + +## The two workflows + +### Measure once, replay forever + +`dump` times every candidate at each length and writes a TSV. `score`, `costs` and `explain` then +replay that file offline, so iterating on the cost model needs no machine after the first run. +This is how the weights were fitted. + +```sh +cargo build --release +./target/release/planner_tuning dump --planner neon --rounds 7 --cap 48 \ + --out dump_neon_f64.tsv 1000 1050 1200 1296 +./target/release/planner_tuning score dump_neon_f64.tsv # regret vs best measured +./target/release/planner_tuning costs dump_neon_f64.tsv # per-candidate cost and ns +./target/release/planner_tuning explain 'rad(rn(7.6,b24))' # cost tree for one recipe +``` + +### Sweep a whole range + +`sweep` times the planner's pick against the cost model's pick at every length in a range. Unlike +`dump` it builds only those two recipes per length, which is what makes a thousand lengths +affordable. Where the two agree the recipe is timed once and reported for both, so agreement +reads as exactly 1.000 rather than as noise. + +This is the instrument that matters. Both cost-model defects fixed so far were invisible to the +33- and 44-length tuning sets and were found only by sweeping 1..1000. + +```sh +./target/release/planner_tuning sweep --planner neon 1..1000 > sweep_neon_f64.tsv +./target/release/planner_tuning sweep --planner neon --f32 1..1000 > sweep_neon_f32.tsv +``` + +Lengths may be a list or an `A..B` range. + +## Plotting + +`plot_sweep.py` draws the two views worth having: the population normalised by `N log2 N`, and +the per-length ratio. It prints the same figures it draws, so it doubles as the reporting tool. +It takes several files at once, which is how cost-model stages and machines get compared. + +matplotlib is the only third-party dependency in this directory, so it lives in a venv: + +```sh +python3 -m venv .venv && .venv/bin/pip install matplotlib +.venv/bin/python plot_sweep.py sweep_neon_f64.tsv # one run +.venv/bin/python plot_sweep.py sweep_*_f64_*.tsv # compare stages or machines +.venv/bin/python plot_sweep.py sweep_neon_f64.tsv --lo 4 --hi 128 # zoom +.venv/bin/python plot_sweep.py sweep_neon_f64.tsv --out fig.png # write instead of show +``` + +## Plan time + +`plantime` compares the cost of *choosing* a recipe against the cost of *building* it. Planning +only produces a `Recipe`; turning that into an `Arc` allocates and precomputes, and +Rader's and Bluestein's both run a full inner FFT inside their constructors. Building is 2x to +900x planning and grows much faster with length, so the estimating planner's overhead is a large +fraction of plan-plus-build at small lengths and a small one at large lengths. + +```sh +./target/release/planner_tuning plantime --planner neon --cap 48 128 1024 1200 10007 +``` + +Single plan-time measurements are noisy, up to 1.6x apart at length 1200. Take medians of three. +The candidate counts it prints are exact. + +`crossover` goes further and asks which recipe would win if construction cost counted: it builds and +times every candidate, then reports the minimiser of `build + k * execute` at several `k`, plus how +many executions the cost model's pick needs to repay its extra build cost. Measured answer is three +to six, so this mostly documents why a construction-aware planner is not worth building. + +```sh +./target/release/planner_tuning crossover --planner neon --cap 24 1260 1009 2018 +``` + +## Two candidate sets, deliberately + +- `candidates` / `candidates_capped` are **exhaustive**: every split in both orders, every + algorithm that can express it, Bluestein's at every length. Used by `dump`, `verify` and + `regret`, because scoring a planner's pick needs the alternatives even where no planner would + look at them. +- `plan_candidates` is **what an estimating planner would really enumerate**. Used by `sweep` and + `plantime`. It skips work that measurement has shown carries no decision: butterfly lengths and + powers of two return the fixed planner's pick immediately, only the smaller-width-first ordering + of each split is emitted, and Bluestein's is offered only where some prime factor has no + butterfly of its own. + +Keep the first set exhaustive. It is how each of those shortcuts was justified, and re-justifying +them after a kernel change needs the same breadth. + +## Running on another machine + +The campaign is SSH-driven. The remote checkouts are rsync copies rather than git worktrees, so +their `.git` points at a path that does not exist there; sync source only and leave the data. + +```sh +rsync -az --delete src/ user@host:~/repos/RustFFT-testing/src/ +rsync -az --delete tools/planner_tuning/src/ \ + user@host:~/repos/RustFFT-testing/tools/planner_tuning/src/ +ssh user@host 'cd ~/repos/RustFFT-testing/tools/planner_tuning && cargo build --release' +``` + +On x86 the tuning crate selects the `sse` feature automatically, so a plain `cargo build` gives +the SSE planner. + +## Traps + +- **`sweep` and `plantime` do not infer the backend.** `score` and `costs` read it from the dump + header, but `sweep` takes the model parameters as given, so an SSE run needs `--backend sse` + explicitly. Without it the model prices SSE recipes with NEON instruction counts and the whole + run looks plausible and is meaningless. +- **`--f32` is needed on `score` as well as on `dump`.** The dump header records the planner, so + the backend is recovered, but it does not record the element type. +- **Do not compile while measuring.** Build first, then run. +- **zsh does not word-split unquoted variables**, so `$LENGTHS` arrives as one argument and the + tool panics. Inline the numbers or use `${=VAR}`. +- **Run `verify` after any change to candidate enumeration.** It checks every enumerated candidate + against a direct DFT, which catches an illegal spec such as a Bluestein's inner shorter than + `2n - 1`. +- Measurement output (`*.tsv`, `*.txt`, `*.log`) and `.venv` are gitignored. + +## The documents + +- `COST-MODEL.md` is how the estimating planner works: what a cost is, how one is computed, which + numbers are counted and which are fitted, and what has to be redone when a kernel changes. Start + here. +- `RESULTS.md` is the evidence: how the cost model was built and what it scores. +- `OP-COUNTS.md` is where the instruction counts come from. +- `NEXT-STEPS.md` is the live plan, and the place new findings get written down. diff --git a/tools/planner_tuning/plot_sweep.py b/tools/planner_tuning/plot_sweep.py new file mode 100755 index 00000000..78102d3f --- /dev/null +++ b/tools/planner_tuning/plot_sweep.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 +"""Plot the output of `planner_tuning sweep`. + +The sweep times the shipping planner's pick against the cost model's at every length in a +range. This draws the two views worth having: + + top cost per transform normalised by N log2 N, both planners, so the shape of the + whole population is visible at once + bottom the ratio planner/model per length, one trace per input file, so several cost + model stages or several machines can be compared directly + +Usage: + plot_sweep.py sweep.tsv # one run + plot_sweep.py a.tsv b.tsv c.tsv # compare runs in the ratio panel + plot_sweep.py *.tsv --out fig.png # write instead of showing + plot_sweep.py sweep.tsv --lo 4 --hi 128 # zoom a length range + +Needs matplotlib, which is the only third-party dependency in this directory: + python3 -m venv .venv && .venv/bin/pip install matplotlib + .venv/bin/python plot_sweep.py sweep_neon_f64.tsv +""" +import argparse +import math +import os +import sys + +try: + import matplotlib.pyplot as plt +except ImportError: + sys.exit(__doc__.rsplit("Needs matplotlib", 1)[0] + + "matplotlib is not installed. From this directory:\n" + " python3 -m venv .venv && .venv/bin/pip install matplotlib\n" + " .venv/bin/python plot_sweep.py ") + + +def load(path): + """Rows of a sweep TSV, plus whatever the header recorded about the run.""" + head, rows = {}, [] + for line in open(path): + if line.startswith("#"): + parts = line[1:].strip().split("\t") + if len(parts) >= 2: + head[parts[0]] = "\t".join(parts[1:]) + continue + if line.startswith("len\t"): + continue + f = line.rstrip("\n").split("\t") + if len(f) < 10: + continue + rows.append(dict(n=int(f[0]), agree=f[2] == "1", + planner=float(f[3]), model=float(f[4]), ratio=float(f[5]), + pnorm=float(f[6]), mnorm=float(f[7]), + pspec=f[8], mspec=f[9])) + if not rows: + sys.exit(f"{path}: no data rows") + return head, rows + + +def gmean(xs): + return math.exp(sum(map(math.log, xs)) / len(xs)) if xs else float("nan") + + +def summarise(path, head, rows): + """The same figures the plot shows, as text, because they are what gets quoted.""" + dis = [r for r in rows if not r["agree"]] + loss = [r for r in dis if r["ratio"] < 1 / 1.02] + win = [r for r in dis if r["ratio"] > 1.02] + worst = min(dis, key=lambda r: r["ratio"]) if dis else None + best = max(dis, key=lambda r: r["ratio"]) if dis else None + print(f"{os.path.basename(path)} [{head.get('planner', '?')}]") + print(f" {len(rows)} lengths, {len(rows) - len(dis)} where both planners agree") + print(f" geometric mean planner/model {gmean([r['ratio'] for r in rows]):.4f}") + print(f" model wins beyond 2% at {len(win)}, loses beyond 2% at {len(loss)}") + if worst: + print(f" best {best['ratio']:.3f}x at N={best['n']} model {best['mspec']}") + print(f" worst {worst['ratio']:.3f}x at N={worst['n']} model {worst['mspec']}" + f" planner {worst['pspec']}") + print() + + +def main(): + ap = argparse.ArgumentParser(description="Plot planner_tuning sweep output.") + ap.add_argument("files", nargs="+", help="sweep TSVs") + ap.add_argument("--lo", type=int, default=4, help="lowest length to plot (default 4)") + ap.add_argument("--hi", type=int, default=None, help="highest length to plot") + ap.add_argument("--out", help="write the figure here instead of showing it") + ap.add_argument("--quiet", action="store_true", help="skip the text summary") + args = ap.parse_args() + + loaded = [] + for path in args.files: + head, rows = load(path) + rows = [r for r in rows + if r["n"] >= args.lo and (args.hi is None or r["n"] <= args.hi)] + if rows: + loaded.append((path, head, rows)) + if not loaded: + sys.exit("nothing left after the length filter") + + if not args.quiet: + for path, head, rows in loaded: + summarise(path, head, rows) + + fig, (ax, bx) = plt.subplots( + 2, 1, figsize=(13, 8), height_ratios=[2, 1], sharex=True, + gridspec_kw=dict(hspace=0.12)) + + # Top: the population, from the first file only. Overlaying several runs here just + # produces mud, and the runs differ in the model's pick rather than the planner's. + path, head, rows = loaded[0] + ax.scatter([r["n"] for r in rows], [r["pnorm"] for r in rows], + s=7, alpha=0.65, label="fixed planner", color="#2a78d6") + ax.scatter([r["n"] for r in rows], [r["mnorm"] for r in rows], + s=7, alpha=0.65, label="estimating planner", color="#eb6834") + ax.set_ylabel("ns per N log2 N") + ax.set_title(f"{os.path.basename(path)} planner {head.get('planner', '?')}") + ax.legend(loc="upper right", frameon=False) + ax.grid(alpha=0.25) + + # Bottom: every file, so stages and machines line up against the same 1.0. + for path, head, rows in loaded: + bx.plot([r["n"] for r in rows], [r["ratio"] for r in rows], + lw=0.9, alpha=0.85, label=os.path.basename(path)) + bx.axhline(1.0, color="0.4", lw=1) + # Log scale so 1.25x up and 1.25x down are the same distance from 1.0, but with + # explicit ticks: the default log locator labels only the decade and 1.0 is the + # only decade anywhere near this data. + bx.set_yscale("log") + ticks = [0.5, 0.67, 0.8, 0.9, 1.0, 1.1, 1.25, 1.5, 2.0] + lo_r = min(r["ratio"] for _, _, rows in loaded for r in rows) + hi_r = max(r["ratio"] for _, _, rows in loaded for r in rows) + ticks = [t for t in ticks if lo_r / 1.05 <= t <= hi_r * 1.05] or [1.0] + bx.set_yticks(ticks) + bx.set_yticklabels([f"{t:g}x" for t in ticks]) + bx.minorticks_off() + bx.set_ylabel("planner / model") + bx.set_xlabel("transform length N") + bx.grid(alpha=0.25, which="both") + if len(loaded) > 1: + bx.legend(loc="upper right", frameon=False, fontsize=8, ncol=2) + bx.annotate("estimate faster", xy=(0.005, 0.92), xycoords="axes fraction", fontsize=8) + bx.annotate("planner faster", xy=(0.005, 0.04), xycoords="axes fraction", fontsize=8) + + if args.out: + fig.savefig(args.out, dpi=140, bbox_inches="tight") + print(f"wrote {args.out}") + else: + plt.show() + + +if __name__ == "__main__": + main() diff --git a/tools/planner_tuning/run_wasm.mjs b/tools/planner_tuning/run_wasm.mjs new file mode 100644 index 00000000..505f964e --- /dev/null +++ b/tools/planner_tuning/run_wasm.mjs @@ -0,0 +1,14 @@ +// Minimal WASI harness: runs planner_tuning.wasm under node with argv/stdio passthrough. +import { readFile } from 'node:fs/promises'; +import { WASI } from 'node:wasi'; +const wasi = new WASI({ + version: 'preview1', + args: ['planner_tuning', ...process.argv.slice(2)], + env: {}, + preopens: { '.': '.' }, + returnOnExit: true, +}); +const bytes = await readFile(new URL('./target/wasm32-wasip1/release/planner_tuning.wasm', import.meta.url)); +const wasm = await WebAssembly.compile(bytes); +const instance = await WebAssembly.instantiate(wasm, wasi.getImportObject()); +process.exitCode = wasi.start(instance); diff --git a/tools/planner_tuning/src/counted.rs b/tools/planner_tuning/src/counted.rs new file mode 100644 index 00000000..f09ae927 --- /dev/null +++ b/tools/planner_tuning/src/counted.rs @@ -0,0 +1,564 @@ +//! A cost model built by counting instructions in the source, not by measuring. +//! +//! Every leaf cost here comes from reading `src/neon/*.rs`; the derivation is written up in +//! `OP-COUNTS.md`. On top of the arithmetic count sits a coarse memory term: each pass over the +//! buffer is charged per element touched, scaled by how it walks memory (sequential, strided, or +//! permuted) and by which level of an *assumed* cache hierarchy the working set lands in. +//! +//! The point of the exercise is that nothing in here needs a machine. The only quantities that +//! are not read off the source are the handful of weights in `Params`, which set the price of a +//! memory access relative to one arithmetic instruction. + +use rustfft::tuning::Spec; + +/// Which backend's kernels to price. The decomposition each butterfly uses is the same on both, +/// but the instruction cost of the primitives is not: SSE4.1 has no FMA, so `fmadd` is a separate +/// multiply and add, and its complex multiply needs six instructions rather than four. +#[derive(Copy, Clone, Debug, PartialEq)] +pub enum Backend { + Neon, + Sse, +} + +/// Which element type. One 128-bit vector holds one complex f64 or two complex f32, so this +/// changes both the instruction counts and how many elements each memory access covers. +#[derive(Copy, Clone, Debug, PartialEq)] +pub enum Elem { + F32, + F64, +} + +impl Elem { + /// Complex numbers per 128-bit vector. + pub fn complex_per_vector(&self) -> f64 { + match self { + Elem::F32 => 2.0, + Elem::F64 => 1.0, + } + } +} + +impl Backend { + pub fn parse(name: &str) -> Option { + match name { + "neon" => Some(Backend::Neon), + "sse" => Some(Backend::Sse), + _ => None, + } + } + + /// `NeonVector::mul_complex` is 4 instructions; `SseVector::mul_complex` is 6 + /// (unpacklo, unpackhi, two muls, shuffle, addsub). + pub fn mul_complex(&self, elem: Elem) -> f64 { + match (self, elem) { + // vcombine + vneg + vmulq_laneq + vfmaq_laneq + (Backend::Neon, Elem::F64) => 4.0, + // vtrn1q + vtrn2q + vnegq + vmulq + vrev64q + vfmaq + (Backend::Neon, Elem::F32) => 6.0, + // unpacklo + unpackhi + 2 mul + shuffle + addsub + (Backend::Sse, _) => 6.0, + } + } + + /// `column_butterfly4` is four `column_butterfly2` plus one `apply_rotate90` on both. + pub fn column_butterfly4(&self) -> f64 { + 10.0 + } + + /// Architectural vector registers: 32 `v` registers on aarch64, 16 `xmm` on x86-64 SSE. + /// + /// This matters because `cross_layer` in `src/simd_radixn.rs` gathers **two** vector columns + /// before transforming either, so a radix-R layer holds 2R rows live at once, plus the + /// butterfly's own temporaries. At radix 7 that is 14 rows before temporaries, which fits + /// comfortably in 32 registers and not at all in 16. + pub fn registers(&self) -> f64 { + match self { + Backend::Neon => 32.0, + Backend::Sse => 16.0, + } + } + + /// Instructions for one `perform_fft_direct`, excluding the load and store of each element. + /// Hand-counted from `src/neon/neon_butterflies.rs` and `src/sse/sse_butterflies.rs`; the + /// derivation is in `OP-COUNTS.md`. + pub fn butterfly_compute(&self, len: usize, elem: Elem) -> Option { + // The generated prime butterflies come out as a closed form, because the generator is a + // pair of loops. Only the weight of the fmadd chain differs between the backends. + // f32 on NEON is counted from `perform_parallel_fft_direct`, which computes two FFTs + // at once, and stored here as the per-FFT figure. See OP-COUNTS.md. + if let (Backend::Neon, Elem::F32) = (self, elem) { + return Some(match len { + 1 => 0.0, 2 => 2.0, 3 => 4.0, 4 => 5.0, 5 => 16.0, 6 => 11.0, 7 => 19.5, + 8 => 19.0, 9 => 36.0, 10 => 37.0, 11 => 44.5, 12 => 31.0, 13 => 61.5, + 15 => 68.0, 16 => 66.0, 17 => 100.0, 19 => 125.5, 23 => 180.0, 24 => 113.5, + 29 => 279.5, 31 => 321.0, 32 => 178.0, + _ => return None, + }); + } + if matches!(len, 7 | 11 | 13 | 17 | 19 | 23 | 29 | 31) { + let h = ((len + 1) / 2) as f64; + return Some(match self { + Backend::Neon => (h - 1.0) * (2.0 * h + 5.0), + Backend::Sse => (h - 1.0) * (4.0 * h + 2.0), + }); + } + let v = match self { + Backend::Neon => match len { + 1 => 0, + 2 => 2, + 3 => 8, + 4 => 10, + 5 => 28, + 6 => 22, + 8 => 38, + 9 => 64, + 10 => 66, + 12 => 62, + 15 => 124, + 16 => 114, + 24 => 201, + 32 => 314, + _ => return None, + }, + // Same decompositions, but bf3 is written without FMA (10, not a re-weighted 8) and + // bf8 reaches for rotate_45/rotate_135 where NEON uses explicit multiplies. + Backend::Sse => match len { + 1 => 0, + 2 => 2, + 3 => 10, + 4 => 10, + 5 => 28, + 6 => 26, + 8 => 38, + 9 => 84, + 10 => 66, + 12 => 70, + 15 => 134, + 16 => 122, + 24 => 233, + 32 => 346, + _ => return None, + }, + }; + Some(v as f64) + } +} + +/// How a pass walks memory. +#[derive(Copy, Clone)] +pub enum Pattern { + /// Contiguous run, one cache line feeding many elements. + Sequential, + /// Fixed stride greater than a line, as in a transpose or a cross-FFT layer. + Strided, + /// Data-dependent scatter or gather: digit reversal, CRT reindexing, Rader's permutation. + Permuted, +} + +/// The weights that are not read off the source. +/// +/// Costs are in units of one arithmetic instruction. `seq` is the price of a single load or +/// store at each level of the hierarchy; the multipliers raise that for less friendly patterns. +#[derive(Copy, Clone, Debug)] +pub struct Params { + /// Complex numbers that fit in the first level of cache. + pub l1_elems: f64, + /// Complex numbers that fit in the last level of cache. + pub l2_elems: f64, + /// Cost of one load or store at L1, L2 and memory. + pub seq: [f64; 3], + pub strided_mult: f64, + pub permuted_mult: f64, + /// Cost of computing one Rader's permutation index, in arithmetic-instruction equivalents. + /// + /// `raders_algorithm.rs` recomputes `index = index * root % len` per element, where `len` is + /// a `StrengthReducedU64`, so that `%` is two 64x64->128 widening multiplies plus a shift and + /// a subtract: about 7 instructions on aarch64. The instruction count alone understates it, + /// because `index` feeds the next iteration, making the chain loop-carried and latency-bound + /// rather than throughput-bound. This weight converts the counted instructions into the + /// effective cost of that serial chain. + /// + /// **Size it from the latency, not the instruction count.** The carried chain is + /// `mul -> umulh -> mul -> sub`, which is about 10 cycles on both an M1 firestorm core and + /// Coffee Lake. Ten cycles of a core that retires 4 to 6 instructions per cycle is 40 to 60 + /// instruction slots, not the 7 instructions counted. The first default of 20 assumed a 3x + /// latency inflation and got 12 of 54 prime lengths wrong, at up to 1.64x. + /// + /// The term is linear in `len` while everything around it grows as `len log len`, so it + /// matters most at small lengths, which is why the original 33 lengths (all 1000 or larger) + /// could not pin it down. It decides the Rader's-versus-Bluestein's call. + /// + /// **It is not one value per machine, so the default is a compromise across machines.** Slots + /// per cycle depend on the core, and the optimum moves with it: 40 on the M1 and 25 to 30 on a + /// Cortex-A76. The default is chosen by the 1..1000 sweep on the M1, the Pi 5 and the + /// ThinkCentre, to be acceptable on all three rather than optimal on one: + /// + /// - f64: **30**. Fewer losses beyond 2% than 45 on every machine (33/28/46 against + /// 54/116/54), for 0.2% of geometric mean on the M1. + /// - f32: **45**. At or near the best on all three; 30 costs 2 to 4% of geometric mean. + /// + /// Why f32 wants a larger value is not understood. `complex_per_vector` already halves the + /// arithmetic per element while this chain stays per element. + /// + /// Negative means "use that per-element default"; see `CountedModel::rader_index`. + pub rader_index: f64, + /// Cost of one spilled vector per unrolled group in a RadixN cross layer, as a store plus a + /// reload. Zero disables the register-pressure term entirely. + pub spill: f64, + /// Extra cost per element per cross-FFT layer for the **generic** `SimdRadixN` driver, + /// over the hand-written `Radix4` kernel doing the same work. + /// + /// They are different code. `cross_layer` in `src/simd_radixn.rs` is generic over the radix + /// and gathers two vector columns before transforming either, so it holds 2R rows live plus + /// the butterfly's temporaries; `sse_radix4.rs` is a hardcoded 2x unroll over six twiddles. + /// 2R rows fits aarch64's 32 vector registers at every radix it supports, and does not fit + /// x86-64's 16 xmm registers, so this is expected to be near zero on NEON and positive on SSE. + /// + /// The defaults, chosen by the 1..1000 sweep on the ThinkCentre after the transpose indices + /// were precomputed (see `NEXT-STEPS.md`): + /// + /// - NEON: **0**, both element types. + /// - SSE f64: **6**. 38 losses beyond 2% against 50 at 5, 84 at 8 and 180 at 12. The + /// 33-length dump alone points at 12 to 16, but every one of those lengths is 1000 or more. + /// - SSE f32: **1**. 20 losses against 25 at both 0 and 2, and a worst case of 0.909 against + /// 0.861 at 2. + /// + /// The SSE values were 5 and 2 before that change. Removing a per-call overhead let the f32 + /// value come down, as expected; the f64 value moved up by one, not down. + /// + /// Negative means "use that per-backend default"; see `CountedModel::radixn_extra`. + pub radixn_extra: f64, + /// Override for the cost of one complex multiply. Negative means "use the backend's counted + /// value". Exists to test whether SSE's shuffle-heavy `mul_complex` costs more than its + /// instruction count suggests: three of its six instructions are shuffle-class, and on Intel + /// those all issue to a single port, whereas NEON spreads them over symmetric pipes. + pub mul_complex: f64, + /// Charge permuted passes per complex number rather than per vector. On by default. + /// + /// A gather or scatter moves one complex number at a time: digit reversal, CRT reindexing and + /// Rader's permutation all compute a destination per element, so there is no contiguous run to + /// fill a vector with. Dividing their access count by `complex_per_vector` therefore + /// under-charges them by exactly that factor, which is invisible at f64 (where the factor is + /// 1) and a factor of 2 at f32. + /// + /// Setting this false restores the original behaviour, which is what `--permuted-vector` is + /// for. That costs nothing at f64, where both f64 datasets score byte-identically either way, + /// and at f32 it forces the fitted `permuted_mult` up from 2.5 to between 4.0 and 6.0 to + /// absorb the same factor. See the f32-on-SSE section of `RESULTS.md`. + pub permuted_scalar: bool, + /// Fixed cost per row of a pass, charged to the **general** MixedRadix and GoodThomas + /// variants and not to their `Small` counterparts. + /// + /// The two pairs move the same data in the same order; they differ in implementation. + /// `MixedRadixSmall` and `GoodThomasAlgorithmSmall` call `array_utils::transpose_small`, + /// a naive strided double loop over the whole rectangle, and read their permutation from a + /// precomputed table. `MixedRadix` and `GoodThomasAlgorithm` call the `transpose` crate's + /// blocked transpose, which keeps both streams cache resident, and `GoodThomasAlgorithm` + /// additionally computes the CRT mapping on the fly with one `StrengthReducedUsize::div_rem` + /// and a branch per row rather than per element. + /// + /// The general form is therefore cheaper per element and dearer per row, which is a + /// crossover, and the measurements are the shape of one. General over small on the M1: + /// + /// ```text + /// len 22 28 45 104 496 992 + /// gt/gts 1.341 1.270 1.182 1.056 1.001 1.027 + /// mr/mrs 1.115 1.059 1.048 1.013 0.931 0.948 + /// ``` + /// + /// The advantage decays towards 1 and MixedRadix crosses under it near len 200. A per-element + /// difference alone could produce neither: it would hold roughly constant in ratio, and it + /// could never change sign. Without this term the model has only per-element costs, so it + /// prices the pair by pattern alone and picks the general form at every length. + /// + /// Every one of the twelve pairs above is called correctly for `general_row` anywhere in + /// 21 to 42; the binding constraints are Good-Thomas at 992 below and MixedRadix at 496 + /// above. 30 sits in the middle of that window. + pub general_row: f64, + /// Cost of one outer-loop iteration of `array_utils::transpose_small`, charged to the + /// **Small** MixedRadix and GoodThomas variants only. + /// + /// This is the term that makes the model prefer one ordering of a factor pair over its + /// reverse. `transpose_small` is a naive double loop: + /// + /// ```text + /// for x in 0..width { for y in 0..height { out[y + x*height] = in[x + y*width] } } + /// ``` + /// + /// The outer loop runs `width` times and the read index strides by `width`, so the cost + /// depends on which dimension is which. The general variants call the `transpose` crate, + /// which tiles the rectangle and so does not care: that contrast is the evidence, because + /// `GoodThomasAlgorithm` and `GoodThomasAlgorithmSmall` perform the *same single transpose + /// in the same orientation* and differ only in the implementation. Over reversed pairs the + /// small form measures smaller-width-faster at 90 to 97% on both machines and both element + /// types, while the general form splits about evenly and its median gap is 0.00 ns. + /// + /// Outer-loop iterations, counted from the source: + /// + /// - `GoodThomasAlgorithmSmall`: one transpose, `(width, height)`, so `width`. + /// - `MixedRadixSmall`: three, `(w,h)`, `(h,w)`, `(w,h)`, so `2*width + height`. + /// + /// Both change by exactly `width - height` when the pair is reversed, which predicts that + /// the two should show the same asymmetry per unit of `w - h` despite having different + /// absolute transpose counts. On SSE they measure 1.48 and 1.47 ns respectively. + /// + /// So the charge is `small_row * max(width - height, 0)`, not the raw iteration count. Both + /// variants differ by exactly `width - height` iterations between the two orderings, so one + /// weight covers both; charging the worse ordering that difference and the better one + /// nothing reproduces it. Only the *difference* is evidenced here, because the absolute + /// level of a Small variant against a general one is what `general_row` already carries, + /// fitted. Charging the difference rather than the count keeps three properties that matter: + /// + /// - a square pair is charged nothing, since there is no ordering to get wrong; + /// - the better ordering keeps exactly the cost it had before this term existed, so + /// `general_row` stays valid and the Small-versus-general balance is untouched; + /// - the cost never goes negative. + /// + /// Charging the raw count instead regressed SSE f64 at length 1215, where the recipe is + /// `mr(b15, mrs(b9,b9))`: the nested square pair was inflated by its 15 repetitions and the + /// whole recipe lost to an `rn(3.3.3.3,b15)` that is 1.21x slower. + /// + /// The value is one outer iteration in instruction-equivalents: about 1.48 ns on the i3 and + /// 0.7 to 1.0 ns on the M1, which at each machine's ns-per-cost-unit is 9 to 13 either way. + /// It barely matters. Scores are byte-identical for anything from 2 to 24 on every dataset, + /// because the term only ever separates two orderings that are otherwise exactly equal in + /// cost. It is a tie-break with a derivation, not a fitted weight. + pub small_row: f64, + /// Which backend's instruction costs to use. + pub backend: Backend, + /// Which element type. + pub elem: Elem, +} + +impl Default for Params { + fn default() -> Self { + // Apple M1 performance core: 128 KiB L1d, 12 MiB L2, at 16 bytes per complex f64. + Self { + l1_elems: 8192.0, + l2_elems: 786432.0, + seq: [1.0, 2.0, 6.0], + strided_mult: 1.5, + permuted_mult: 2.5, + rader_index: -1.0, + spill: 0.0, + radixn_extra: -1.0, + mul_complex: -1.0, + permuted_scalar: true, + general_row: 30.0, + small_row: 10.0, + backend: Backend::Neon, + elem: Elem::F64, + } + } +} + +pub struct CountedModel { + pub params: Params, +} + +impl CountedModel { + pub fn new(params: Params) -> Self { + Self { params } + } + + /// The complex-multiply cost actually in force. + fn mul_complex(&self) -> f64 { + if self.params.mul_complex >= 0.0 { + self.params.mul_complex + } else { + self.params.backend.mul_complex(self.params.elem) + } + } + + /// The Rader's index cost actually in force. See `Params::rader_index` for the values. + fn rader_index(&self) -> f64 { + if self.params.rader_index >= 0.0 { + self.params.rader_index + } else { + match self.params.elem { + Elem::F64 => 30.0, + Elem::F32 => 45.0, + } + } + } + + /// The RadixN driver cost actually in force. See `Params::radixn_extra` for the values. + fn radixn_extra(&self) -> f64 { + if self.params.radixn_extra >= 0.0 { + self.params.radixn_extra + } else { + match (self.params.backend, self.params.elem) { + (Backend::Neon, _) => 0.0, + (Backend::Sse, Elem::F64) => 6.0, + (Backend::Sse, Elem::F32) => 1.0, + } + } + } + + /// Cost of touching `accesses` elements (counting each load and each store once) with the + /// given pattern, when the enclosing buffer holds `ws` complex numbers. + /// Rows walked by the three passes of a width x height decomposition. + /// + /// The passes run over `height`, `width` and `height` rows respectively, so the exact total is + /// `2h + w`. The model deliberately ties `mr(A,B)` with `mr(B,A)`, so use the mean of the two + /// orderings, `1.5 * (w + h)`, rather than introduce an asymmetry here alone. + fn rows(&self, left: &Spec, right: &Spec) -> f64 { + 1.5 * (left.len() as f64 + right.len() as f64) + } + + fn mem(&self, accesses: f64, pattern: Pattern, ws: f64) -> f64 { + let p = &self.params; + // One load or store moves a whole vector, which is one complex f64 or two complex f32, + // except under a permutation, where each element's address is computed separately. + let per_access = match (pattern, p.permuted_scalar) { + (Pattern::Permuted, true) => 1.0, + _ => p.elem.complex_per_vector(), + }; + let accesses = accesses / per_access; + let level = if ws <= p.l1_elems { + 0 + } else if ws <= p.l2_elems { + 1 + } else { + 2 + }; + let mult = match pattern { + Pattern::Sequential => 1.0, + Pattern::Strided => p.strided_mult, + Pattern::Permuted => p.permuted_mult, + }; + accesses * p.seq[level] * mult + } + + /// Estimated cost of one FFT of this recipe, in arithmetic-instruction equivalents. + /// + /// `None` if a butterfly length has no counted entry, so a gap fails loudly. + pub fn cost(&self, spec: &Spec) -> Option { + self.cost_ws(spec, spec.len() as f64) + } + + /// `ws` is the working set of the whole transform, threaded down unchanged: every pass of + /// every nested algorithm walks the same top-level buffer, so that is what decides which + /// cache level the traffic is served from. + fn cost_ws(&self, spec: &Spec, ws: f64) -> Option { + let p = &self.params; + Some(match spec { + Spec::Dft(n) => { + let n = *n as f64; + 100.0 * n * n + } + Spec::Butterfly(len) => { + p.backend.butterfly_compute(*len, p.elem)? + self.mem(2.0 * *len as f64, Pattern::Sequential, ws) + } + Spec::Radix4 { k, base } => { + let len = spec.len() as f64; + let reps = len / base.len() as f64; + // One digit-reversal transpose, then the base FFTs, then k cross layers. + let mut c = self.mem(2.0 * len, Pattern::Permuted, ws); + c += reps * self.cost_ws(base, ws)?; + for _ in 0..*k { + // len/4 column_butterfly4, each with three twiddle multiplies. + c += (len / (4.0 * p.elem.complex_per_vector())) + * (p.backend.column_butterfly4() + 3.0 * self.mul_complex()); + c += self.mem(2.0 * len, Pattern::Strided, ws); + } + c + } + Spec::RadixN { radixes, base } => { + let len = spec.len() as f64; + let reps = len / base.len() as f64; + let mut c = self.mem(2.0 * len, Pattern::Permuted, ws); + c += reps * self.cost_ws(base, ws)?; + for r in radixes.iter() { + let rf = *r as f64; + // The cross-FFT layers call the very same butterfly kernels, so the counted + // table applies directly. Row 0 needs no twiddle, hence r - 1. + c += (len / rf) * (p.backend.butterfly_compute(*r, p.elem)? + (rf - 1.0) * self.mul_complex()); + c += self.mem(2.0 * len, Pattern::Strided, ws); + c += len * self.radixn_extra(); + // Register pressure: the layer keeps 2R rows live across the two-column + // unroll. Anything past the architectural register file becomes a spill and a + // reload, once per element of the group. + if p.spill > 0.0 { + let live = 2.0 * rf; + let over = (live - p.backend.registers()).max(0.0); + c += over * p.spill * (len / rf); + } + } + c + } + Spec::MixedRadix { left, right, small } => { + let len = spec.len() as f64; + // Three transposes, one full twiddle pass, two inner dimensions. + // + // Both variants transpose the same rectangle three times, but not the same way. + // `MixedRadixSmall` calls `transpose_small`, whose read index strides by `width` + // and so touches a fresh cache line per element once `width` exceeds a line: + // line-wasting, which is what `Permuted` prices. `MixedRadix` hands the job to + // the `transpose` crate, which tiles the rectangle to get that reuse back, and + // pays `general_row` per row of setup for it. + let pat = if *small { Pattern::Permuted } else { Pattern::Strided }; + let mut c = 3.0 * self.mem(2.0 * len, pat, ws); + if *small { + // transpose_small at (w,h), (h,w), (w,h) is 2*width + height outer + // iterations; reversing the pair gives 2*height + width, so the two + // orderings differ by width - height. See `small_row`. + c += p.small_row * (left.len() as f64 - right.len() as f64).max(0.0); + } else { + c += p.general_row * self.rows(left, right); + } + c += (len / p.elem.complex_per_vector()) * self.mul_complex() + + self.mem(2.0 * len, Pattern::Sequential, ws); + c += right.len() as f64 * self.cost_ws(left, ws)?; + c += left.len() as f64 * self.cost_ws(right, ws)?; + c + } + Spec::GoodThomas { left, right, small } => { + let len = spec.len() as f64; + // Two CRT reindexing passes and one transpose, but no twiddle multiplies at all: + // dropping them is the whole point of Good-Thomas, and it pays in index work. + // + // Both reindexing passes are `Permuted` in either variant. The small one gathers + // through a precomputed table; the general one walks `destination_index` forward + // by `width + 1` and wraps modulo `len`, which cycles over the whole buffer and + // is no friendlier to a cache than a table would be. The transpose splits the two + // exactly as in MixedRadix, and the general form pays the same per-row setup. + let pat = if *small { Pattern::Permuted } else { Pattern::Strided }; + let mut c = 2.0 * self.mem(2.0 * len, Pattern::Permuted, ws); + c += self.mem(2.0 * len, pat, ws); + if *small { + // One transpose_small at (width, height): `width` outer iterations, against + // `height` reversed. The same width - height difference as MixedRadixSmall, + // which is why one weight serves both. + c += p.small_row * (left.len() as f64 - right.len() as f64).max(0.0); + } else { + c += p.general_row * self.rows(left, right); + } + c += right.len() as f64 * self.cost_ws(left, ws)?; + c += left.len() as f64 * self.cost_ws(right, ws)?; + c + } + Spec::Raders { inner } => { + let len = spec.len() as f64; + // The inner FFT runs twice, and the permutation is precomputed into a u32 table, + // so it is a gather and a scatter rather than a modular multiply per element. + let mut c = 2.0 * self.cost_ws(inner, ws)?; + // Two permutation passes, each a scatter or gather whose index comes from a + // serial modular-multiply chain rather than from a table. + c += 2.0 * (self.mem(2.0 * len, Pattern::Permuted, ws) + len * self.rader_index()); + c += len * self.mul_complex() + self.mem(2.0 * len, Pattern::Sequential, ws); + c + } + Spec::Bluesteins { len, inner } => { + let outer = *len as f64; + let ilen = inner.len() as f64; + // Inner FFT twice, pointwise multiply over the padded inner length, and a + // twiddle-and-pad pass in and out over the outer length. + let mut c = 2.0 * self.cost_ws(inner, ws)?; + c += ilen * self.mul_complex() + self.mem(2.0 * ilen, Pattern::Sequential, ws); + c += 2.0 * (outer * self.mul_complex() + self.mem(2.0 * outer, Pattern::Sequential, ws)); + c + } + }) + } +} diff --git a/tools/planner_tuning/src/emit.rs b/tools/planner_tuning/src/emit.rs new file mode 100644 index 00000000..848b590a --- /dev/null +++ b/tools/planner_tuning/src/emit.rs @@ -0,0 +1,77 @@ +//! Emit a fitted model as a Rust source file for the Neon planner to include. +//! +//! The generated file is data, not logic: measured primitive costs and the fitted per-element +//! overheads. Keeping it separate from the planner keeps the boundary between "measured on some +//! machine" and "written by a person" obvious, which is what ejmahler asked for on PR #38. + +use crate::model::Model; + +/// Costs are stored relative to Butterfly32 rather than in nanoseconds, so that a table fitted +/// on one machine is at least meaningfully comparable on another. The planner only ever +/// compares costs, so the unit cancels. +const REFERENCE_BUTTERFLY: usize = 32; + +pub fn emit(model: &Model, element_type: &str, planner: &str) -> String { + let unit = model + .butterfly + .get(&REFERENCE_BUTTERFLY) + .copied() + .expect("butterfly 32 must be measured, it is the normalisation reference"); + + let mut out = String::new(); + out.push_str(&format!( + "// Generated by tools/planner_tuning. Do not edit by hand.\n\ + //\n\ + // Costs for the {} planner, relative to Butterfly{} for {} elements. Regenerate with:\n\ + // planner_tuning emit --planner {} LENGTHS...\n\n", + planner, REFERENCE_BUTTERFLY, element_type, planner + )); + + let mut butterflies: Vec<(&usize, &f64)> = model.butterfly.iter().collect(); + butterflies.sort_by_key(|(len, _)| **len); + out.push_str("/// Measured cost of one butterfly, by length.\n"); + out.push_str("pub(crate) const BUTTERFLY_COST: &[(usize, f32)] = &[\n"); + for (len, cost) in butterflies { + out.push_str(&format!(" ({}, {:.5}),\n", len, cost / unit)); + } + out.push_str("];\n\n"); + + let mut shapes: Vec<(&(usize, u32), &f64)> = model.radix4.iter().collect(); + shapes.sort_by_key(|((base, k), _)| (*base, *k)); + out.push_str("/// Measured cost of one Radix4, by (base length, k).\n"); + out.push_str("pub(crate) const RADIX4_COST: &[(usize, u32, f32)] = &[\n"); + for ((base, k), cost) in shapes { + out.push_str(&format!(" ({}, {}, {:.5}),\n", base, k, cost / unit)); + } + out.push_str("];\n\n"); + + out.push_str( + "/// Fitted overhead per element, as (log2 of working set, cost), to be interpolated.\n\ + ///\n\ + /// These are curves rather than constants because the scattered reindexing in\n\ + /// GoodThomas and Rader's costs sharply more once the array outgrows L1.\n", + ); + for (name, key) in [ + ("MIXEDRADIX", "mr"), + ("MIXEDRADIX_SMALL", "mrs"), + ("GOODTHOMAS", "gt"), + ("GOODTHOMAS_SMALL", "gts"), + ("RADERS", "rad"), + ("BLUESTEINS", "bs"), + ] { + let table = model + .overhead + .get(key) + .unwrap_or_else(|| panic!("overhead '{}' was never fitted", key)); + let entries: Vec = table + .iter() + .map(|(bucket, value)| format!("({}, {:.6})", bucket, value / unit)) + .collect(); + out.push_str(&format!( + "pub(crate) const {}_OVERHEAD: &[(u32, f32)] = &[{}];\n", + name, + entries.join(", ") + )); + } + out +} diff --git a/tools/planner_tuning/src/main.rs b/tools/planner_tuning/src/main.rs new file mode 100644 index 00000000..5d817137 --- /dev/null +++ b/tools/planner_tuning/src/main.rs @@ -0,0 +1,1509 @@ +//! Measurement tool for RustFFT's planners. +//! +//! Works against any planner that implements `TunablePlanner`, selected with `--planner`. +//! Recipes are built through the planner's own internals, so what is measured here is exactly +//! what the planner would construct. +//! +//! Subcommands: +//! time SPEC... time recipes against each other; a '*reps' suffix runs a recipe +//! over a len*reps buffer, which is how inner FFTs are invoked +//! regret LEN... measure how far the planner's pick is from the best available +//! sweep LEN... | A..B time the planner's pick against the counted model's pick, as TSV +//! model TRAIN... 0 TEST... calibrate a cost model and score its picks the same way +//! residuals LEN... show how per-element overhead varies with working set +//! verify LEN... check every enumerated candidate against a direct DFT +//! emit LEN... print the fitted model as a Rust source file + +mod counted; +mod emit; +mod model; + +use model::{bucket_of, overhead_scale, Model, FIT_ORDER}; +use rustfft::num_complex::Complex; +use rustfft::num_traits::{ToPrimitive, Zero}; +use rustfft::tuning::{ + candidates_capped, parse, plan_candidates, to_spec_string, ScalarTuner, Spec, + TunablePlanner, +}; +use rustfft::{Fft, FftDirection, FftNum}; +use std::sync::Arc; +use std::time::Instant; + +// --------------------------------------------------------------------------- +// Timing +// --------------------------------------------------------------------------- + +struct Subject { + name: String, + fft: Arc>, + reps: usize, + buffer: Vec>, + scratch: Vec>, + rounds: Vec, +} + +impl Subject { + fn new(name: String, fft: Arc>, reps: usize) -> Self { + let buffer = vec![Complex::zero(); fft.len() * reps]; + let scratch = vec![Complex::zero(); fft.get_inplace_scratch_len()]; + Subject { + name, + fft, + reps, + buffer, + scratch, + rounds: Vec::new(), + } + } + + /// Wall-clock nanoseconds for `iters` passes over the whole buffer. + fn time_block(&mut self, iters: usize) -> f64 { + let start = Instant::now(); + for _ in 0..iters { + self.fft + .process_with_scratch(&mut self.buffer, &mut self.scratch); + } + start.elapsed().as_secs_f64() * 1e9 + } + + /// Nanoseconds per individual FFT, ie per chunk of `fft.len()`. + fn time_per_fft(&mut self, iters: usize) -> f64 { + let reps = self.reps; + self.time_block(iters) / (iters * reps) as f64 + } + + fn best(&self) -> f64 { + self.rounds.iter().cloned().fold(f64::INFINITY, f64::min) + } + + fn median(&self) -> f64 { + let mut values = self.rounds.clone(); + values.sort_by(|a, b| a.partial_cmp(b).unwrap()); + values[values.len() / 2] + } +} + +/// Time a set of subjects against each other, round-robin. +/// +/// Round-robin rather than one at a time so any drift over the run hits every subject equally, +/// and min-of-rounds rather than a mean because the quantity of interest is the cost with +/// nothing else interfering. +fn measure(subjects: &mut [Subject], rounds: usize, block_ms: f64) { + let iters: Vec = subjects + .iter_mut() + .map(|subject| { + let probe = subject.time_block(1); + (((block_ms * 1e6) / probe.max(1.0)).ceil() as usize).clamp(1, 20_000_000) + }) + .collect(); + + for (subject, &n) in subjects.iter_mut().zip(iters.iter()) { + subject.time_block(1 + n / 4); + } + + for _ in 0..rounds { + for (subject, &n) in subjects.iter_mut().zip(iters.iter()) { + let per_fft = subject.time_per_fft(n); + subject.rounds.push(per_fft); + } + } +} + +/// Keep this thread on a performance core. Without it the macOS scheduler is free to park a +/// long-running thread on an efficiency core, which shows up as bimodal timings. +#[cfg(target_os = "macos")] +fn request_performance_core() { + const QOS_CLASS_USER_INTERACTIVE: u32 = 0x21; + extern "C" { + fn pthread_set_qos_class_self_np(qos_class: u32, relative_priority: i32) -> i32; + } + unsafe { + pthread_set_qos_class_self_np(QOS_CLASS_USER_INTERACTIVE, 0); + } +} + +#[cfg(not(target_os = "macos"))] +fn request_performance_core() {} + +fn median_of(mut values: Vec) -> f64 { + values.sort_by(|a, b| a.partial_cmp(b).unwrap()); + values[values.len() / 2] +} + +// --------------------------------------------------------------------------- +// Calibration +// --------------------------------------------------------------------------- + +/// Measure every primitive the model needs: each butterfly, and each valid Radix4 shape. +/// +/// Each is timed over a buffer of at least 8192 elements, because primitives are almost always +/// invoked many times over a larger buffer rather than standalone, and a single cold call would +/// price in a startup cost they do not really pay in use. +fn calibrate_primitives>( + rounds: usize, + block_ms: f64, + max_len: usize, +) -> Model { + const REFERENCE_ELEMENTS: usize = 8192; + let mut planner = P::new(); + let mut model = Model::default(); + + let mut shapes: Vec> = Vec::new(); + let mut subjects: Vec> = Vec::new(); + + for len in P::butterfly_lens() { + let spec = Spec::Butterfly(len); + let fft = planner.build(&spec, FftDirection::Forward); + let reps = (REFERENCE_ELEMENTS / len).max(1); + shapes.push(None); + subjects.push(Subject::new(format!("b{}", len), fft, reps)); + } + + for base in P::radix4_bases() { + let mut k = 1u32; + while base * (1usize << (2 * k)) <= max_len { + let spec = Spec::Radix4 { + k, + base: Arc::new(Spec::Butterfly(base)), + }; + let len = spec.len(); + let fft = planner.build(&spec, FftDirection::Forward); + let reps = (REFERENCE_ELEMENTS / len).max(1); + shapes.push(Some((base, k))); + subjects.push(Subject::new(to_spec_string(&spec), fft, reps)); + k += 1; + } + } + + measure(&mut subjects, rounds, block_ms); + + for (shape, subject) in shapes.iter().zip(subjects.iter()) { + match shape { + Some(key) => { + model.radix4.insert(*key, subject.best()); + } + None => { + model.butterfly.insert(subject.fft.len(), subject.best()); + } + } + } + model +} + +/// Fit one overhead curve for each composing algorithm. +/// +/// Kinds are fitted in dependency order, since the residual of a MixedRadix that contains a +/// MixedRadixSmall only means anything once the Small's own overhead is known. +fn fit_overheads>( + model: &mut Model, + lengths: &[usize], + rounds: usize, + block_ms: f64, + cap: usize, + bucketed: bool, +) -> Vec<(Arc, f64)> { + let mut samples: Vec<(Arc, f64)> = Vec::new(); + for &len in lengths { + let mut planner = P::new(); + let specs = candidates_capped(&mut planner, len, cap); + let mut subjects: Vec> = specs + .iter() + .map(|spec| { + let fft = planner.build(spec, FftDirection::Forward); + Subject::new(to_spec_string(spec), fft, 1) + }) + .collect(); + measure(&mut subjects, rounds, block_ms); + for (spec, subject) in specs.iter().zip(subjects.iter()) { + samples.push((Arc::clone(spec), subject.best())); + } + } + + let mut fitted: Vec<&'static str> = Vec::new(); + for target in FIT_ORDER { + let usable: Vec<(u32, f64)> = samples + .iter() + .filter(|(spec, _)| spec.kind() == target && model.descendants_known(spec, &fitted)) + .filter_map(|(spec, measured)| { + model + .inner_cost(spec) + .map(|inner| (bucket_of(spec), (measured - inner) / overhead_scale(spec))) + }) + .collect(); + if usable.is_empty() { + eprintln!("warning: no calibration samples for '{}'", target); + continue; + } + + // One median per log2 bucket, but only for buckets with enough samples to mean anything. + // Sparse buckets are dropped and filled in by interpolation instead. + let mut table: Vec<(u32, f64)> = Vec::new(); + if bucketed { + let mut by_bucket: std::collections::BTreeMap> = Default::default(); + for (bucket, residual) in usable.iter() { + by_bucket.entry(*bucket).or_default().push(*residual); + } + const MIN_PER_BUCKET: usize = 3; + table = by_bucket + .iter() + .filter(|(_, values)| values.len() >= MIN_PER_BUCKET) + .map(|(bucket, values)| (*bucket, median_of(values.clone()))) + .collect(); + } + + // Too little data to describe a curve, or curves not wanted, so use one constant. + if table.len() < 2 { + table = vec![(0, median_of(usable.iter().map(|(_, r)| *r).collect()))]; + } + + let rendered: Vec = table + .iter() + .map(|(bucket, value)| format!("{}:{:.2}", 1usize << bucket, value)) + .collect(); + println!( + " {:<5} {:>3} buckets from {:>4} samples {}", + target, + table.len(), + usable.len(), + rendered.join(" ") + ); + model.overhead.insert(target, table); + fitted.push(target); + } + samples +} + +// --------------------------------------------------------------------------- +// Subcommands +// --------------------------------------------------------------------------- + +struct Options { + rounds: usize, + block_ms: f64, + cap: usize, + verbose: bool, + bucketed: bool, + /// Where `dump` writes its rows. + out: Option, + /// Weights for the counted model. + params: counted::Params, + backend_explicit: bool, +} + +fn cmd_time>(specs: &[String], opts: &Options) { + let mut planner = P::new(); + let mut subjects: Vec> = specs + .iter() + .map(|text| { + let (spec_text, reps) = match text.rsplit_once('*') { + Some((spec, reps)) => (spec, reps.parse().expect("bad repeat count")), + None => (text.as_str(), 1usize), + }; + let spec = parse(spec_text).unwrap_or_else(|e| { + eprintln!("error in spec '{}': {}", spec_text, e); + std::process::exit(2); + }); + let fft = planner.build(&spec, FftDirection::Forward); + Subject::new(text.clone(), fft, reps) + }) + .collect(); + + measure(&mut subjects, opts.rounds, opts.block_ms); + + println!( + "{:<46} {:>9} {:>7} {:>13} {:>9}", + "recipe", "len", "reps", "min ns", "spread" + ); + for subject in subjects.iter() { + let (min, med) = (subject.best(), subject.median()); + println!( + "{:<46} {:>9} {:>7} {:>13.1} {:>8.2}%", + subject.name, + subject.fft.len(), + subject.reps, + min, + (med - min) / min * 100.0 + ); + } + + if subjects.len() > 1 { + let best = subjects + .iter() + .map(|s| s.best()) + .fold(f64::INFINITY, f64::min); + println!("\nrelative to best:"); + for subject in subjects.iter() { + println!(" {:<46} {:>6.3}x", subject.name, subject.best() / best); + } + } +} + +fn cmd_regret>(lengths: &[usize], opts: &Options) { + println!( + "{:>9} {:>8} {:>10} {:>10} {}", + "len", "regret", "planner ns", "best ns", "best recipe (when it differs)" + ); + + let mut regrets: Vec<(f64, usize, String, String)> = Vec::new(); + + for &len in lengths { + let mut planner = P::new(); + let specs = candidates_capped(&mut planner, len, opts.cap); + let planner_spec = to_spec_string(&specs[0]); + + let mut subjects: Vec> = specs + .iter() + .map(|spec| { + let fft = planner.build(spec, FftDirection::Forward); + Subject::new(to_spec_string(spec), fft, 1) + }) + .collect(); + + measure(&mut subjects, opts.rounds, opts.block_ms); + + let planner_time = subjects[0].best(); + let best_index = subjects + .iter() + .enumerate() + .min_by(|a, b| a.1.best().partial_cmp(&b.1.best()).unwrap()) + .map(|(i, _)| i) + .unwrap(); + let best_time = subjects[best_index].best(); + let best_spec = subjects[best_index].name.clone(); + let regret = planner_time / best_time; + + println!( + "{:>9} {:>7.3}x {:>10.0} {:>10.0} {}", + len, + regret, + planner_time, + best_time, + if best_index == 0 { + "= planner".to_string() + } else { + best_spec.clone() + } + ); + if opts.verbose { + let mut ranked: Vec<&Subject> = subjects.iter().collect(); + ranked.sort_by(|a, b| a.best().partial_cmp(&b.best()).unwrap()); + for subject in ranked.iter().take(6) { + println!( + " {:>6.3}x {}", + subject.best() / best_time, + subject.name + ); + } + println!(" ({} candidates measured)", subjects.len()); + } + + regrets.push((regret, len, planner_spec, best_spec)); + } + + regrets.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + let n = regrets.len(); + let mean = regrets.iter().map(|r| r.0).sum::() / n as f64; + println!("\n--- regret: planner's pick divided by the best recipe measured ---"); + println!(" lengths {}", n); + println!(" mean {:.4}", mean); + println!(" median {:.4}", regrets[n / 2].0); + println!(" p90 {:.4}", regrets[(n * 9) / 10].0); + println!(" worst {:.4}", regrets[n - 1].0); + let losing = regrets.iter().filter(|r| r.0 > 1.02).count(); + println!(" more than 2% off the best: {} of {} lengths", losing, n); + println!("\nworst offenders:"); + for (regret, len, planner_spec, best_spec) in regrets.iter().rev().take(10) { + println!(" {:>8} {:.3}x", len, regret); + println!(" planner: {}", planner_spec); + println!(" best: {}", best_spec); + } +} + +/// Time the shipping planner's pick against the counted model's pick, at every length in a range. +/// +/// Unlike `regret`, this builds and times only two recipes per length rather than the whole +/// candidate set, which is what makes a thousand-length sweep affordable. When both agree, the +/// recipe is timed once and reported for both, so an agreement shows as exactly 1.000 rather than +/// as timing noise. +/// +/// Output is TSV on stdout, one row per length, ready to plot. +fn cmd_sweep>(lengths: &[usize], opts: &Options) { + let model = counted::CountedModel::new(opts.params); + + println!("# planner\t{}", P::label()); + println!("# params\t{:?}", opts.params); + println!("# rounds\t{}\tblock_ms\t{}\tcap\t{}", opts.rounds, opts.block_ms, opts.cap); + println!( + "len\tcands\tagree\tplanner_ns\tmodel_ns\tratio\tplanner_norm\tmodel_norm\tplanner_spec\tmodel_spec" + ); + + for &len in lengths { + let mut planner = P::new(); + let specs = plan_candidates(&mut planner, len, opts.cap); + if specs.is_empty() { + eprintln!("len {}: no candidates", len); + continue; + } + + let model_index = specs + .iter() + .enumerate() + .filter_map(|(i, spec)| model.cost(spec).map(|c| (i, c))) + .min_by(|a, b| a.1.partial_cmp(&b.1).unwrap()) + .map(|(i, _)| i) + .unwrap_or(0); + + let agree = model_index == 0; + let planner_spec = to_spec_string(&specs[0]); + let model_spec = to_spec_string(&specs[model_index]); + + let mut subjects: Vec> = if agree { + vec![Subject::new( + planner_spec.clone(), + planner.build(&specs[0], FftDirection::Forward), + 1, + )] + } else { + vec![ + Subject::new( + planner_spec.clone(), + planner.build(&specs[0], FftDirection::Forward), + 1, + ), + Subject::new( + model_spec.clone(), + planner.build(&specs[model_index], FftDirection::Forward), + 1, + ), + ] + }; + measure(&mut subjects, opts.rounds, opts.block_ms); + + let planner_ns = subjects[0].best(); + let model_ns = if agree { planner_ns } else { subjects[1].best() }; + + // n log2 n, the work a radix-2 FFT of this length would do. Undefined at len 1, where + // there is no work to normalise by. + let nlogn = (len as f64) * (len as f64).log2(); + let norm = |ns: f64| if nlogn > 0.0 { ns / nlogn } else { f64::NAN }; + + println!( + "{}\t{}\t{}\t{:.2}\t{:.2}\t{:.4}\t{:.5}\t{:.5}\t{}\t{}", + len, + specs.len(), + if agree { 1 } else { 0 }, + planner_ns, + model_ns, + planner_ns / model_ns, + norm(planner_ns), + norm(model_ns), + planner_spec, + model_spec + ); + } +} + +/// Check that every enumerated candidate actually computes a correct FFT. +/// +/// Timing a recipe says nothing about whether it is valid, and an invalid one would happily +/// produce fast wrong answers. Each candidate is compared against a direct DFT. +/// Error budget for a correct FFT of this length in this precision. +/// +/// Two terms, because two things are inexact. +/// +/// The recipe's own error grows like `eps * sqrt(log2 len)`: an FFT accumulates rounding over its +/// `log2 len` passes, not over its elements. The margin of 20 leaves room for recipes that compose +/// several algorithms, and still catches a real defect, which shows up as orders of magnitude +/// rather than as a factor of two. +/// +/// The reference's error grows like `eps * sqrt(len)`, because a naive DFT really does sum `len` +/// terms per output. The reference runs in f64, so for f32 recipes this term is negligible and the +/// budget is tight. For f64 recipes the reference is no better than the thing it judges, and this +/// term dominates: at len 100000 it is 2.8e-13 against the recipes' 1.8e-14. +fn tolerance(len: usize) -> f64 { + let eps = if std::mem::size_of::() == 4 { + f32::EPSILON as f64 + } else { + f64::EPSILON + }; + let len = len as f64; + 20.0 * eps * len.log2().max(1.0).sqrt() + 4.0 * f64::EPSILON * len.sqrt() +} + +fn cmd_verify>(lengths: &[usize], opts: &Options) { + let mut worst_overall: f64 = 0.0; + let mut failures = 0usize; + + for &len in lengths { + let input: Vec> = (0..len) + .map(|i| { + let x = ((i * 2654435761usize) % 1000) as f64 / 500.0 - 1.0; + let y = ((i * 40503usize) % 1000) as f64 / 500.0 - 1.0; + Complex::new(T::from_f64(x).unwrap(), T::from_f64(y).unwrap()) + }) + .collect(); + + // Reference: a direct DFT, which shares no code with the recipes under test. + // + // It is computed in f64 even when the recipes run in f32. A same-precision reference is + // useless at large f32 lengths: the naive DFT sums `len` terms, so its own error grows + // with length and swamps what is being measured. At len 100000 an f32 reference DFT is + // off by 1.6e-5, which is 30x the error of the recipes it is judging. + let reference_fft = rustfft::algorithm::Dft::::new(len, FftDirection::Forward); + let mut reference: Vec> = input + .iter() + .map(|c| Complex::new(c.re.to_f64().unwrap(), c.im.to_f64().unwrap())) + .collect(); + let mut reference_scratch = vec![Complex::zero(); reference_fft.get_inplace_scratch_len()]; + reference_fft.process_with_scratch(&mut reference, &mut reference_scratch); + let reference_norm: f64 = reference.iter().map(|c| c.norm_sqr()).sum::().sqrt(); + + let mut planner = P::new(); + let specs = candidates_capped(&mut planner, len, opts.cap); + let mut worst_here: f64 = 0.0; + let mut worst_spec = String::new(); + + for spec in specs.iter() { + let fft = planner.build(spec, FftDirection::Forward); + let mut buffer = input.clone(); + let mut scratch = vec![Complex::zero(); fft.get_inplace_scratch_len()]; + fft.process_with_scratch(&mut buffer, &mut scratch); + + let error: f64 = buffer + .iter() + .zip(reference.iter()) + .map(|(a, b)| { + let d = Complex::new( + a.re.to_f64().unwrap() - b.re, + a.im.to_f64().unwrap() - b.im, + ); + d.norm_sqr() + }) + .sum::() + .sqrt() + / reference_norm; + if error > worst_here { + worst_here = error; + worst_spec = to_spec_string(spec); + } + } + + let bad = worst_here > tolerance::(len); + if bad { + failures += 1; + } + println!( + "{:>8} {} candidates, worst relative error {:.3e} (budget {:.1e}) {}{}", + len, + specs.len(), + worst_here, + tolerance::(len), + if bad { "FAIL " } else { "" }, + worst_spec + ); + worst_overall = worst_overall.max(worst_here); + } + + println!( + "\nworst relative error over all lengths: {:.3e} ({} lengths failed)", + worst_overall, failures + ); + if failures > 0 { + std::process::exit(1); + } +} + +/// Show how each algorithm's per-element overhead varies with size. +fn cmd_residuals>(lengths: &[usize], opts: &Options) { + let max_len = lengths.iter().copied().max().unwrap_or(1024) * 4; + eprintln!("measuring primitives..."); + let mut model = calibrate_primitives::(opts.rounds, opts.block_ms, max_len); + eprintln!("fitting overheads..."); + let samples = fit_overheads::( + &mut model, + lengths, + opts.rounds, + opts.block_ms, + opts.cap, + opts.bucketed, + ); + + let mut binned: std::collections::BTreeMap< + &'static str, + std::collections::BTreeMap>, + > = Default::default(); + for (spec, measured) in samples.iter() { + if !FIT_ORDER.contains(&spec.kind()) { + continue; + } + if let Some(inner) = model.inner_cost(spec) { + let residual = (measured - inner) / overhead_scale(spec); + binned + .entry(spec.kind()) + .or_default() + .entry(bucket_of(spec)) + .or_default() + .push(residual); + } + } + + println!( + "{:<5} {:>8} {:>8} {:>9} {:>7}", + "kind", "len>=", "median", "p25..p75", "n" + ); + for (kind, buckets) in binned { + for (bucket, values) in buckets { + if values.len() < 3 { + continue; + } + let mut sorted = values.clone(); + sorted.sort_by(|a, b| a.partial_cmp(b).unwrap()); + let n = sorted.len(); + println!( + "{:<5} {:>8} {:>8.3} {:>9} {:>7}", + kind, + 1usize << bucket, + sorted[n / 2], + format!("{:.2}..{:.2}", sorted[n / 4], sorted[(n * 3) / 4]), + n + ); + } + println!(); + } +} + +/// Measure every candidate at every length and write the raw timings to a file. +/// +/// This exists so that a cost model can be scored and refitted offline, against a frozen +/// dataset, instead of needing the machine for every iteration. Rows are +/// `lenspecnspass`, where pass is 1 for the wide sweep and 2 for the careful +/// re-timing the fastest few get. +fn cmd_dump>(lengths: &[usize], opts: &Options) { + use std::io::Write; + const FINALISTS: usize = 8; + + let path = opts.out.clone().unwrap_or_else(|| "dump.tsv".to_string()); + let mut f = std::io::BufWriter::new(std::fs::File::create(&path).expect("cannot create output")); + writeln!(f, "# planner\t{}", P::label()).unwrap(); + writeln!(f, "# rounds\t{}\tblock_ms\t{}\tcap\t{}", opts.rounds, opts.block_ms, opts.cap).unwrap(); + writeln!(f, "len\tspec\tns\tpass\tplanner_pick").unwrap(); + + for &len in lengths { + let mut planner = P::new(); + let picked = to_spec_string(&planner.plan(len)); + let specs = candidates_capped(&mut planner, len, opts.cap); + let mut subjects: Vec> = specs + .iter() + .map(|spec| { + let fft = planner.build(spec, FftDirection::Forward); + Subject::new(to_spec_string(spec), fft, 1) + }) + .collect(); + measure(&mut subjects, opts.rounds, opts.block_ms); + + // Re-time the fastest few for longer. The minimum of many noisy draws is biased low, so + // without this the best candidate looks faster than it is and every regret is flattered. + let mut order: Vec = (0..subjects.len()).collect(); + order.sort_by(|a, b| subjects[*a].best().partial_cmp(&subjects[*b].best()).unwrap()); + let finalists: Vec = order.into_iter().take(FINALISTS).collect(); + let mut finals: Vec> = finalists + .iter() + .map(|&i| { + let fft = planner.build(&specs[i], FftDirection::Forward); + Subject::new(subjects[i].name.clone(), fft, 1) + }) + .collect(); + measure(&mut finals, opts.rounds * 4, opts.block_ms); + + for (i, s) in subjects.iter().enumerate() { + let is_pick = if s.name == picked { "1" } else { "0" }; + writeln!(f, "{}\t{}\t{:.3}\t1\t{}", len, s.name, s.best(), is_pick).unwrap(); + let _ = i; + } + for s in finals.iter() { + let is_pick = if s.name == picked { "1" } else { "0" }; + writeln!(f, "{}\t{}\t{:.3}\t2\t{}", len, s.name, s.best(), is_pick).unwrap(); + } + println!("{:>9} {} candidates", len, subjects.len()); + } + println!("wrote {}", path); +} + +/// Score the counted model against a dump file. No measurement, no planner, no machine. +fn cmd_score(path: &str, opts: &Options) { + use std::collections::BTreeMap; + let text = std::fs::read_to_string(path).expect("cannot read dump"); + + // len -> [(spec, best ns, is planner pick)]. Pass 2 overwrites pass 1. + // + // The inner container must preserve the dump's own order, which is the order + // `candidates_capped` enumerated in. The planner takes the first minimum in that order, so + // replay has to break cost ties the same way or it does not model the planner. Keying by + // spec string instead sorts `gts(b10,b3)` ahead of `gts(b3,b10)`, which is the opposite of + // the enumeration order and silently reverses every width/height tie. + let mut data: BTreeMap> = BTreeMap::new(); + for line in text.lines() { + if line.starts_with('#') || line.starts_with("len\t") { + continue; + } + let f: Vec<&str> = line.split('\t').collect(); + if f.len() < 5 { + continue; + } + let (len, spec, ns, pass, pick) = ( + f[0].parse::().unwrap(), + f[1].to_string(), + f[2].parse::().unwrap(), + f[3].parse::().unwrap(), + f[4] == "1", + ); + let e: &mut Vec<(String, f64, bool)> = data.entry(len).or_default(); + match e.iter_mut().find(|r| r.0 == spec) { + // pass 2 is the careful re-timing, so it wins + Some(row) if pass != 1 => { + row.1 = ns; + row.2 = pick; + } + Some(_) => {} + None => e.push((spec, ns, pick)), + } + } + + // Take the backend from the dump header unless it was given explicitly, so an SSE dump is + // never priced with NEON instruction costs by accident. + let mut params = opts.params; + if !opts.backend_explicit { + if let Some(label) = text + .lines() + .find_map(|l| l.strip_prefix("# planner\t")) + .and_then(counted::Backend::parse) + { + params.backend = label; + } + } + let model = counted::CountedModel::new(params); + println!("params: {:?}", params); + println!( + "{:>9} {:>9} {:>9} {}", + "len", "counted", "planner", "counted model's pick (when it is not the best)" + ); + + let (mut mr, mut pr) = (Vec::new(), Vec::new()); + for (&len, rows) in &data { + let best = rows.iter().map(|v| v.1).fold(f64::INFINITY, f64::min); + let planner_ns = rows.iter().find(|v| v.2).map(|v| v.1); + + let mut scored: Vec<(String, f64, f64)> = Vec::new(); + for (spec_text, ns, _) in rows { + let spec = match parse(spec_text) { + Ok(s) => s, + Err(_) => continue, + }; + if let Some(c) = model.cost(&spec) { + scored.push((spec_text.clone(), c, *ns)); + } + } + let pick = scored + .iter() + .min_by(|a, b| a.1.partial_cmp(&b.1).unwrap()); + let (pick_name, pick_ns) = match pick { + Some((n, _, ns)) => (n.clone(), *ns), + None => { + println!("{:>9} no candidate priced", len); + continue; + } + }; + let m = pick_ns / best; + mr.push(m); + let p = planner_ns.map(|n| n / best); + if let Some(p) = p { + pr.push(p); + } + let shown = if m <= 1.0001 { "= best".to_string() } else { pick_name }; + println!( + "{:>9} {:>8.3}x {:>8} {}", + len, + m, + p.map(|p| format!("{:.3}x", p)).unwrap_or_else(|| "-".into()), + shown + ); + } + + let stat = |v: &mut Vec, label: &str| { + v.sort_by(|a, b| a.partial_cmp(b).unwrap()); + let mean = v.iter().sum::() / v.len() as f64; + let median = v[v.len() / 2]; + let p90 = v[((v.len() as f64 * 0.9) as usize).min(v.len() - 1)]; + let worst = *v.last().unwrap(); + println!( + " {:<9} n={:<4} mean {:.4} median {:.4} p90 {:.4} worst {:.4}", + label, + v.len(), + mean, + median, + p90, + worst + ); + }; + println!("\n--- regret, pick divided by best measured ---"); + stat(&mut mr, "counted"); + stat(&mut pr, "planner"); +} + +/// Print the counted model's cost tree for one recipe, so the recursion can be checked by eye. +fn cmd_explain(spec_text: &str, opts: &Options) { + let spec = parse(spec_text).expect("could not parse spec"); + let model = counted::CountedModel::new(opts.params); + let root_ws = spec.len() as f64; + + fn walk( + model: &counted::CountedModel, + spec: &Spec, + root_ws: f64, + mult: f64, + depth: usize, + out: &mut Vec, + ) { + let own = model.cost(spec).unwrap_or(f64::NAN); + // cost of the children alone, at the multiplicity the parent runs them + let (kids, child_total): (Vec<(&Spec, f64)>, f64) = match spec { + Spec::MixedRadix { left, right, .. } | Spec::GoodThomas { left, right, .. } => { + let v = vec![ + (left.as_ref(), right.len() as f64), + (right.as_ref(), left.len() as f64), + ]; + let t = v + .iter() + .map(|(c, m)| m * model.cost(c).unwrap_or(0.0)) + .sum(); + (v, t) + } + Spec::RadixN { radixes, base } => { + let m = radixes.iter().product::() as f64; + ( + vec![(base.as_ref(), m)], + m * model.cost(base).unwrap_or(0.0), + ) + } + Spec::Radix4 { k, base } => { + let m = (1u64 << (2 * k)) as f64; + ( + vec![(base.as_ref(), m)], + m * model.cost(base).unwrap_or(0.0), + ) + } + Spec::Raders { inner } | Spec::Bluesteins { inner, .. } => { + (vec![(inner.as_ref(), 2.0)], 2.0 * model.cost(inner).unwrap_or(0.0)) + } + _ => (vec![], 0.0), + }; + out.push(format!( + "{:indent$}{:<34} len {:>7} x{:<8.0} cost {:>14.0} own {:>12.0}", + "", + to_spec_string(spec), + spec.len(), + mult, + mult * own, + mult * (own - child_total), + indent = depth * 2 + )); + for (child, m) in kids { + walk(model, child, root_ws, mult * m, depth + 1, out); + } + } + + let mut out = Vec::new(); + walk(&model, &spec, root_ws, 1.0, 0, &mut out); + println!("params: {:?}", opts.params); + println!("{:<36} {:>11} {:<9} {:>19} {:>16}", "recipe", "len", "times", "total cost", "own cost"); + for line in out { + println!("{}", line); + } +} + +/// Print the model's cost for every measured candidate, for offline analysis. +/// +/// Columns: len, spec, measured ns, model cost. Pure replay, no machine needed. +fn cmd_costs(path: &str, opts: &Options) { + use std::collections::BTreeMap; + let text = std::fs::read_to_string(path).expect("cannot read dump"); + let mut params = opts.params; + if !opts.backend_explicit { + if let Some(b) = text + .lines() + .find_map(|l| l.strip_prefix("# planner\t")) + .and_then(counted::Backend::parse) + { + params.backend = b; + } + } + let model = counted::CountedModel::new(params); + + let mut data: BTreeMap> = BTreeMap::new(); + for line in text.lines() { + if line.starts_with('#') || line.starts_with("len\t") { + continue; + } + let f: Vec<&str> = line.split('\t').collect(); + if f.len() < 5 { + continue; + } + let (len, spec, ns, pass, pick) = ( + f[0].parse::().unwrap(), + f[1].to_string(), + f[2].parse::().unwrap(), + f[3].parse::().unwrap(), + f[4] == "1", + ); + let e = data.entry(len).or_default(); + match e.get(&spec) { + Some(_) if pass == 1 => {} + _ => { + e.insert(spec, (ns, pick)); + } + } + } + + println!("len\tspec\tns\tcost\tplanner_pick"); + for (len, rows) in &data { + for (spec_text, (ns, pick)) in rows { + let cost = parse(spec_text) + .ok() + .and_then(|s| model.cost(&s)) + .map(|c| format!("{:.1}", c)) + .unwrap_or_else(|| "NA".into()); + println!( + "{}\t{}\t{:.3}\t{}\t{}", + len, + spec_text, + ns, + cost, + if *pick { 1 } else { 0 } + ); + } + } +} + +/// Compare the plan-time cost of the fixed planner against enumerate-and-price. +/// +/// The fixed planner answers from a few integer operations. A cost model has to enumerate the +/// candidate set and price every member, which is real work the fixed planner never does. This +/// is the one axis where the fixed planner is unambiguously ahead, so it should be measured. +/// Compare the cost of *choosing* a recipe against the cost of *building* it. +/// +/// Planning only produces a `Recipe`. Turning that into an `Arc` is a separate and much +/// larger job: every algorithm precomputes twiddles, and Rader's and Bluestein's both run a full +/// inner FFT inside their constructors. So the question that decides whether an estimating +/// planner is affordable is not how much slower it is than the fixed planner, but how much it +/// adds to plan-plus-build, which is what a caller actually pays before the first transform. +/// How the best recipe changes once construction cost is counted, which is the question a +/// quick-and-dirty planner for one-shot transforms exists to answer. +/// +/// For every candidate this measures build time and execution time, then reports which recipe +/// minimises `build + k * execute` at several values of `k`. If the same recipe wins at every `k` +/// there is nothing for a construction-aware planner to choose, and the idea is dead. Plan time is +/// deliberately excluded: it is the same constant for every candidate at one length, so it cannot +/// change which recipe wins. +fn cmd_crossover>(lengths: &[usize], opts: &Options) { + let model = counted::CountedModel::new(opts.params); + println!( + "{:>7} {:>6} {:>6} {:>11} {:>11} {:>12} {}", + "len", "cands", "k", "build ns", "exec ns", "total ns", "recipe" + ); + + for &len in lengths { + let mut planner = P::new(); + let specs = plan_candidates(&mut planner, len, opts.cap); + if specs.is_empty() { + eprintln!("len {}: no candidates", len); + continue; + } + + let mut subjects: Vec> = specs + .iter() + .map(|spec| { + let fft = planner.build(spec, FftDirection::Forward); + Subject::new(to_spec_string(spec), fft, 1) + }) + .collect(); + measure(&mut subjects, opts.rounds, opts.block_ms); + let exec: Vec = subjects.iter().map(|s| s.best()).collect(); + + // A fresh planner per repetition, so no inner FFT is served from a cache. + let build_reps = if len > 4096 { 5 } else if len > 256 { 20 } else { 100 }; + let build: Vec = specs + .iter() + .map(|spec| { + let t = Instant::now(); + for _ in 0..build_reps { + let mut pl = P::new(); + std::hint::black_box(pl.build(spec, FftDirection::Forward)); + } + t.elapsed().as_secs_f64() * 1e9 / build_reps as f64 + }) + .collect(); + + let pick = |k: f64| -> usize { + (0..specs.len()) + .min_by(|&a, &b| { + (build[a] + k * exec[a]) + .partial_cmp(&(build[b] + k * exec[b])) + .unwrap() + }) + .unwrap() + }; + + let mut first = true; + for k in [1.0, 10.0, 100.0, 1000.0] { + let i = pick(k); + println!( + "{:>7} {:>6} {:>6} {:>11.0} {:>11.1} {:>12.0} {}", + if first { len.to_string() } else { String::new() }, + if first { specs.len().to_string() } else { String::new() }, + k as usize, + build[i], + exec[i], + build[i] + k * exec[i], + subjects[i].name + ); + first = false; + } + + // What the counted model picks, which optimises execution alone. + let mi = specs + .iter() + .enumerate() + .filter_map(|(i, sp)| model.cost(sp).map(|c| (i, c))) + .min_by(|a, b| a.1.partial_cmp(&b.1).unwrap()) + .map(|(i, _)| i) + .unwrap_or(0); + let one = pick(1.0); + println!( + "{:>7} {:>6} {:>6} {:>11.0} {:>11.1} {:>12} {}", + "", "", "model", build[mi], exec[mi], "", subjects[mi].name + ); + // Crossover: how many executions before the model's pick repays its extra build cost. + if one != mi && exec[mi] < exec[one] { + let k = (build[mi] - build[one]) / (exec[one] - exec[mi]); + println!( + "{:>7} {:>6} {:>6} {:>11} {:>11} {:>12} model repays its build cost after {:.0} executions", + "", "", "", "", "", "", k.max(0.0) + ); + } else if one == mi { + println!( + "{:>7} {:>6} {:>6} {:>11} {:>11} {:>12} same recipe at k=1 and by cost model: nothing to choose", + "", "", "", "", "", "" + ); + } + println!(); + } +} + +fn cmd_plantime>(lengths: &[usize], opts: &Options) { + let model = counted::CountedModel::new(opts.params); + println!( + "{:>7} {:>6} {:>12} {:>14} {:>12} {:>9} {:>11}", + "len", "cands", "plan fixed", "plan+price", "build", "build/plan", "extra vs" + ); + println!( + "{:>7} {:>6} {:>12} {:>14} {:>12} {:>9} {:>11}", + "", "", "ns", "ns", "ns", "", "plan+build" + ); + let (mut tot_a, mut tot_b, mut tot_c) = (0.0, 0.0, 0.0); + for &len in lengths { + // fixed planner: design only, with a fresh planner each time so nothing is cached + let reps = 200; + let t0 = Instant::now(); + for _ in 0..reps { + let mut pl = P::new(); + std::hint::black_box(pl.plan(len)); + } + let a = t0.elapsed().as_secs_f64() * 1e9 / reps as f64; + + let mut pl = P::new(); + let n = plan_candidates(&mut pl, len, opts.cap).len(); + let t1 = Instant::now(); + for _ in 0..reps { + let mut pl = P::new(); + let specs = plan_candidates(&mut pl, len, opts.cap); + let best = specs + .iter() + .filter_map(|sp| model.cost(sp).map(|c| (c, sp))) + .min_by(|x, y| x.0.partial_cmp(&y.0).unwrap()); + std::hint::black_box(best); + } + let b = t1.elapsed().as_secs_f64() * 1e9 / reps as f64; + + // Build the recipe the fixed planner chose. Construction allocates and precomputes, so + // it is far slower than planning; use fewer repetitions and a fresh planner each time so + // nothing is served from a cache. + let mut pl = P::new(); + let spec = pl.plan(len); + let build_reps = if len > 4096 { 5 } else if len > 256 { 20 } else { 100 }; + let t2 = Instant::now(); + for _ in 0..build_reps { + let mut pl = P::new(); + std::hint::black_box(pl.build(&spec, FftDirection::Forward)); + } + let c = t2.elapsed().as_secs_f64() * 1e9 / build_reps as f64; + + tot_a += a; + tot_b += b; + tot_c += c; + println!( + "{:>7} {:>6} {:>12.0} {:>14.0} {:>12.0} {:>8.0}x {:>10.1}%", + len, n, a, b, c, c / a, 100.0 * (b - a) / (a + c) + ); + } + println!( + "\n totals: plan fixed {:.0} ns, plan+price {:.0} ns ({:.1}x), build {:.0} ns", + tot_a, tot_b, tot_b / tot_a, tot_c + ); + println!( + " building is {:.0}x planning; the estimating planner adds {:.2}% to plan-plus-build", + tot_c / tot_a, + 100.0 * (tot_b - tot_a) / (tot_a + tot_c) + ); +} + +fn cmd_model>(train: &[usize], test: &[usize], opts: &Options) { + let max_len = test.iter().chain(train.iter()).copied().max().unwrap_or(1024) * 4; + + println!("planner: {}", P::label()); + println!("measuring primitives..."); + let mut model = calibrate_primitives::(opts.rounds, opts.block_ms, max_len); + + println!("fitting overheads on {} training lengths...", train.len()); + fit_overheads::( + &mut model, + train, + opts.rounds, + opts.block_ms, + opts.cap, + opts.bucketed, + ); + println!("\n{}", model.describe()); + + println!( + "{:>9} {:>9} {:>9} {}", + "len", "model", "planner", "model's pick (when it is not the best)" + ); + let mut model_regrets = Vec::new(); + let mut planner_regrets = Vec::new(); + + for &len in test { + let mut planner = P::new(); + let specs = candidates_capped(&mut planner, len, opts.cap); + let mut subjects: Vec> = specs + .iter() + .map(|spec| { + let fft = planner.build(spec, FftDirection::Forward); + Subject::new(to_spec_string(spec), fft, 1) + }) + .collect(); + measure(&mut subjects, opts.rounds, opts.block_ms); + + // Fastest of the first pass, used only to pick which recipes deserve a careful re-timing. + let best_index = subjects + .iter() + .enumerate() + .min_by(|a, b| a.1.best().partial_cmp(&b.1.best()).unwrap()) + .map(|(i, _)| i) + .unwrap(); + + let model_index = specs + .iter() + .enumerate() + .filter_map(|(i, spec)| model.cost(spec).map(|c| (i, c))) + .min_by(|a, b| a.1.partial_cmp(&b.1).unwrap()) + .map(|(i, _)| i); + + let model_spec = match model_index { + Some(i) => subjects[i].name.clone(), + None => "".to_string(), + }; + + // Second pass. The winner of a wide comparison is biased fast: with sub-percent noise, + // the minimum of many draws lands below the true minimum, which puts a floor under any + // regret measured against it. Re-timing just the recipes of interest, for longer, + // removes most of that bias. + let finalists: Vec = { + let mut picked = vec![0usize, best_index]; + if let Some(i) = model_index { + picked.push(i); + } + picked.sort_unstable(); + picked.dedup(); + picked + }; + let mut finals: Vec> = finalists + .iter() + .map(|&i| { + let fft = planner.build(&specs[i], FftDirection::Forward); + Subject::new(to_spec_string(&specs[i]), fft, 1) + }) + .collect(); + measure(&mut finals, opts.rounds * 4, opts.block_ms); + + let time_of = + |index: usize| -> f64 { finals[finalists.iter().position(|&i| i == index).unwrap()].best() }; + let planner_time = time_of(0); + let best_time = finalists + .iter() + .map(|&i| time_of(i)) + .fold(f64::INFINITY, f64::min); + let model_time = match model_index { + Some(i) => time_of(i), + None => f64::NAN, + }; + + let model_regret = model_time / best_time; + let planner_regret = planner_time / best_time; + model_regrets.push(model_regret); + planner_regrets.push(planner_regret); + + println!( + "{:>9} {:>8.3}x {:>8.3}x {}", + len, + model_regret, + planner_regret, + if model_regret <= 1.001 { + "= best".to_string() + } else { + model_spec + } + ); + } + + for (label, mut values) in [("model ", model_regrets), ("planner", planner_regrets)] { + values.sort_by(|a, b| a.partial_cmp(b).unwrap()); + let n = values.len(); + let mean = values.iter().sum::() / n as f64; + println!( + "{}: mean {:.4} median {:.4} p90 {:.4} worst {:.4}", + label, + mean, + values[n / 2], + values[(n * 9) / 10], + values[n - 1] + ); + } +} + +fn cmd_emit>(lengths: &[usize], opts: &Options, element: &str) { + let max_len = lengths.iter().copied().max().unwrap_or(1024) * 4; + eprintln!("measuring primitives..."); + let mut model = calibrate_primitives::(opts.rounds, opts.block_ms, max_len); + eprintln!("fitting overheads on {} lengths...", lengths.len()); + fit_overheads::( + &mut model, + lengths, + opts.rounds, + opts.block_ms, + opts.cap, + opts.bucketed, + ); + print!("{}", emit::emit(&model, element, P::label())); +} + +// --------------------------------------------------------------------------- + +enum Command { + Time(Vec), + Regret(Vec), + Model(Vec, Vec), + Residuals(Vec), + Verify(Vec), + Emit(Vec), + Dump(Vec), + Score(String), + Explain(String), + Costs(String), + Plantime(Vec), + Crossover(Vec), + Sweep(Vec), +} + +fn run>( + command: &Command, + opts: &Options, + element: &str, +) { + match command { + Command::Time(specs) => cmd_time::(specs, opts), + Command::Regret(lengths) => cmd_regret::(lengths, opts), + Command::Sweep(lengths) => cmd_sweep::(lengths, opts), + Command::Model(train, test) => cmd_model::(train, test, opts), + Command::Residuals(lengths) => cmd_residuals::(lengths, opts), + Command::Verify(lengths) => cmd_verify::(lengths, opts), + Command::Emit(lengths) => cmd_emit::(lengths, opts, element), + Command::Dump(lengths) => cmd_dump::(lengths, opts), + Command::Score(path) => cmd_score(path, opts), + Command::Explain(spec) => cmd_explain(spec, opts), + Command::Costs(path) => cmd_costs(path, opts), + Command::Plantime(l) => cmd_plantime::(l, opts), + Command::Crossover(l) => cmd_crossover::(l, opts), + } +} + +fn dispatch(planner: &str, command: &Command, opts: &Options, el: &str) { + match planner { + "scalar" => run::>(command, opts, el), + // The manifest gives rustfft the SIMD feature matching the target, so architecture + // alone decides which of these exists. A tool-crate `feature = ...` cfg would refer to + // the tool's own features and always be false. + #[cfg(target_arch = "aarch64")] + "neon" => run::>(command, opts, el), + #[cfg(target_arch = "x86_64")] + "sse" => run::>(command, opts, el), + #[cfg(target_arch = "wasm32")] + "wasm_simd" => run::>(command, opts, el), + other => { + eprintln!( + "unknown or unavailable planner '{}' on this build; try 'scalar'", + other + ); + std::process::exit(2); + } + } +} + +fn main() { + let args: Vec = std::env::args().skip(1).collect(); + if args.is_empty() { + eprintln!("usage: planner_tuning [options] ARGS..."); + eprintln!("commands: time SPEC... | regret LEN... | model TRAIN... 0 TEST..."); + eprintln!(" residuals LEN... | verify LEN... | emit LEN..."); + eprintln!(" --planner NAME scalar (default), neon, sse"); + eprintln!(" --rounds N timing rounds per subject (default 9)"); + eprintln!(" --block-ms MS wall-clock time per timed block (default 10)"); + eprintln!(" --cap N max candidates per length (default 48)"); + eprintln!(" --f32 measure f32 instead of f64"); + eprintln!(" --permuted-vector charge permuted passes per vector, not per element"); + eprintln!(" --bucketed fit overhead curves over working set, not constants"); + eprintln!(" --verbose for 'regret', list the top candidates per length"); + std::process::exit(2); + } + + let command_name = args[0].clone(); + let mut planner = "scalar".to_string(); + let mut opts = Options { + rounds: 9, + block_ms: 10.0, + cap: 48, + verbose: false, + bucketed: false, + out: None, + params: counted::Params::default(), + backend_explicit: false, + }; + let mut f32_mode = false; + let mut rest: Vec = Vec::new(); + + let mut i = 1; + while i < args.len() { + match args[i].as_str() { + "--planner" => { + i += 1; + planner = args[i].clone(); + } + "--rounds" => { + i += 1; + opts.rounds = args[i].parse().expect("--rounds wants a number"); + } + "--block-ms" => { + i += 1; + opts.block_ms = args[i].parse().expect("--block-ms wants a number"); + } + "--cap" => { + i += 1; + opts.cap = args[i].parse().expect("--cap wants a number"); + } + "--f32" => f32_mode = true, + "--bucketed" => opts.bucketed = true, + "--out" => { + i += 1; + opts.out = Some(args[i].clone()); + } + "--seq-l1" => { i += 1; opts.params.seq[0] = args[i].parse().unwrap(); } + "--seq-l2" => { i += 1; opts.params.seq[1] = args[i].parse().unwrap(); } + "--seq-dram" => { i += 1; opts.params.seq[2] = args[i].parse().unwrap(); } + "--strided" => { i += 1; opts.params.strided_mult = args[i].parse().unwrap(); } + "--permuted" => { i += 1; opts.params.permuted_mult = args[i].parse().unwrap(); } + "--rader-index" => { i += 1; opts.params.rader_index = args[i].parse().unwrap(); } + "--radixn-extra" => { i += 1; opts.params.radixn_extra = args[i].parse().unwrap(); } + "--mul-complex" => { i += 1; opts.params.mul_complex = args[i].parse().unwrap(); } + "--spill" => { i += 1; opts.params.spill = args[i].parse().unwrap(); } + "--general-row" => { i += 1; opts.params.general_row = args[i].parse().unwrap(); } + "--small-row" => { i += 1; opts.params.small_row = args[i].parse().unwrap(); } + "--permuted-vector" => opts.params.permuted_scalar = false, + "--f64" => opts.params.elem = counted::Elem::F64, + "--backend" => { i += 1; opts.params.backend = counted::Backend::parse(&args[i]).expect("--backend wants neon or sse"); opts.backend_explicit = true; } + "--l1-elems" => { i += 1; opts.params.l1_elems = args[i].parse().unwrap(); } + "--l2-elems" => { i += 1; opts.params.l2_elems = args[i].parse().unwrap(); } + "--verbose" => opts.verbose = true, + other => rest.push(other.to_string()), + } + i += 1; + } + + let numbers = |values: &[String]| -> Vec { + values + .iter() + .map(|s| s.parse().expect("lengths must be numbers")) + .collect() + }; + + // `sweep` takes a thousand lengths, so accept "A..B" as well as a list. + let range_or_numbers = |values: &[String]| -> Vec { + let mut out = Vec::new(); + for value in values { + match value.split_once("..") { + Some((lo, hi)) => { + let lo: usize = lo.parse().expect("bad range start"); + let hi: usize = hi.parse().expect("bad range end"); + out.extend(lo..=hi); + } + None => out.push(value.parse().expect("lengths must be numbers")), + } + } + out + }; + + let command = match command_name.as_str() { + "time" => Command::Time(rest.clone()), + "regret" => Command::Regret(numbers(&rest)), + "sweep" => Command::Sweep(range_or_numbers(&rest)), + "residuals" => Command::Residuals(numbers(&rest)), + "verify" => Command::Verify(numbers(&rest)), + "emit" => Command::Emit(numbers(&rest)), + "dump" => Command::Dump(numbers(&rest)), + "score" => Command::Score(rest[0].clone()), + "explain" => Command::Explain(rest[0].clone()), + "costs" => Command::Costs(rest[0].clone()), + "plantime" => Command::Plantime(numbers(&rest)), + "crossover" => Command::Crossover(numbers(&rest)), + "model" => { + let lengths = numbers(&rest); + let split = lengths + .iter() + .position(|&l| l == 0) + .expect("model wants TRAIN... 0 TEST..."); + let (train, test) = lengths.split_at(split); + Command::Model(train.to_vec(), test[1..].to_vec()) + } + other => { + eprintln!("unknown command '{}'", other); + std::process::exit(2); + } + }; + + request_performance_core(); + + if f32_mode { + opts.params.elem = counted::Elem::F32; + dispatch::(&planner, &command, &opts, "f32"); + } else { + dispatch::(&planner, &command, &opts, "f64"); + } +} diff --git a/tools/planner_tuning/src/model.rs b/tools/planner_tuning/src/model.rs new file mode 100644 index 00000000..6246e104 --- /dev/null +++ b/tools/planner_tuning/src/model.rs @@ -0,0 +1,164 @@ +//! A cost model for FFT recipes. +//! +//! Table-driven rather than curve-fitted. Each planner has a finite and small set of primitives +//! (a couple of dozen butterflies, a few dozen valid Radix4 shapes), so their costs are simply +//! measured and stored. That removes the extrapolation error a fitted closed form introduces, +//! which is what made the 2021 scalar attempt mis-rank a direct Radix4 against a split one. +//! +//! Only the composing algorithms need fitted numbers, and each needs one: the cost per element +//! of the transposes and twiddle multiplies they add on top of their inner FFTs. + +use rustfft::tuning::Spec; +use std::collections::HashMap; + +/// What an algorithm's per-element overhead scales with. +/// +/// Usually its own length. Two exceptions: Bluestein's pointwise multiply and zero-padding run +/// over the padded inner length, which can be nearly four times the outer length; and RadixN +/// makes one pass per factor, so its overhead scales with length times the number of levels. +pub fn overhead_scale(spec: &Spec) -> f64 { + match spec { + Spec::Bluesteins { inner, .. } => inner.len() as f64, + Spec::RadixN { radixes, .. } => (spec.len() * radixes.len()) as f64, + other => other.len() as f64, + } +} + +/// Which log2 bucket a spec's overhead belongs in, bucketed on the same quantity the overhead is +/// charged per. +pub fn bucket_of(spec: &Spec) -> u32 { + overhead_scale(spec).log2() as u32 +} + +#[derive(Default, Clone)] +pub struct Model { + /// Measured nanoseconds for one FFT, by butterfly length. + pub butterfly: HashMap, + /// Measured nanoseconds for one FFT, by (base length, k). + pub radix4: HashMap<(usize, u32), f64>, + /// Fitted nanoseconds per element of overhead, by algorithm, as a curve over log2 of the + /// working set. Entries are sorted by bucket. A single entry means a flat constant. + pub overhead: HashMap<&'static str, Vec<(u32, f64)>>, +} + +impl Model { + /// Overhead per element at a working set of `scale` elements, linearly interpolated between + /// measured buckets and clamped outside the measured range. + pub fn overhead_at(&self, kind: &str, scale: f64) -> Option { + let table = self.overhead.get(kind)?; + match table.len() { + 0 => None, + 1 => Some(table[0].1), + _ => { + let x = scale.log2(); + if x <= table[0].0 as f64 { + return Some(table[0].1); + } + if x >= table[table.len() - 1].0 as f64 { + return Some(table[table.len() - 1].1); + } + for pair in table.windows(2) { + let (lo_bucket, lo_value) = pair[0]; + let (hi_bucket, hi_value) = pair[1]; + if x <= hi_bucket as f64 { + let span = (hi_bucket - lo_bucket) as f64; + let t = if span > 0.0 { + (x - lo_bucket as f64) / span + } else { + 0.0 + }; + return Some(lo_value + t * (hi_value - lo_value)); + } + } + Some(table[table.len() - 1].1) + } + } + } + + /// Estimated nanoseconds for one FFT of this recipe. + /// + /// `None` if the recipe needs a primitive that was never measured or an overhead that was + /// never fitted, so a gap is a loud failure rather than a silently wrong ranking. + pub fn cost(&self, spec: &Spec) -> Option { + Some(match spec { + Spec::Dft(n) => { + // Only reached for degenerate sizes; a quadratic keeps it last. + let n = *n as f64; + 100.0 * n * n + } + Spec::Butterfly(len) => *self.butterfly.get(len)?, + Spec::Radix4 { k, base } => *self.radix4.get(&(base.len(), *k))?, + Spec::RadixN { radixes, base } => { + let scale = overhead_scale(spec); + radixes.iter().product::() as f64 * self.cost(base)? + + self.overhead_at("rn", scale)? * scale + } + Spec::MixedRadix { left, right, .. } | Spec::GoodThomas { left, right, .. } => { + let scale = overhead_scale(spec); + right.len() as f64 * self.cost(left)? + + left.len() as f64 * self.cost(right)? + + self.overhead_at(spec.kind(), scale)? * scale + } + // Rader's runs its inner FFT twice per transform, once forward and once to invert + // the convolution, exactly like Bluestein's. + Spec::Raders { inner } => { + let scale = overhead_scale(spec); + 2.0 * self.cost(inner)? + self.overhead_at("rad", scale)? * scale + } + Spec::Bluesteins { inner, .. } => { + let scale = overhead_scale(spec); + 2.0 * self.cost(inner)? + self.overhead_at("bs", scale)? * scale + } + }) + } + + /// The sum of the inner FFT costs only, with no overhead for `spec` itself. Subtracting this + /// from a measurement is what isolates one algorithm's overhead. + pub fn inner_cost(&self, spec: &Spec) -> Option { + Some(match spec { + Spec::MixedRadix { left, right, .. } | Spec::GoodThomas { left, right, .. } => { + right.len() as f64 * self.cost(left)? + left.len() as f64 * self.cost(right)? + } + Spec::RadixN { radixes, base } => { + radixes.iter().product::() as f64 * self.cost(base)? + } + Spec::Raders { inner } => 2.0 * self.cost(inner)?, + Spec::Bluesteins { inner, .. } => 2.0 * self.cost(inner)?, + other => self.cost(other)?, + }) + } + + /// True if every strict descendant is a primitive or an already-fitted kind. + /// + /// Overheads are fitted in dependency order, because the residual for a MixedRadix + /// containing a MixedRadixSmall only means anything once the Small's overhead is known. + pub fn descendants_known(&self, spec: &Spec, fitted: &[&'static str]) -> bool { + spec.children().iter().all(|child| { + let k = child.kind(); + let ok = matches!(k, "butterfly" | "r4" | "dft") || fitted.contains(&k); + ok && self.descendants_known(child, fitted) + }) + } + + pub fn describe(&self) -> String { + let mut kinds: Vec<(&&str, &Vec<(u32, f64)>)> = self.overhead.iter().collect(); + kinds.sort_by_key(|(k, _)| **k); + let mut out = format!( + "{} butterflies, {} radix4 shapes measured\n", + self.butterfly.len(), + self.radix4.len() + ); + out.push_str("overhead, ns per element, by working set:\n"); + for (kind, table) in kinds { + let rendered: Vec = table + .iter() + .map(|(bucket, value)| format!("{}:{:.2}", 1usize << bucket, value)) + .collect(); + out.push_str(&format!(" {:<5} {}\n", kind, rendered.join(" "))); + } + out + } +} + +/// The order overheads must be fitted in, so that each kind's inners are already known. +pub const FIT_ORDER: [&str; 7] = ["mrs", "gts", "mr", "gt", "rn", "rad", "bs"]; From cfb41852a2d4f1c72584775fedfc1bc91bf30f8c Mon Sep 17 00:00:00 2001 From: Henrik Date: Wed, 16 Sep 2026 22:44:48 +0200 Subject: [PATCH 13/22] Estimate the cheapest recipe in the NEON, SSE and wasm_simd planners Each planner now enumerates the recipes that could compute a length, prices them with a cost model read off the source, and keeps the cheapest. Inner FFTs go through the recipe cache, and the best cost per length is cached next to it, so the search recurses over divisors. The fixed planner stays reachable behind the tuning feature for comparison while this is a draft. wasm_simd borrows NEON's instruction counts as a placeholder. --- src/neon/neon_planner.rs | 253 ++++++++++++- src/plan.rs | 6 +- src/simd/mod.rs | 1 + src/simd/simd_estimate.rs | 588 +++++++++++++++++++++++++++++ src/sse/sse_planner.rs | 253 ++++++++++++- src/tuning/adapters.rs | 11 +- src/wasm_simd/wasm_simd_planner.rs | 263 ++++++++++++- 7 files changed, 1325 insertions(+), 50 deletions(-) create mode 100644 src/simd/simd_estimate.rs diff --git a/src/neon/neon_planner.rs b/src/neon/neon_planner.rs index bea1e657..6e5043bd 100644 --- a/src/neon/neon_planner.rs +++ b/src/neon/neon_planner.rs @@ -18,6 +18,7 @@ use crate::neon::neon_radixn::*; use crate::Fft; use crate::math_utils::{PrimeFactor, PrimeFactors}; +use crate::simd::simd_estimate::{self, CostModel, InstructionSet, Shape}; use crate::simd::simd_planner::{self, RadixNPlan}; const MIN_RADIX4_BITS: u32 = 6; // smallest size to consider radix 4 an option is 2^6 = 64 @@ -155,10 +156,21 @@ impl Recipe { /// /// Each FFT instance owns [`Arc`s](std::sync::Arc) to its internal data, rather than borrowing it from the planner, so it's perfectly /// safe to drop the planner after creating Fft instances. +/// +/// For lengths with more than one plausible recipe, the planner estimates the cost of each +/// candidate from instruction counts and a model of memory access, and picks the cheapest. This +/// means planning a new length takes longer than building a fixed recipe would, but the planner +/// caches every length it has planned, including the inner lengths of composite FFTs. pub struct FftPlannerNeon { algorithm_cache: FftCache, recipe_cache: HashMap>, all_butterflies: Box<[usize]>, + // The estimated cost of each length in `recipe_cache`, when estimating. + cost_cache: HashMap, + cost_model: CostModel, + // False plans with the fixed planner, which is kept for comparison while the + // estimating planner is a draft. + estimating: bool, } impl FftPlannerNeon { @@ -205,6 +217,9 @@ impl FftPlannerNeon { algorithm_cache: FftCache::new(), recipe_cache: HashMap::new(), all_butterflies, + cost_cache: HashMap::new(), + cost_model: CostModel::for_type::(InstructionSet::Neon), + estimating: true, }); } } @@ -238,7 +253,7 @@ impl FftPlannerNeon { self.plan_fft(len, FftDirection::Inverse) } - // Make a recipe for a length + // Make a recipe for a length, by estimating or with the fixed planner. pub(crate) fn design_fft_for_len(&mut self, len: usize) -> Arc { if len < 1 { Arc::new(Recipe::Dft(len)) @@ -246,14 +261,208 @@ impl FftPlannerNeon { Arc::clone(&recipe) } else { let factors = PrimeFactors::compute(len); - let recipe = self.design_fft_with_factors(len, factors); + let recipe = if self.estimating { + self.estimate_fft_with_factors(len, factors) + } else { + self.design_fft_with_factors(len, factors) + }; self.recipe_cache.insert(len, Arc::clone(&recipe)); recipe } } + // Price every recipe worth considering for this length, and keep the cheapest. The fixed + // planner's pick is always one of them, and wins ties. + // + // Inner FFTs are planned through `design_fft_for_len`, so each length is estimated once, and + // the chosen recipe's cost is recorded for the larger lengths built on top of it. + fn estimate_fft_with_factors(&mut self, len: usize, factors: PrimeFactors) -> Arc { + let fixed = Self::shape_of(&self.design_fft_with_factors(len, factors.clone())); + let shapes = if simd_estimate::has_choice(len, &self.all_butterflies) { + simd_estimate::candidates( + len, + &factors, + fixed, + &self.all_butterflies, + simd_planner::complex_per_vector::(), + ) + } else { + vec![fixed] + }; + + let mut best: Option<(f64, Shape)> = None; + for shape in shapes { + if let Some(cost) = self.price(&shape) { + if best + .as_ref() + .map_or(true, |(best_cost, _)| cost < *best_cost) + { + best = Some((cost, shape)); + } + } + } + let (cost, shape) = best + .expect("the cost model is missing the counts for one of this planner's butterflies"); + self.cost_cache.insert(len, cost); + self.recipe_for_shape(shape) + } + + // The estimated cost of one FFT of this shape, planning its inner FFTs first. + fn price(&mut self, shape: &Shape) -> Option { + for child_len in shape.child_lens() { + self.design_fft_for_len(child_len); + } + let costs = &self.cost_cache; + self.cost_model.cost(shape, |child_len| costs[&child_len]) + } + + // The top level of a recipe the fixed planner made, so it can be priced like any other. + fn shape_of(recipe: &Recipe) -> Shape { + match recipe { + Recipe::Dft(_) => unreachable!("the planner only uses a Dft for length 0"), + Recipe::Radix4 { k, base_fft } => Shape::Radix4 { + k: *k, + base_len: base_fft.len(), + }, + Recipe::RadixN { factors, base_fft } => Shape::RadixN { + factors: factors.clone(), + base_len: base_fft.len(), + }, + Recipe::MixedRadix { + left_fft, + right_fft, + } => Shape::MixedRadix { + left_len: left_fft.len(), + right_len: right_fft.len(), + small: false, + }, + Recipe::MixedRadixSmall { + left_fft, + right_fft, + } => Shape::MixedRadix { + left_len: left_fft.len(), + right_len: right_fft.len(), + small: true, + }, + Recipe::GoodThomasAlgorithm { + left_fft, + right_fft, + } => Shape::GoodThomas { + left_len: left_fft.len(), + right_len: right_fft.len(), + small: false, + }, + Recipe::GoodThomasAlgorithmSmall { + left_fft, + right_fft, + } => Shape::GoodThomas { + left_len: left_fft.len(), + right_len: right_fft.len(), + small: true, + }, + Recipe::RadersAlgorithm { inner_fft } => Shape::Raders { + len: inner_fft.len() + 1, + }, + Recipe::BluesteinsAlgorithm { len, inner_fft } => Shape::Bluesteins { + len: *len, + inner_len: inner_fft.len(), + }, + butterfly => Shape::Butterfly(butterfly.len()), + } + } + + // Turn a shape into a recipe, with estimated inner FFTs. + fn recipe_for_shape(&mut self, shape: Shape) -> Arc { + Arc::new(match shape { + Shape::Butterfly(len) => { + return self + .design_butterfly_algorithm(len) + .expect("a butterfly shape should have a butterfly") + } + Shape::Radix4 { k, base_len } => Recipe::Radix4 { + k, + base_fft: self.design_fft_for_len(base_len), + }, + Shape::RadixN { factors, base_len } => Recipe::RadixN { + factors, + base_fft: self.design_fft_for_len(base_len), + }, + Shape::MixedRadix { + left_len, + right_len, + small, + } => { + let left_fft = self.design_fft_for_len(left_len); + let right_fft = self.design_fft_for_len(right_len); + if small { + Recipe::MixedRadixSmall { + left_fft, + right_fft, + } + } else { + Recipe::MixedRadix { + left_fft, + right_fft, + } + } + } + Shape::GoodThomas { + left_len, + right_len, + small, + } => { + let left_fft = self.design_fft_for_len(left_len); + let right_fft = self.design_fft_for_len(right_len); + if small { + Recipe::GoodThomasAlgorithmSmall { + left_fft, + right_fft, + } + } else { + Recipe::GoodThomasAlgorithm { + left_fft, + right_fft, + } + } + } + Shape::Raders { len } => Recipe::RadersAlgorithm { + inner_fft: self.design_fft_for_len(len - 1), + }, + Shape::Bluesteins { len, inner_len } => Recipe::BluesteinsAlgorithm { + len, + inner_fft: self.design_fft_for_len(inner_len), + }, + }) + } + + /// Switch between the estimating planner and the fixed planner it replaces, for comparing the + /// two. Clears every cache, so nothing planned one way is reused the other. + #[cfg(any(test, feature = "tuning"))] + pub(crate) fn set_estimating(&mut self, estimating: bool) { + self.estimating = estimating; + self.clear_caches(); + } + + /// Replace the cost model's weights, for fitting them. Clears every cache. + #[cfg(feature = "tuning")] + pub(crate) fn set_cost_model(&mut self, cost_model: CostModel) { + self.cost_model = cost_model; + self.clear_caches(); + } + + #[cfg(any(test, feature = "tuning"))] + fn clear_caches(&mut self) { + self.algorithm_cache = FftCache::new(); + self.recipe_cache.clear(); + self.cost_cache.clear(); + } + // Create the fft from a recipe, take from cache if possible - pub(crate) fn build_fft(&mut self, recipe: &Recipe, direction: FftDirection) -> Arc> { + pub(crate) fn build_fft( + &mut self, + recipe: &Recipe, + direction: FftDirection, + ) -> Arc> { let len = recipe.len(); if let Some(instance) = self.algorithm_cache.get(len, direction) { instance @@ -667,6 +876,24 @@ impl FftPlannerNeon { mod unit_tests { use super::*; + // The recipe tests pin down the fixed planner's decisions. It stays available for comparison + // while the estimating planner is a draft. + fn fixed(mut planner: FftPlannerNeon) -> FftPlannerNeon { + planner.set_estimating(false); + planner + } + + #[test] + fn test_estimated_recipes_have_the_planned_length() { + // Checks the whole recursion, including Bluestein's inner lengths above the planned one. + let mut planner32 = FftPlannerNeon::::new().unwrap(); + let mut planner64 = FftPlannerNeon::::new().unwrap(); + for len in 0..2000 { + assert_eq!(planner32.design_fft_for_len(len).len(), len); + assert_eq!(planner64.design_fft_for_len(len).len(), len); + } + } + fn is_mixedradix(plan: &Recipe) -> bool { match plan { &Recipe::MixedRadix { .. } => true, @@ -712,7 +939,7 @@ mod unit_tests { #[test] fn test_plan_neon_trivial() { // Length 0 and 1 should use Dft - let mut planner = FftPlannerNeon::::new().unwrap(); + let mut planner = fixed(FftPlannerNeon::::new().unwrap()); for len in 0..1 { let plan = planner.design_fft_for_len(len); assert_eq!(*plan, Recipe::Dft(len)); @@ -723,7 +950,7 @@ mod unit_tests { #[test] fn test_plan_neon_largepoweroftwo() { // Powers of 2 above 6 should use Radix4 - let mut planner = FftPlannerNeon::::new().unwrap(); + let mut planner = fixed(FftPlannerNeon::::new().unwrap()); for pow in 6..32 { let len = 1 << pow; let plan = planner.design_fft_for_len(len); @@ -735,7 +962,7 @@ mod unit_tests { #[test] fn test_plan_neon_butterflies() { // Check that all butterflies are used - let mut planner = FftPlannerNeon::::new().unwrap(); + let mut planner = fixed(FftPlannerNeon::::new().unwrap()); assert_eq!(*planner.design_fft_for_len(2), Recipe::Butterfly2); assert_eq!(*planner.design_fft_for_len(3), Recipe::Butterfly3); assert_eq!(*planner.design_fft_for_len(4), Recipe::Butterfly4); @@ -761,7 +988,7 @@ mod unit_tests { fn test_plan_neon_mixedradix() { // Products of several primes that are all too big for a RadixN cross-FFT layer should // become MixedRadix - let mut planner = FftPlannerNeon::::new().unwrap(); + let mut planner = fixed(FftPlannerNeon::::new().unwrap()); for len in [11 * 11 * 13, 11 * 13 * 17, 17 * 19 * 23, 11 * 13 * 17 * 19] { let plan = planner.design_fft_for_len(len); assert!(is_mixedradix(&plan), "Expected MixedRadix, got {:?}", plan); @@ -772,7 +999,7 @@ mod unit_tests { #[test] fn test_plan_neon_radixn() { // Products of several small primes should become RadixN - let mut planner = FftPlannerNeon::::new().unwrap(); + let mut planner = fixed(FftPlannerNeon::::new().unwrap()); for pow2 in 2..5 { for pow3 in 2..5 { for pow5 in 2..5 { @@ -794,8 +1021,8 @@ mod unit_tests { fn test_plan_neon_radixn_f32_needs_an_even_base() { // An f32 vector holds two complex numbers, so RadixN needs an even column count and can // never take an odd length. Those have to keep falling back to mixed radix. - let mut planner32 = FftPlannerNeon::::new().unwrap(); - let mut planner64 = FftPlannerNeon::::new().unwrap(); + let mut planner32 = fixed(FftPlannerNeon::::new().unwrap()); + let mut planner64 = fixed(FftPlannerNeon::::new().unwrap()); for len in [1215, 10125, 3125] { let plan32 = planner32.design_fft_for_len(len); assert!(!is_radixn(&plan32), "Expected no RadixN, got {:?}", plan32); @@ -810,7 +1037,7 @@ mod unit_tests { #[test] fn test_plan_neon_mixedradixsmall() { // Products of two "small" lengths < 31 that have a common divisor >1, and isn't a power of 2 should be MixedRadixSmall - let mut planner = FftPlannerNeon::::new().unwrap(); + let mut planner = fixed(FftPlannerNeon::::new().unwrap()); for len in [5 * 20, 6 * 9, 12 * 15, 10 * 15].iter() { let plan = planner.design_fft_for_len(*len); assert!( @@ -824,7 +1051,7 @@ mod unit_tests { #[test] fn test_plan_neon_goodthomasbutterfly() { - let mut planner = FftPlannerNeon::::new().unwrap(); + let mut planner = fixed(FftPlannerNeon::::new().unwrap()); for len in [3 * 7, 5 * 7, 11 * 13, 2 * 29].iter() { let plan = planner.design_fft_for_len(*len); assert!( @@ -844,7 +1071,7 @@ mod unit_tests { 181, 191, 193, 197, 199, ]; - let mut planner = FftPlannerNeon::::new().unwrap(); + let mut planner = fixed(FftPlannerNeon::::new().unwrap()); for len in difficultprimes.iter() { let plan = planner.design_fft_for_len(*len); assert!( diff --git a/src/plan.rs b/src/plan.rs index 7055060b..738fd7a5 100644 --- a/src/plan.rs +++ b/src/plan.rs @@ -323,7 +323,11 @@ impl FftPlannerScalar { } // Create the fft from a recipe, take from cache if possible - pub(crate) fn build_fft(&mut self, recipe: &Recipe, direction: FftDirection) -> Arc> { + pub(crate) fn build_fft( + &mut self, + recipe: &Recipe, + direction: FftDirection, + ) -> Arc> { let len = recipe.len(); if let Some(instance) = self.algorithm_cache.get(len, direction) { instance diff --git a/src/simd/mod.rs b/src/simd/mod.rs index 22d7c175..c7a30358 100644 --- a/src/simd/mod.rs +++ b/src/simd/mod.rs @@ -1,5 +1,6 @@ //! Code shared by every SIMD backend, written once against the `SimdVector` trait. +pub mod simd_estimate; pub mod simd_planner; pub mod simd_radixn; pub mod simd_vector; diff --git a/src/simd/simd_estimate.rs b/src/simd/simd_estimate.rs new file mode 100644 index 00000000..0f1fafe5 --- /dev/null +++ b/src/simd/simd_estimate.rs @@ -0,0 +1,588 @@ +//! Estimating planning for the SIMD planners: enumerate the recipes that could compute a length, +//! price each one with a cost model, and keep the cheapest. +//! +//! The cost model is built by reading the source rather than by measuring. Every arithmetic count +//! comes from the butterflies and vector traits of the backend (derived in +//! `tools/planner_tuning/OP-COUNTS.md`), and on top of that sits a coarse memory term: each pass +//! over the buffer is charged per element touched, scaled by how it walks memory. The handful of +//! weights that cannot be read off the source, which set the price of a memory access relative to +//! one arithmetic instruction, were fitted by sweeping lengths 1 to 1000 on several machines with +//! the harness in `tools/planner_tuning`. `COST-MODEL.md` there explains the whole model. +//! +//! Costs are priced on a [`Shape`], which names a recipe's top-level algorithm and the lengths of +//! its inner FFTs, but not the inner recipes themselves. The caller supplies the cost of each +//! inner length. That lets a planner memoise the best cost per length next to its recipe cache, +//! so the search recurses over the divisors of a length rather than over whole recipe trees. +//! Memoising by length alone is sound because a node's cost depends only on its own subtree, +//! never on the transform it is nested in. + +use crate::common::RadixFactor; +use crate::math_utils::PrimeFactors; +use crate::FftNum; + +use super::simd_planner::complex_per_vector; + +/// The most candidates priced at one length. A highly composite length has hundreds of two-way +/// splits; see `cap_candidates` for which ones are dropped. +const MAX_CANDIDATES: usize = 48; + +/// Bases a SIMD `Radix4` can be built on, before the per element type vector pair filter. +const RADIX4_BASES: [usize; 10] = [1, 2, 4, 8, 16, 32, 3, 6, 12, 24]; + +/// Bases a SIMD `RadixN` can be built on, before the per element type vector filter. Wider than +/// the `Radix4` set because `SimdRadixN` needs only a whole number of vectors in the base. +const RADIXN_BASES: [usize; 12] = [4, 5, 6, 7, 8, 9, 10, 12, 15, 16, 24, 32]; + +/// Bluestein's inner lengths are each of these, doubled until at least `2 * len - 1`. +const BLUESTEIN_MULTIPLIERS: [usize; 6] = [1, 3, 5, 7, 9, 15]; + +/// The top level of a recipe, with its inner FFTs given only by length. +#[derive(Debug, Clone, PartialEq)] +pub enum Shape { + Butterfly(usize), + Radix4 { + k: u32, + base_len: usize, + }, + RadixN { + factors: Box<[RadixFactor]>, + base_len: usize, + }, + MixedRadix { + left_len: usize, + right_len: usize, + small: bool, + }, + GoodThomas { + left_len: usize, + right_len: usize, + small: bool, + }, + Raders { + len: usize, + }, + Bluesteins { + len: usize, + inner_len: usize, + }, +} + +impl Shape { + /// The lengths of the inner FFTs, which have to be planned before this shape can be priced. + pub fn child_lens(&self) -> impl Iterator { + let (first, second) = match self { + Shape::Butterfly(_) => (None, None), + Shape::Radix4 { base_len, .. } | Shape::RadixN { base_len, .. } => { + (Some(*base_len), None) + } + Shape::MixedRadix { + left_len, + right_len, + .. + } + | Shape::GoodThomas { + left_len, + right_len, + .. + } => (Some(*left_len), Some(*right_len)), + Shape::Raders { len } => (Some(len - 1), None), + Shape::Bluesteins { inner_len, .. } => (Some(*inner_len), None), + }; + first.into_iter().chain(second) + } +} + +/// Which backend's instruction counts to use. +/// +/// The decomposition each butterfly uses is the same on NEON and SSE, but the cost of the +/// primitives is not: SSE4.1 has no FMA, so `fmadd` is a separate multiply and add, and its +/// complex multiply needs six instructions rather than four. +// Only one backend is compiled for a given target, so the other variant always looks unused. +#[allow(dead_code)] +#[derive(Copy, Clone, Debug, PartialEq)] +pub enum InstructionSet { + Neon, + Sse, +} + +impl InstructionSet { + /// `mul_complex` in the backend's vector trait. + fn mul_complex(self, complex_per_vector: usize) -> f64 { + match (self, complex_per_vector) { + // vcombine + vneg + vmulq_laneq + vfmaq_laneq + (InstructionSet::Neon, 1) => 4.0, + // vtrn1q + vtrn2q + vnegq + vmulq + vrev64q + vfmaq + (InstructionSet::Neon, _) => 6.0, + // unpacklo + unpackhi + 2 mul + shuffle + addsub + (InstructionSet::Sse, _) => 6.0, + } + } + + /// Instructions for one `perform_fft_direct`, excluding the load and store of each element. + /// Hand-counted from the backend's butterflies; the derivation is in `OP-COUNTS.md`. + fn butterfly_compute(self, len: usize, complex_per_vector: usize) -> Option { + // f32 on NEON is counted from `perform_parallel_fft_direct`, which computes two FFTs at + // once, and stored here as the per-FFT figure. + if let (InstructionSet::Neon, 2) = (self, complex_per_vector) { + return Some(match len { + 1 => 0.0, + 2 => 2.0, + 3 => 4.0, + 4 => 5.0, + 5 => 16.0, + 6 => 11.0, + 7 => 19.5, + 8 => 19.0, + 9 => 36.0, + 10 => 37.0, + 11 => 44.5, + 12 => 31.0, + 13 => 61.5, + 15 => 68.0, + 16 => 66.0, + 17 => 100.0, + 19 => 125.5, + 23 => 180.0, + 24 => 113.5, + 29 => 279.5, + 31 => 321.0, + 32 => 178.0, + _ => return None, + }); + } + // The generated prime butterflies come out as a closed form, because the generator is a + // pair of loops. Only the weight of the fmadd chain differs between the backends. + if matches!(len, 7 | 11 | 13 | 17 | 19 | 23 | 29 | 31) { + let h = ((len + 1) / 2) as f64; + return Some(match self { + InstructionSet::Neon => (h - 1.0) * (2.0 * h + 5.0), + InstructionSet::Sse => (h - 1.0) * (4.0 * h + 2.0), + }); + } + let count = match self { + InstructionSet::Neon => match len { + 1 => 0, + 2 => 2, + 3 => 8, + 4 => 10, + 5 => 28, + 6 => 22, + 8 => 38, + 9 => 64, + 10 => 66, + 12 => 62, + 15 => 124, + 16 => 114, + 24 => 201, + 32 => 314, + _ => return None, + }, + // Same decompositions, but bf3 is written without FMA and bf8 reaches for + // rotate_45/rotate_135 where NEON uses explicit multiplies. + InstructionSet::Sse => match len { + 1 => 0, + 2 => 2, + 3 => 10, + 4 => 10, + 5 => 28, + 6 => 26, + 8 => 38, + 9 => 84, + 10 => 66, + 12 => 70, + 15 => 134, + 16 => 122, + 24 => 233, + 32 => 346, + _ => return None, + }, + }; + Some(count as f64) + } +} + +/// How a pass walks memory. +#[derive(Copy, Clone)] +enum Pattern { + /// Contiguous run, one cache line feeding many elements. + Sequential, + /// Fixed stride greater than a line, as in a transpose or a cross-FFT layer. + Strided, + /// Data-dependent scatter or gather: digit reversal, CRT reindexing, Rader's permutation. + Permuted, +} + +/// Prices recipes for one backend and element type, in units of one arithmetic instruction. +/// +/// The public fields are the weights that are not read off the source. They are public so the +/// tuning harness can override them. +#[derive(Copy, Clone, Debug)] +pub struct CostModel { + pub instruction_set: InstructionSet, + /// Complex numbers in one 128-bit vector: 1 for f64, 2 for f32. + pub complex_per_vector: usize, + /// A fixed-stride pass against a sequential one. + pub strided: f64, + /// A gather or scatter against a sequential pass. + pub permuted: f64, + /// One element of Rader's permutation passes, on top of the gather or scatter itself. + /// + /// This was fitted as the latency of a loop-carried modular-multiply chain, which ejmahler#178 + /// has since replaced with a precomputed table. It needs refitting. + pub rader_index: f64, + /// Extra cost per element per cross-FFT layer of the generic `SimdRadixN` driver, over the + /// hand-written `Radix4` kernel doing the same work. Expected near zero on NEON's 32 vector + /// registers and positive on SSE's 16, since a layer holds two rows per radix live at once. + pub radixn_extra: f64, + /// Fixed cost per row charged to the general MixedRadix and GoodThomas, for the blocked + /// transpose and on-the-fly CRT mapping that the `Small` variants do not have. + pub general_row: f64, + /// One outer iteration of `transpose_small`, charged to the `Small` variants for the width + /// by which the pair is wider than it is tall. A tie-break between the two orderings. + pub small_row: f64, +} + +impl CostModel { + /// The fitted weights for this instruction set and element type. + pub fn new(instruction_set: InstructionSet, complex_per_vector: usize) -> Self { + let f32 = complex_per_vector == 2; + let (strided, radixn_extra) = match (instruction_set, f32) { + (InstructionSet::Neon, _) => (1.5, 0.0), + (InstructionSet::Sse, false) => (1.5, 6.0), + (InstructionSet::Sse, true) => (2.5, 1.0), + }; + Self { + instruction_set, + complex_per_vector, + strided, + permuted: 2.5, + rader_index: if f32 { 45.0 } else { 30.0 }, + radixn_extra, + general_row: 30.0, + small_row: 10.0, + } + } + + /// The fitted weights for this instruction set and the element type `T`. + pub fn for_type(instruction_set: InstructionSet) -> Self { + Self::new(instruction_set, complex_per_vector::()) + } + + /// Cost of touching `accesses` elements, counting each load and each store once. + fn mem(&self, accesses: f64, pattern: Pattern) -> f64 { + match pattern { + // A gather or scatter computes an address per complex number, so it cannot fill a + // vector. + Pattern::Permuted => accesses * self.permuted, + Pattern::Strided => accesses / self.complex_per_vector as f64 * self.strided, + Pattern::Sequential => accesses / self.complex_per_vector as f64, + } + } + + fn mul_complex(&self) -> f64 { + self.instruction_set.mul_complex(self.complex_per_vector) + } + + fn butterfly(&self, len: usize) -> Option { + self.instruction_set + .butterfly_compute(len, self.complex_per_vector) + } + + /// Estimated cost of one FFT of this shape. `child_cost` gives the cost of an inner FFT by + /// length, for every length `shape.child_lens()` returns. + /// + /// `None` if a butterfly has no counted entry, so a gap in the tables fails loudly rather than + /// pricing a recipe as free. + pub fn cost(&self, shape: &Shape, child_cost: impl Fn(usize) -> f64) -> Option { + let cpv = self.complex_per_vector as f64; + Some(match shape { + Shape::Butterfly(len) => { + self.butterfly(*len)? + self.mem(2.0 * *len as f64, Pattern::Sequential) + } + Shape::Radix4 { k, base_len } => { + let len = (*base_len << (2 * k)) as f64; + let reps = len / *base_len as f64; + // One digit-reversal transpose, then the base FFTs, then k cross layers of + // len/4 column_butterfly4, each four butterfly2 plus a rotate90 (10 instructions) + // and three twiddle multiplies. + let mut c = self.mem(2.0 * len, Pattern::Permuted); + c += reps * child_cost(*base_len); + c += *k as f64 + * ((len / (4.0 * cpv)) * (10.0 + 3.0 * self.mul_complex()) + + self.mem(2.0 * len, Pattern::Strided)); + c + } + Shape::RadixN { factors, base_len } => { + let radix_product: usize = factors.iter().map(|f| f.radix()).product(); + let len = (*base_len * radix_product) as f64; + let reps = len / *base_len as f64; + let mut c = self.mem(2.0 * len, Pattern::Permuted); + c += reps * child_cost(*base_len); + for factor in factors.iter() { + let radix = factor.radix(); + // The cross-FFT layers call the very same butterfly kernels, so the counted + // table applies directly. Row 0 needs no twiddle, hence radix - 1. + c += (len / radix as f64) + * (self.butterfly(radix)? + (radix as f64 - 1.0) * self.mul_complex()); + c += self.mem(2.0 * len, Pattern::Strided); + c += len * self.radixn_extra; + } + c + } + Shape::MixedRadix { + left_len, + right_len, + small, + } => { + let len = (left_len * right_len) as f64; + // Three transposes, one full twiddle pass, and the two inner dimensions. + // + // `MixedRadixSmall` calls `transpose_small`, whose read index strides by `width` + // and so touches a fresh cache line per element: that is what `Permuted` prices. + // `MixedRadix` hands the job to the `transpose` crate, which tiles the rectangle + // to get cache reuse back and pays `general_row` per row for it. + let mut c = if *small { + 3.0 * self.mem(2.0 * len, Pattern::Permuted) + + self.small_row * left_len.saturating_sub(*right_len) as f64 + } else { + 3.0 * self.mem(2.0 * len, Pattern::Strided) + + self.general_row * 1.5 * (left_len + right_len) as f64 + }; + c += (len / cpv) * self.mul_complex() + self.mem(2.0 * len, Pattern::Sequential); + c += *right_len as f64 * child_cost(*left_len); + c += *left_len as f64 * child_cost(*right_len); + c + } + Shape::GoodThomas { + left_len, + right_len, + small, + } => { + let len = (left_len * right_len) as f64; + // Two CRT reindexing passes and one transpose, but no twiddle multiplies. Both + // reindexing passes are permuted in either variant; the transpose splits the two + // exactly as in MixedRadix. + let mut c = 2.0 * self.mem(2.0 * len, Pattern::Permuted); + c += if *small { + self.mem(2.0 * len, Pattern::Permuted) + + self.small_row * left_len.saturating_sub(*right_len) as f64 + } else { + self.mem(2.0 * len, Pattern::Strided) + + self.general_row * 1.5 * (left_len + right_len) as f64 + }; + c += *right_len as f64 * child_cost(*left_len); + c += *left_len as f64 * child_cost(*right_len); + c + } + Shape::Raders { len } => { + let len_f = *len as f64; + // The inner FFT runs twice, around a permuting gather and scatter. + let mut c = 2.0 * child_cost(len - 1); + c += 2.0 * (self.mem(2.0 * len_f, Pattern::Permuted) + len_f * self.rader_index); + c += len_f * self.mul_complex() + self.mem(2.0 * len_f, Pattern::Sequential); + c + } + Shape::Bluesteins { len, inner_len } => { + let outer = *len as f64; + let inner = *inner_len as f64; + // Inner FFT twice, pointwise multiply over the padded inner length, and a + // twiddle-and-pad pass in and out over the outer length. + let mut c = 2.0 * child_cost(*inner_len); + c += inner * self.mul_complex() + self.mem(2.0 * inner, Pattern::Sequential); + c += + 2.0 * (outer * self.mul_complex() + self.mem(2.0 * outer, Pattern::Sequential)); + c + } + }) + } +} + +/// Whether a length has anything to decide. +/// +/// Two classes do not, and both were checked against measurement rather than assumed. A length +/// with its own butterfly: over lengths 8 to 128 on NEON the bare butterfly is fastest at every +/// one. And a power of two: at every power of two from 64 up, across four datasets, the fixed +/// planner's Radix4 is exactly the fastest measured candidate. Skipping enumeration there matters, +/// because plan time is the largest fraction of plan-plus-build at exactly these short lengths. +pub fn has_choice(len: usize, all_butterflies: &[usize]) -> bool { + !len.is_power_of_two() && all_butterflies.binary_search(&len).is_err() +} + +/// The shapes worth pricing at `len`, with `fixed` (the fixed planner's pick) first. +/// +/// The order is significant: the planner keeps the first of equally cheap shapes, so the fixed +/// planner's pick wins every tie. +/// +/// Each two-way split is offered only with the smaller side as the width. The reverse ordering +/// roughly doubles the candidate count at a highly composite length for almost no information: +/// the smaller-width ordering is the better one in 90 to 97% of measured pairs for the `Small` +/// variants, and the general variants are usually indistinguishable. +pub fn candidates( + len: usize, + factors: &PrimeFactors, + fixed: Shape, + all_butterflies: &[usize], + complex_per_vector: usize, +) -> Vec { + let mut out = vec![fixed]; + let push = |shape: Shape, out: &mut Vec| { + if !out.contains(&shape) { + out.push(shape); + } + }; + + for left_len in 2..=(len / 2) { + if len % left_len != 0 { + continue; + } + let right_len = len / left_len; + if left_len > right_len { + break; + } + let coprime = num_integer::gcd(left_len, right_len) == 1; + let small_allowed = left_len < 33 && right_len < 33; + for small in [false, true] { + if small && !small_allowed { + continue; + } + push( + Shape::MixedRadix { + left_len, + right_len, + small, + }, + &mut out, + ); + if coprime { + push( + Shape::GoodThomas { + left_len, + right_len, + small, + }, + &mut out, + ); + } + } + } + + // Radix4 on every base that leaves a power of four. The kernels need a whole number of + // vector pairs in the base. + for base_len in RADIX4_BASES { + if base_len % (2 * complex_per_vector) != 0 || len % base_len != 0 { + continue; + } + let cross = len / base_len; + if cross.is_power_of_two() && cross.trailing_zeros() % 2 == 0 { + let k = cross.trailing_zeros() / 2; + push(Shape::Radix4 { k, base_len }, &mut out); + } + } + + // RadixN on every base that leaves only radixes a cross layer can take. `SimdRadixN` needs + // only a whole number of vectors in the base. + for base_len in RADIXN_BASES { + if base_len % complex_per_vector != 0 || len % base_len != 0 || len == base_len { + continue; + } + let mut cross = len / base_len; + let mut radixes = Vec::new(); + for (radix, factor) in [ + (7, RadixFactor::Factor7), + (6, RadixFactor::Factor6), + (5, RadixFactor::Factor5), + (4, RadixFactor::Factor4), + (3, RadixFactor::Factor3), + (2, RadixFactor::Factor2), + ] { + while cross % radix == 0 { + cross /= radix; + radixes.push(factor); + } + } + if cross == 1 { + // Benchmarking upstream suggests the 4s want to go last. + radixes.sort_by_key(|factor| *factor == RadixFactor::Factor4); + push( + Shape::RadixN { + factors: radixes.into_boxed_slice(), + base_len, + }, + &mut out, + ); + } + } + + if len > 3 && factors.is_prime() { + push(Shape::Raders { len }, &mut out); + } + + // Bluestein's needs no prime, only an inner FFT of at least 2 * len - 1, so it is offered at + // composite lengths too: 671 = 11 x 61 measured 1.46x faster as Bluestein's than as a split + // around a Rader's for 61. But only where some prime factor has no butterfly of its own. If + // every one has, the whole length decomposes into butterflies, and across four 1..1000 + // sweeps not one such length was won by Bluestein's. + let uncovered_factor = factors + .get_other_factors() + .iter() + .any(|factor| all_butterflies.binary_search(&factor.value).is_err()); + if len > 3 && uncovered_factor { + let min_inner_len = 2 * len - 1; + let mut inner_lens: Vec = BLUESTEIN_MULTIPLIERS + .iter() + .map(|&multiplier| { + let mut inner_len = multiplier; + while inner_len < min_inner_len { + inner_len *= 2; + } + inner_len + }) + .collect(); + inner_lens.sort_unstable(); + inner_lens.dedup(); + for inner_len in inner_lens { + push(Shape::Bluesteins { len, inner_len }, &mut out); + } + } + + cap_candidates(out) +} + +/// Trim a candidate list to `MAX_CANDIDATES`. +/// +/// Everything structural is kept (the fixed planner's pick, Radix4 and RadixN, Rader's and +/// Bluestein's), and only splits are dropped, most lopsided first, on the grounds that a split +/// with a tiny side is mostly its large side plus a transpose. The kept splits move behind the +/// structural shapes, ordered from most to least balanced. +fn cap_candidates(all: Vec) -> Vec { + if all.len() <= MAX_CANDIDATES { + return all; + } + let imbalance = |shape: &Shape| match shape { + Shape::MixedRadix { + left_len, + right_len, + .. + } + | Shape::GoodThomas { + left_len, + right_len, + .. + } => Some(((*right_len as f64).ln() - (*left_len as f64).ln()).abs()), + _ => None, + }; + + let mut kept = Vec::with_capacity(MAX_CANDIDATES); + let mut splits = Vec::new(); + for (index, shape) in all.into_iter().enumerate() { + if index == 0 || imbalance(&shape).is_none() { + kept.push(shape); + } else { + splits.push(shape); + } + } + splits.sort_by(|a, b| imbalance(a).unwrap().total_cmp(&imbalance(b).unwrap())); + let room = MAX_CANDIDATES.saturating_sub(kept.len()); + kept.extend(splits.into_iter().take(room)); + kept +} diff --git a/src/sse/sse_planner.rs b/src/sse/sse_planner.rs index 34e821cd..58e55089 100644 --- a/src/sse/sse_planner.rs +++ b/src/sse/sse_planner.rs @@ -18,6 +18,7 @@ use crate::sse::sse_radixn::*; use crate::Fft; use crate::math_utils::{PrimeFactor, PrimeFactors}; +use crate::simd::simd_estimate::{self, CostModel, InstructionSet, Shape}; use crate::simd::simd_planner::{self, RadixNPlan}; const MIN_RADIX4_BITS: u32 = 6; // smallest size to consider radix 4 an option is 2^6 = 64 @@ -155,10 +156,21 @@ impl Recipe { /// /// Each FFT instance owns [`Arc`s](std::sync::Arc) to its internal data, rather than borrowing it from the planner, so it's perfectly /// safe to drop the planner after creating Fft instances. +/// +/// For lengths with more than one plausible recipe, the planner estimates the cost of each +/// candidate from instruction counts and a model of memory access, and picks the cheapest. This +/// means planning a new length takes longer than building a fixed recipe would, but the planner +/// caches every length it has planned, including the inner lengths of composite FFTs. pub struct FftPlannerSse { algorithm_cache: FftCache, recipe_cache: HashMap>, all_butterflies: Box<[usize]>, + // The estimated cost of each length in `recipe_cache`, when estimating. + cost_cache: HashMap, + cost_model: CostModel, + // False plans with the fixed planner, which is kept for comparison while the + // estimating planner is a draft. + estimating: bool, } impl FftPlannerSse { @@ -205,6 +217,9 @@ impl FftPlannerSse { algorithm_cache: FftCache::new(), recipe_cache: HashMap::new(), all_butterflies, + cost_cache: HashMap::new(), + cost_model: CostModel::for_type::(InstructionSet::Sse), + estimating: true, }); } } @@ -238,7 +253,7 @@ impl FftPlannerSse { self.plan_fft(len, FftDirection::Inverse) } - // Make a recipe for a length + // Make a recipe for a length, by estimating or with the fixed planner. pub(crate) fn design_fft_for_len(&mut self, len: usize) -> Arc { if len < 1 { Arc::new(Recipe::Dft(len)) @@ -246,14 +261,208 @@ impl FftPlannerSse { Arc::clone(&recipe) } else { let factors = PrimeFactors::compute(len); - let recipe = self.design_fft_with_factors(len, factors); + let recipe = if self.estimating { + self.estimate_fft_with_factors(len, factors) + } else { + self.design_fft_with_factors(len, factors) + }; self.recipe_cache.insert(len, Arc::clone(&recipe)); recipe } } + // Price every recipe worth considering for this length, and keep the cheapest. The fixed + // planner's pick is always one of them, and wins ties. + // + // Inner FFTs are planned through `design_fft_for_len`, so each length is estimated once, and + // the chosen recipe's cost is recorded for the larger lengths built on top of it. + fn estimate_fft_with_factors(&mut self, len: usize, factors: PrimeFactors) -> Arc { + let fixed = Self::shape_of(&self.design_fft_with_factors(len, factors.clone())); + let shapes = if simd_estimate::has_choice(len, &self.all_butterflies) { + simd_estimate::candidates( + len, + &factors, + fixed, + &self.all_butterflies, + simd_planner::complex_per_vector::(), + ) + } else { + vec![fixed] + }; + + let mut best: Option<(f64, Shape)> = None; + for shape in shapes { + if let Some(cost) = self.price(&shape) { + if best + .as_ref() + .map_or(true, |(best_cost, _)| cost < *best_cost) + { + best = Some((cost, shape)); + } + } + } + let (cost, shape) = best + .expect("the cost model is missing the counts for one of this planner's butterflies"); + self.cost_cache.insert(len, cost); + self.recipe_for_shape(shape) + } + + // The estimated cost of one FFT of this shape, planning its inner FFTs first. + fn price(&mut self, shape: &Shape) -> Option { + for child_len in shape.child_lens() { + self.design_fft_for_len(child_len); + } + let costs = &self.cost_cache; + self.cost_model.cost(shape, |child_len| costs[&child_len]) + } + + // The top level of a recipe the fixed planner made, so it can be priced like any other. + fn shape_of(recipe: &Recipe) -> Shape { + match recipe { + Recipe::Dft(_) => unreachable!("the planner only uses a Dft for length 0"), + Recipe::Radix4 { k, base_fft } => Shape::Radix4 { + k: *k, + base_len: base_fft.len(), + }, + Recipe::RadixN { factors, base_fft } => Shape::RadixN { + factors: factors.clone(), + base_len: base_fft.len(), + }, + Recipe::MixedRadix { + left_fft, + right_fft, + } => Shape::MixedRadix { + left_len: left_fft.len(), + right_len: right_fft.len(), + small: false, + }, + Recipe::MixedRadixSmall { + left_fft, + right_fft, + } => Shape::MixedRadix { + left_len: left_fft.len(), + right_len: right_fft.len(), + small: true, + }, + Recipe::GoodThomasAlgorithm { + left_fft, + right_fft, + } => Shape::GoodThomas { + left_len: left_fft.len(), + right_len: right_fft.len(), + small: false, + }, + Recipe::GoodThomasAlgorithmSmall { + left_fft, + right_fft, + } => Shape::GoodThomas { + left_len: left_fft.len(), + right_len: right_fft.len(), + small: true, + }, + Recipe::RadersAlgorithm { inner_fft } => Shape::Raders { + len: inner_fft.len() + 1, + }, + Recipe::BluesteinsAlgorithm { len, inner_fft } => Shape::Bluesteins { + len: *len, + inner_len: inner_fft.len(), + }, + butterfly => Shape::Butterfly(butterfly.len()), + } + } + + // Turn a shape into a recipe, with estimated inner FFTs. + fn recipe_for_shape(&mut self, shape: Shape) -> Arc { + Arc::new(match shape { + Shape::Butterfly(len) => { + return self + .design_butterfly_algorithm(len) + .expect("a butterfly shape should have a butterfly") + } + Shape::Radix4 { k, base_len } => Recipe::Radix4 { + k, + base_fft: self.design_fft_for_len(base_len), + }, + Shape::RadixN { factors, base_len } => Recipe::RadixN { + factors, + base_fft: self.design_fft_for_len(base_len), + }, + Shape::MixedRadix { + left_len, + right_len, + small, + } => { + let left_fft = self.design_fft_for_len(left_len); + let right_fft = self.design_fft_for_len(right_len); + if small { + Recipe::MixedRadixSmall { + left_fft, + right_fft, + } + } else { + Recipe::MixedRadix { + left_fft, + right_fft, + } + } + } + Shape::GoodThomas { + left_len, + right_len, + small, + } => { + let left_fft = self.design_fft_for_len(left_len); + let right_fft = self.design_fft_for_len(right_len); + if small { + Recipe::GoodThomasAlgorithmSmall { + left_fft, + right_fft, + } + } else { + Recipe::GoodThomasAlgorithm { + left_fft, + right_fft, + } + } + } + Shape::Raders { len } => Recipe::RadersAlgorithm { + inner_fft: self.design_fft_for_len(len - 1), + }, + Shape::Bluesteins { len, inner_len } => Recipe::BluesteinsAlgorithm { + len, + inner_fft: self.design_fft_for_len(inner_len), + }, + }) + } + + /// Switch between the estimating planner and the fixed planner it replaces, for comparing the + /// two. Clears every cache, so nothing planned one way is reused the other. + #[cfg(any(test, feature = "tuning"))] + pub(crate) fn set_estimating(&mut self, estimating: bool) { + self.estimating = estimating; + self.clear_caches(); + } + + /// Replace the cost model's weights, for fitting them. Clears every cache. + #[cfg(feature = "tuning")] + pub(crate) fn set_cost_model(&mut self, cost_model: CostModel) { + self.cost_model = cost_model; + self.clear_caches(); + } + + #[cfg(any(test, feature = "tuning"))] + fn clear_caches(&mut self) { + self.algorithm_cache = FftCache::new(); + self.recipe_cache.clear(); + self.cost_cache.clear(); + } + // Create the fft from a recipe, take from cache if possible - pub(crate) fn build_fft(&mut self, recipe: &Recipe, direction: FftDirection) -> Arc> { + pub(crate) fn build_fft( + &mut self, + recipe: &Recipe, + direction: FftDirection, + ) -> Arc> { let len = recipe.len(); if let Some(instance) = self.algorithm_cache.get(len, direction) { instance @@ -668,6 +877,24 @@ impl FftPlannerSse { mod unit_tests { use super::*; + // The recipe tests pin down the fixed planner's decisions. It stays available for comparison + // while the estimating planner is a draft. + fn fixed(mut planner: FftPlannerSse) -> FftPlannerSse { + planner.set_estimating(false); + planner + } + + #[test] + fn test_estimated_recipes_have_the_planned_length() { + // Checks the whole recursion, including Bluestein's inner lengths above the planned one. + let mut planner32 = FftPlannerSse::::new().unwrap(); + let mut planner64 = FftPlannerSse::::new().unwrap(); + for len in 0..2000 { + assert_eq!(planner32.design_fft_for_len(len).len(), len); + assert_eq!(planner64.design_fft_for_len(len).len(), len); + } + } + fn is_mixedradix(plan: &Recipe) -> bool { match plan { &Recipe::MixedRadix { .. } => true, @@ -713,7 +940,7 @@ mod unit_tests { #[test] fn test_plan_sse_trivial() { // Length 0 and 1 should use Dft - let mut planner = FftPlannerSse::::new().unwrap(); + let mut planner = fixed(FftPlannerSse::::new().unwrap()); for len in 0..1 { let plan = planner.design_fft_for_len(len); assert_eq!(*plan, Recipe::Dft(len)); @@ -724,7 +951,7 @@ mod unit_tests { #[test] fn test_plan_sse_largepoweroftwo() { // Powers of 2 above 6 should use Radix4 - let mut planner = FftPlannerSse::::new().unwrap(); + let mut planner = fixed(FftPlannerSse::::new().unwrap()); for pow in 6..32 { let len = 1 << pow; let plan = planner.design_fft_for_len(len); @@ -736,7 +963,7 @@ mod unit_tests { #[test] fn test_plan_sse_butterflies() { // Check that all butterflies are used - let mut planner = FftPlannerSse::::new().unwrap(); + let mut planner = fixed(FftPlannerSse::::new().unwrap()); assert_eq!(*planner.design_fft_for_len(2), Recipe::Butterfly2); assert_eq!(*planner.design_fft_for_len(3), Recipe::Butterfly3); assert_eq!(*planner.design_fft_for_len(4), Recipe::Butterfly4); @@ -762,7 +989,7 @@ mod unit_tests { fn test_plan_sse_mixedradix() { // Products of several primes that are all too big for a RadixN cross-FFT layer should // become MixedRadix - let mut planner = FftPlannerSse::::new().unwrap(); + let mut planner = fixed(FftPlannerSse::::new().unwrap()); for len in [11 * 11 * 13, 11 * 13 * 17, 17 * 19 * 23, 11 * 13 * 17 * 19] { let plan = planner.design_fft_for_len(len); assert!(is_mixedradix(&plan), "Expected MixedRadix, got {:?}", plan); @@ -773,7 +1000,7 @@ mod unit_tests { #[test] fn test_plan_sse_radixn() { // Products of several small primes should become RadixN - let mut planner = FftPlannerSse::::new().unwrap(); + let mut planner = fixed(FftPlannerSse::::new().unwrap()); for pow2 in 2..5 { for pow3 in 2..5 { for pow5 in 2..5 { @@ -795,8 +1022,8 @@ mod unit_tests { fn test_plan_sse_radixn_f32_needs_an_even_base() { // An f32 vector holds two complex numbers, so RadixN needs an even column count and can // never take an odd length. Those have to keep falling back to mixed radix. - let mut planner32 = FftPlannerSse::::new().unwrap(); - let mut planner64 = FftPlannerSse::::new().unwrap(); + let mut planner32 = fixed(FftPlannerSse::::new().unwrap()); + let mut planner64 = fixed(FftPlannerSse::::new().unwrap()); for len in [1215, 10125, 3125] { let plan32 = planner32.design_fft_for_len(len); assert!(!is_radixn(&plan32), "Expected no RadixN, got {:?}", plan32); @@ -811,7 +1038,7 @@ mod unit_tests { #[test] fn test_plan_sse_mixedradixsmall() { // Products of two "small" lengths < 31 that have a common divisor >1, and isn't a power of 2 should be MixedRadixSmall - let mut planner = FftPlannerSse::::new().unwrap(); + let mut planner = fixed(FftPlannerSse::::new().unwrap()); for len in [5 * 20, 6 * 9, 12 * 15, 10 * 15].iter() { let plan = planner.design_fft_for_len(*len); assert!( @@ -825,7 +1052,7 @@ mod unit_tests { #[test] fn test_plan_sse_goodthomasbutterfly() { - let mut planner = FftPlannerSse::::new().unwrap(); + let mut planner = fixed(FftPlannerSse::::new().unwrap()); for len in [3 * 7, 5 * 7, 11 * 13, 2 * 29].iter() { let plan = planner.design_fft_for_len(*len); assert!( @@ -845,7 +1072,7 @@ mod unit_tests { 181, 191, 193, 197, 199, ]; - let mut planner = FftPlannerSse::::new().unwrap(); + let mut planner = fixed(FftPlannerSse::::new().unwrap()); for len in difficultprimes.iter() { let plan = planner.design_fft_for_len(*len); assert!( diff --git a/src/tuning/adapters.rs b/src/tuning/adapters.rs index d9e01aba..73643299 100644 --- a/src/tuning/adapters.rs +++ b/src/tuning/adapters.rs @@ -313,9 +313,7 @@ macro_rules! simd_adapter { 5 => crate::common::RadixFactor::Factor5, 6 => crate::common::RadixFactor::Factor6, 7 => crate::common::RadixFactor::Factor7, - other => { - unsupported($label, &format!("a radix of {}", other)) - } + other => unsupported($label, &format!("a radix of {}", other)), }) .collect::>() .into_boxed_slice(), @@ -396,11 +394,10 @@ macro_rules! simd_adapter { } fn build(&mut self, spec: &Spec, direction: FftDirection) -> Arc> { - super::check_unambiguous(spec, &mut Default::default()) - .expect("ambiguous recipe"); + super::check_unambiguous(spec, &mut Default::default()).expect("ambiguous recipe"); let recipe = Self::from_spec(spec); - let mut planner = <$planner>::new() - .expect(concat!("this machine does not support ", $label)); + let mut planner = + <$planner>::new().expect(concat!("this machine does not support ", $label)); planner.build_fft(&recipe, direction) } diff --git a/src/wasm_simd/wasm_simd_planner.rs b/src/wasm_simd/wasm_simd_planner.rs index d2a02e3d..7130f7e9 100644 --- a/src/wasm_simd/wasm_simd_planner.rs +++ b/src/wasm_simd/wasm_simd_planner.rs @@ -6,6 +6,7 @@ use crate::algorithm::{ }; use crate::common::RadixFactor; use crate::math_utils::PrimeFactor; +use crate::simd::simd_estimate::{self, CostModel, InstructionSet, Shape}; use crate::simd::simd_planner::{self, RadixNPlan}; use crate::wasm_simd::*; use crate::{fft_cache::FftCache, math_utils::PrimeFactors, Fft, FftDirection, FftNum}; @@ -146,10 +147,21 @@ impl Recipe { /// /// Each FFT instance owns [`Arc`s](std::sync::Arc) to its internal data, rather than borrowing it from the planner, so it's perfectly /// safe to drop the planner after creating Fft instances. +/// +/// For lengths with more than one plausible recipe, the planner estimates the cost of each +/// candidate from instruction counts and a model of memory access, and picks the cheapest. This +/// means planning a new length takes longer than building a fixed recipe would, but the planner +/// caches every length it has planned, including the inner lengths of composite FFTs. pub struct FftPlannerWasmSimd { algorithm_cache: FftCache, recipe_cache: HashMap>, all_butterflies: Box<[usize]>, + // The estimated cost of each length in `recipe_cache`, when estimating. + cost_cache: HashMap, + cost_model: CostModel, + // False plans with the fixed planner, which is kept for comparison while the + // estimating planner is a draft. + estimating: bool, } impl FftPlannerWasmSimd { /// Creates a new `FftPlannerWasmSimd` instance. @@ -184,6 +196,12 @@ impl FftPlannerWasmSimd { algorithm_cache: FftCache::new(), recipe_cache: HashMap::new(), all_butterflies, + cost_cache: HashMap::new(), + // Placeholder: wasm has no instruction counts of its own yet, so it borrows + // NEON's. A quick test on an M1 showed wasm tracking NEON closely, but that is the + // same machine, and nothing has been counted or fitted for wasm itself. + cost_model: CostModel::for_type::(InstructionSet::Neon), + estimating: true, }) } /// Returns a `Fft` instance which uses WebAssembly SIMD instructions to compute FFTs of size `len`. @@ -210,6 +228,7 @@ impl FftPlannerWasmSimd { } impl FftPlannerWasmSimd { + // Make a recipe for a length, by estimating or with the fixed planner. pub(crate) fn design_fft_for_len(&mut self, len: usize) -> Arc { if len < 1 { Arc::new(Recipe::Dft(len)) @@ -217,13 +236,207 @@ impl FftPlannerWasmSimd { Arc::clone(&recipe) } else { let factors = PrimeFactors::compute(len); - let recipe = self.design_fft_with_factors(len, factors); + let recipe = if self.estimating { + self.estimate_fft_with_factors(len, factors) + } else { + self.design_fft_with_factors(len, factors) + }; self.recipe_cache.insert(len, Arc::clone(&recipe)); recipe } } - pub(crate) fn build_fft(&mut self, recipe: &Recipe, direction: FftDirection) -> Arc> { + // Price every recipe worth considering for this length, and keep the cheapest. The fixed + // planner's pick is always one of them, and wins ties. + // + // Inner FFTs are planned through `design_fft_for_len`, so each length is estimated once, and + // the chosen recipe's cost is recorded for the larger lengths built on top of it. + fn estimate_fft_with_factors(&mut self, len: usize, factors: PrimeFactors) -> Arc { + let fixed = Self::shape_of(&self.design_fft_with_factors(len, factors.clone())); + let shapes = if simd_estimate::has_choice(len, &self.all_butterflies) { + simd_estimate::candidates( + len, + &factors, + fixed, + &self.all_butterflies, + simd_planner::complex_per_vector::(), + ) + } else { + vec![fixed] + }; + + let mut best: Option<(f64, Shape)> = None; + for shape in shapes { + if let Some(cost) = self.price(&shape) { + if best + .as_ref() + .map_or(true, |(best_cost, _)| cost < *best_cost) + { + best = Some((cost, shape)); + } + } + } + let (cost, shape) = best + .expect("the cost model is missing the counts for one of this planner's butterflies"); + self.cost_cache.insert(len, cost); + self.recipe_for_shape(shape) + } + + // The estimated cost of one FFT of this shape, planning its inner FFTs first. + fn price(&mut self, shape: &Shape) -> Option { + for child_len in shape.child_lens() { + self.design_fft_for_len(child_len); + } + let costs = &self.cost_cache; + self.cost_model.cost(shape, |child_len| costs[&child_len]) + } + + // The top level of a recipe the fixed planner made, so it can be priced like any other. + fn shape_of(recipe: &Recipe) -> Shape { + match recipe { + Recipe::Dft(_) => unreachable!("the planner only uses a Dft for length 0"), + Recipe::Radix4 { k, base_fft } => Shape::Radix4 { + k: *k, + base_len: base_fft.len(), + }, + Recipe::RadixN { factors, base_fft } => Shape::RadixN { + factors: factors.clone(), + base_len: base_fft.len(), + }, + Recipe::MixedRadix { + left_fft, + right_fft, + } => Shape::MixedRadix { + left_len: left_fft.len(), + right_len: right_fft.len(), + small: false, + }, + Recipe::MixedRadixSmall { + left_fft, + right_fft, + } => Shape::MixedRadix { + left_len: left_fft.len(), + right_len: right_fft.len(), + small: true, + }, + Recipe::GoodThomasAlgorithm { + left_fft, + right_fft, + } => Shape::GoodThomas { + left_len: left_fft.len(), + right_len: right_fft.len(), + small: false, + }, + Recipe::GoodThomasAlgorithmSmall { + left_fft, + right_fft, + } => Shape::GoodThomas { + left_len: left_fft.len(), + right_len: right_fft.len(), + small: true, + }, + Recipe::RadersAlgorithm { inner_fft } => Shape::Raders { + len: inner_fft.len() + 1, + }, + Recipe::BluesteinsAlgorithm { len, inner_fft } => Shape::Bluesteins { + len: *len, + inner_len: inner_fft.len(), + }, + butterfly => Shape::Butterfly(butterfly.len()), + } + } + + // Turn a shape into a recipe, with estimated inner FFTs. + fn recipe_for_shape(&mut self, shape: Shape) -> Arc { + Arc::new(match shape { + Shape::Butterfly(len) => { + return self + .design_butterfly_algorithm(len) + .expect("a butterfly shape should have a butterfly") + } + Shape::Radix4 { k, base_len } => Recipe::Radix4 { + k, + base_fft: self.design_fft_for_len(base_len), + }, + Shape::RadixN { factors, base_len } => Recipe::RadixN { + factors, + base_fft: self.design_fft_for_len(base_len), + }, + Shape::MixedRadix { + left_len, + right_len, + small, + } => { + let left_fft = self.design_fft_for_len(left_len); + let right_fft = self.design_fft_for_len(right_len); + if small { + Recipe::MixedRadixSmall { + left_fft, + right_fft, + } + } else { + Recipe::MixedRadix { + left_fft, + right_fft, + } + } + } + Shape::GoodThomas { + left_len, + right_len, + small, + } => { + let left_fft = self.design_fft_for_len(left_len); + let right_fft = self.design_fft_for_len(right_len); + if small { + Recipe::GoodThomasAlgorithmSmall { + left_fft, + right_fft, + } + } else { + Recipe::GoodThomasAlgorithm { + left_fft, + right_fft, + } + } + } + Shape::Raders { len } => Recipe::RadersAlgorithm { + inner_fft: self.design_fft_for_len(len - 1), + }, + Shape::Bluesteins { len, inner_len } => Recipe::BluesteinsAlgorithm { + len, + inner_fft: self.design_fft_for_len(inner_len), + }, + }) + } + + /// Switch between the estimating planner and the fixed planner it replaces, for comparing the + /// two. Clears every cache, so nothing planned one way is reused the other. + #[cfg(any(test, feature = "tuning"))] + pub(crate) fn set_estimating(&mut self, estimating: bool) { + self.estimating = estimating; + self.clear_caches(); + } + + /// Replace the cost model's weights, for fitting them. Clears every cache. + #[cfg(feature = "tuning")] + pub(crate) fn set_cost_model(&mut self, cost_model: CostModel) { + self.cost_model = cost_model; + self.clear_caches(); + } + + #[cfg(any(test, feature = "tuning"))] + fn clear_caches(&mut self) { + self.algorithm_cache = FftCache::new(); + self.recipe_cache.clear(); + self.cost_cache.clear(); + } + + pub(crate) fn build_fft( + &mut self, + recipe: &Recipe, + direction: FftDirection, + ) -> Arc> { let len = recipe.len(); if let Some(instance) = self.algorithm_cache.get(len, direction) { instance @@ -635,6 +848,24 @@ impl FftPlannerWasmSimd { #[cfg(test)] mod unit_tests { use super::*; + + // The recipe tests pin down the fixed planner's decisions. It stays available for comparison + // while the estimating planner is a draft. + fn fixed(mut planner: FftPlannerWasmSimd) -> FftPlannerWasmSimd { + planner.set_estimating(false); + planner + } + + #[test] + fn test_estimated_recipes_have_the_planned_length() { + // Checks the whole recursion, including Bluestein's inner lengths above the planned one. + let mut planner32 = FftPlannerWasmSimd::::new().unwrap(); + let mut planner64 = FftPlannerWasmSimd::::new().unwrap(); + for len in 0..2000 { + assert_eq!(planner32.design_fft_for_len(len).len(), len); + assert_eq!(planner64.design_fft_for_len(len).len(), len); + } + } use wasm_bindgen_test::*; fn is_mixedradix(plan: &Recipe) -> bool { @@ -682,7 +913,7 @@ mod unit_tests { #[wasm_bindgen_test] fn test_plan_wasm_simd_trivial() { // Length 0 and 1 should use Dft - let mut planner = FftPlannerWasmSimd::::new().unwrap(); + let mut planner = fixed(FftPlannerWasmSimd::::new().unwrap()); for len in 0..1 { let plan = planner.design_fft_for_len(len); assert_eq!(*plan, Recipe::Dft(len)); @@ -693,7 +924,7 @@ mod unit_tests { #[wasm_bindgen_test] fn test_plan_wasm_simd_largepoweroftwo() { // Powers of 2 above 6 should use Radix4 - let mut planner = FftPlannerWasmSimd::::new().unwrap(); + let mut planner = fixed(FftPlannerWasmSimd::::new().unwrap()); for pow in 6..32 { let len = 1 << pow; let plan = planner.design_fft_for_len(len); @@ -705,7 +936,7 @@ mod unit_tests { #[wasm_bindgen_test] fn test_plan_wasm_simd_butterflies() { // Check that all butterflies are used - let mut planner = FftPlannerWasmSimd::::new().unwrap(); + let mut planner = fixed(FftPlannerWasmSimd::::new().unwrap()); assert_eq!(*planner.design_fft_for_len(2), Recipe::Butterfly2); assert_eq!(*planner.design_fft_for_len(3), Recipe::Butterfly3); assert_eq!(*planner.design_fft_for_len(4), Recipe::Butterfly4); @@ -731,7 +962,7 @@ mod unit_tests { fn test_plan_wasm_simd_mixedradix() { // Products of several primes that are all too big for a RadixN cross-FFT layer should // become MixedRadix - let mut planner = FftPlannerWasmSimd::::new().unwrap(); + let mut planner = fixed(FftPlannerWasmSimd::::new().unwrap()); for len in [11 * 11 * 13, 11 * 13 * 17, 17 * 19 * 23, 11 * 13 * 17 * 19] { let plan = planner.design_fft_for_len(len); assert!(is_mixedradix(&plan), "Expected MixedRadix, got {:?}", plan); @@ -742,7 +973,7 @@ mod unit_tests { #[wasm_bindgen_test] fn test_plan_wasm_simd_radixn() { // Products of several small primes should become RadixN - let mut planner = FftPlannerWasmSimd::::new().unwrap(); + let mut planner = fixed(FftPlannerWasmSimd::::new().unwrap()); for pow2 in 2..5 { for pow3 in 2..5 { for pow5 in 2..5 { @@ -764,8 +995,8 @@ mod unit_tests { fn test_plan_wasm_simd_radixn_f32_needs_an_even_base() { // An f32 vector holds two complex numbers, so RadixN needs an even column count and can // never take an odd length. Those have to keep falling back to mixed radix. - let mut planner32 = FftPlannerWasmSimd::::new().unwrap(); - let mut planner64 = FftPlannerWasmSimd::::new().unwrap(); + let mut planner32 = fixed(FftPlannerWasmSimd::::new().unwrap()); + let mut planner64 = fixed(FftPlannerWasmSimd::::new().unwrap()); for len in [1215, 10125, 3125] { let plan32 = planner32.design_fft_for_len(len); assert!(!is_radixn(&plan32), "Expected no RadixN, got {:?}", plan32); @@ -780,7 +1011,7 @@ mod unit_tests { #[wasm_bindgen_test] fn test_plan_wasm_simd_mixedradixsmall() { // Products of two "small" lengths < 31 that have a common divisor >1, and isn't a power of 2 should be MixedRadixSmall - let mut planner = FftPlannerWasmSimd::::new().unwrap(); + let mut planner = fixed(FftPlannerWasmSimd::::new().unwrap()); for len in [5 * 20, 6 * 9, 12 * 15, 10 * 15].iter() { let plan = planner.design_fft_for_len(*len); assert!( @@ -794,7 +1025,7 @@ mod unit_tests { #[wasm_bindgen_test] fn test_plan_wasm_simd_goodthomasbutterfly() { - let mut planner = FftPlannerWasmSimd::::new().unwrap(); + let mut planner = fixed(FftPlannerWasmSimd::::new().unwrap()); for len in [3 * 7, 5 * 7, 11 * 13, 2 * 29].iter() { let plan = planner.design_fft_for_len(*len); assert!( @@ -814,7 +1045,7 @@ mod unit_tests { 181, 191, 193, 197, 199, ]; - let mut planner = FftPlannerWasmSimd::::new().unwrap(); + let mut planner = fixed(FftPlannerWasmSimd::::new().unwrap()); for len in difficultprimes.iter() { let plan = planner.design_fft_for_len(*len); assert!( @@ -835,21 +1066,21 @@ mod unit_tests { fn test_wasm_simd_fft_cache() { { // Check that FFTs are reused if they're both forward - let mut planner = FftPlannerWasmSimd::::new().unwrap(); + let mut planner = fixed(FftPlannerWasmSimd::::new().unwrap()); let fft_a = planner.plan_fft(1234, FftDirection::Forward); let fft_b = planner.plan_fft(1234, FftDirection::Forward); assert!(Arc::ptr_eq(&fft_a, &fft_b), "Existing fft was not reused"); } { // Check that FFTs are reused if they're both inverse - let mut planner = FftPlannerWasmSimd::::new().unwrap(); + let mut planner = fixed(FftPlannerWasmSimd::::new().unwrap()); let fft_a = planner.plan_fft(1234, FftDirection::Inverse); let fft_b = planner.plan_fft(1234, FftDirection::Inverse); assert!(Arc::ptr_eq(&fft_a, &fft_b), "Existing fft was not reused"); } { // Check that FFTs are NOT resued if they don't both have the same direction - let mut planner = FftPlannerWasmSimd::::new().unwrap(); + let mut planner = fixed(FftPlannerWasmSimd::::new().unwrap()); let fft_a = planner.plan_fft(1234, FftDirection::Forward); let fft_b = planner.plan_fft(1234, FftDirection::Inverse); assert!( @@ -862,7 +1093,7 @@ mod unit_tests { #[wasm_bindgen_test] fn test_wasm_simd_recipe_cache() { // Check that all butterflies are used - let mut planner = FftPlannerWasmSimd::::new().unwrap(); + let mut planner = fixed(FftPlannerWasmSimd::::new().unwrap()); let fft_a = planner.design_fft_for_len(1234); let fft_b = planner.design_fft_for_len(1234); assert!( From 6c86535e34d1d89b4dddb63bcd03e2d1816554ef Mon Sep 17 00:00:00 2001 From: Henrik Date: Wed, 16 Sep 2026 22:52:22 +0200 Subject: [PATCH 14/22] Drive the tuning harness from the library's planners and cost model The harness now asks the fixed and the estimating planner for their picks instead of enumerating and pricing on its own, and prices offline dumps with the library's cost model, so there is one copy of it. Adds a survey command that sweeps random lengths and reports percentiles of the runtime change. Drops the measured-table model the counted one replaced. --- src/neon/neon_planner.rs | 6 + src/simd/simd_estimate.rs | 6 +- src/sse/sse_planner.rs | 6 + src/tuning/adapters.rs | 22 +- src/tuning/cost.rs | 70 ++ src/tuning/mod.rs | 157 +-- src/wasm_simd/wasm_simd_planner.rs | 6 + tools/planner_tuning/src/counted.rs | 564 --------- tools/planner_tuning/src/emit.rs | 77 -- tools/planner_tuning/src/main.rs | 1644 +++++++++++---------------- tools/planner_tuning/src/model.rs | 164 --- 11 files changed, 844 insertions(+), 1878 deletions(-) create mode 100644 src/tuning/cost.rs delete mode 100644 tools/planner_tuning/src/counted.rs delete mode 100644 tools/planner_tuning/src/emit.rs delete mode 100644 tools/planner_tuning/src/model.rs diff --git a/src/neon/neon_planner.rs b/src/neon/neon_planner.rs index 6e5043bd..07f1dd8f 100644 --- a/src/neon/neon_planner.rs +++ b/src/neon/neon_planner.rs @@ -443,6 +443,12 @@ impl FftPlannerNeon { self.clear_caches(); } + /// The cost model the estimating planner uses. + #[cfg(feature = "tuning")] + pub(crate) fn cost_model(&self) -> CostModel { + self.cost_model + } + /// Replace the cost model's weights, for fitting them. Clears every cache. #[cfg(feature = "tuning")] pub(crate) fn set_cost_model(&mut self, cost_model: CostModel) { diff --git a/src/simd/simd_estimate.rs b/src/simd/simd_estimate.rs index 0f1fafe5..d2ba74f1 100644 --- a/src/simd/simd_estimate.rs +++ b/src/simd/simd_estimate.rs @@ -38,7 +38,7 @@ const BLUESTEIN_MULTIPLIERS: [usize; 6] = [1, 3, 5, 7, 9, 15]; /// The top level of a recipe, with its inner FFTs given only by length. #[derive(Debug, Clone, PartialEq)] -pub enum Shape { +pub(crate) enum Shape { Butterfly(usize), Radix4 { k: u32, @@ -293,7 +293,7 @@ impl CostModel { /// /// `None` if a butterfly has no counted entry, so a gap in the tables fails loudly rather than /// pricing a recipe as free. - pub fn cost(&self, shape: &Shape, child_cost: impl Fn(usize) -> f64) -> Option { + pub(crate) fn cost(&self, shape: &Shape, child_cost: impl Fn(usize) -> f64) -> Option { let cpv = self.complex_per_vector as f64; Some(match shape { Shape::Butterfly(len) => { @@ -417,7 +417,7 @@ pub fn has_choice(len: usize, all_butterflies: &[usize]) -> bool { /// roughly doubles the candidate count at a highly composite length for almost no information: /// the smaller-width ordering is the better one in 90 to 97% of measured pairs for the `Small` /// variants, and the general variants are usually indistinguishable. -pub fn candidates( +pub(crate) fn candidates( len: usize, factors: &PrimeFactors, fixed: Shape, diff --git a/src/sse/sse_planner.rs b/src/sse/sse_planner.rs index 58e55089..495dc77b 100644 --- a/src/sse/sse_planner.rs +++ b/src/sse/sse_planner.rs @@ -443,6 +443,12 @@ impl FftPlannerSse { self.clear_caches(); } + /// The cost model the estimating planner uses. + #[cfg(feature = "tuning")] + pub(crate) fn cost_model(&self) -> CostModel { + self.cost_model + } + /// Replace the cost model's weights, for fitting them. Clears every cache. #[cfg(feature = "tuning")] pub(crate) fn set_cost_model(&mut self, cost_model: CostModel) { diff --git a/src/tuning/adapters.rs b/src/tuning/adapters.rs index 73643299..26d85119 100644 --- a/src/tuning/adapters.rs +++ b/src/tuning/adapters.rs @@ -236,7 +236,10 @@ const SIMD_RADIXN_BASES: [usize; 12] = [4, 5, 6, 7, 8, 9, 10, 12, 15, 16, 24, 32 macro_rules! simd_adapter { ($tuner:ident, $label:literal, $planner:path, $recipe:path, $prime_lens:path) => { pub struct $tuner { + // The fixed planner, whose picks `plan` reports. planner: $planner, + // The estimating planner, whose picks `estimate` reports. + estimator: $planner, } impl $tuner { @@ -382,8 +385,12 @@ macro_rules! simd_adapter { } fn new() -> Self { + let mut planner = + <$planner>::new().expect(concat!("this machine does not support ", $label)); + planner.set_estimating(false); Self { - planner: <$planner>::new() + planner, + estimator: <$planner>::new() .expect(concat!("this machine does not support ", $label)), } } @@ -393,6 +400,19 @@ macro_rules! simd_adapter { Self::to_spec(&recipe) } + fn estimate(&mut self, len: usize) -> Option> { + let recipe = self.estimator.design_fft_for_len(len); + Some(Self::to_spec(&recipe)) + } + + fn cost_model(&self) -> Option { + Some(self.estimator.cost_model()) + } + + fn set_cost_model(&mut self, cost_model: super::CostModel) { + self.estimator.set_cost_model(cost_model); + } + fn build(&mut self, spec: &Spec, direction: FftDirection) -> Arc> { super::check_unambiguous(spec, &mut Default::default()).expect("ambiguous recipe"); let recipe = Self::from_spec(spec); diff --git a/src/tuning/cost.rs b/src/tuning/cost.rs new file mode 100644 index 00000000..9f122712 --- /dev/null +++ b/src/tuning/cost.rs @@ -0,0 +1,70 @@ +//! Pricing whole [`Spec`] trees with the estimating planner's cost model. +//! +//! The planners price a recipe one level at a time, taking the cost of each inner length from +//! their cache. Measurement tools need the cost of an arbitrary recipe instead, such as every +//! candidate in a dump, so this walks the tree and feeds the model the costs of the actual inner +//! recipes. + +use std::collections::HashMap; + +use super::Spec; +use crate::common::RadixFactor; +use crate::simd::simd_estimate::Shape; + +pub use crate::simd::simd_estimate::{CostModel, InstructionSet}; + +/// Estimated cost of one FFT of `spec`, or `None` if it contains a butterfly the model has no +/// counts for. +pub fn spec_cost(model: &CostModel, spec: &Spec) -> Option { + let shape = match spec { + // Only ever used at length 0 by the planners. Priced as its quadratic work, so that it + // never looks attractive. + Spec::Dft(len) => return Some(100.0 * (*len as f64) * (*len as f64)), + Spec::Butterfly(len) => Shape::Butterfly(*len), + Spec::Radix4 { k, base } => Shape::Radix4 { + k: *k, + base_len: base.len(), + }, + Spec::RadixN { radixes, base } => Shape::RadixN { + factors: radixes + .iter() + .map(|radix| match radix { + 2 => Some(RadixFactor::Factor2), + 3 => Some(RadixFactor::Factor3), + 4 => Some(RadixFactor::Factor4), + 5 => Some(RadixFactor::Factor5), + 6 => Some(RadixFactor::Factor6), + 7 => Some(RadixFactor::Factor7), + _ => None, + }) + .collect::>>()? + .into_boxed_slice(), + base_len: base.len(), + }, + Spec::MixedRadix { left, right, small } => Shape::MixedRadix { + left_len: left.len(), + right_len: right.len(), + small: *small, + }, + Spec::GoodThomas { left, right, small } => Shape::GoodThomas { + left_len: left.len(), + right_len: right.len(), + small: *small, + }, + Spec::Raders { inner } => Shape::Raders { + len: inner.len() + 1, + }, + Spec::Bluesteins { len, inner } => Shape::Bluesteins { + len: *len, + inner_len: inner.len(), + }, + }; + + // A spec's children have distinct lengths, except a square split, whose two sides are the + // same recipe anyway. + let mut child_costs = HashMap::new(); + for child in spec.children() { + child_costs.insert(child.len(), spec_cost(model, child)?); + } + model.cost(&shape, |child_len| child_costs[&child_len]) +} diff --git a/src/tuning/mod.rs b/src/tuning/mod.rs index 4acc1164..4e1555a9 100644 --- a/src/tuning/mod.rs +++ b/src/tuning/mod.rs @@ -20,6 +20,19 @@ use crate::{Fft, FftDirection, FftNum}; mod adapters; pub use adapters::*; +#[cfg(any( + all(target_arch = "aarch64", feature = "neon"), + all(target_arch = "x86_64", feature = "sse"), + all(target_arch = "wasm32", feature = "wasm_simd"), +))] +mod cost; +#[cfg(any( + all(target_arch = "aarch64", feature = "neon"), + all(target_arch = "x86_64", feature = "sse"), + all(target_arch = "wasm32", feature = "wasm_simd"), +))] +pub use cost::{spec_cost, CostModel, InstructionSet}; + // --------------------------------------------------------------------------- // Planner-independent recipe description // --------------------------------------------------------------------------- @@ -288,6 +301,31 @@ pub trait TunablePlanner: Sized { fn has_radixn() -> bool { false } + + /// The recipe the estimating planner picks for `len`, or `None` for a planner that does not + /// estimate. [`plan`](Self::plan) is always the fixed planner's pick. + fn estimate(&mut self, _len: usize) -> Option> { + None + } + + /// The estimating planner's current cost model, or `None` for a planner that does not + /// estimate. + #[cfg(any( + all(target_arch = "aarch64", feature = "neon"), + all(target_arch = "x86_64", feature = "sse"), + all(target_arch = "wasm32", feature = "wasm_simd"), + ))] + fn cost_model(&self) -> Option { + None + } + + /// Replace the estimating planner's cost model, for fitting its weights. + #[cfg(any( + all(target_arch = "aarch64", feature = "neon"), + all(target_arch = "x86_64", feature = "sse"), + all(target_arch = "wasm32", feature = "wasm_simd"), + ))] + fn set_cost_model(&mut self, _cost_model: CostModel) {} } // --------------------------------------------------------------------------- @@ -317,30 +355,20 @@ pub fn check_unambiguous(spec: &Spec, seen: &mut HashMap) -> Resu Ok(()) } -/// Plausible alternatives to the planner's choice for `len`, the planner's own pick first. +/// Plausible alternatives to the planner's choice for `len`, the fixed planner's pick first. /// -/// Deliberately broader than what any planner would consider: the point is to find what the best -/// available recipe actually is, so a planner's pick can be scored against it. -pub fn candidates>(planner: &mut P, len: usize) -> Vec> { - candidates_inner(planner, len, false) -} - -/// `candidates`, optionally skipping the wider-of-the-two-first ordering of each split. +/// Deliberately broader than what the estimating planner considers: every split in both orders, +/// and Bluestein's at every length. The point is to find what the best available recipe actually +/// is, so a planner's pick can be scored against it, and it is how each of the estimating +/// planner's shortcuts was justified. Re-justifying them after a kernel change needs the same +/// breadth. /// -/// The prune has to happen here rather than as a filter afterwards. Generating a candidate costs -/// far more than pricing one: it renders the spec to a string, scans the seen-list linearly, and -/// walks the tree to check it is unambiguous. Filtering after the fact at length 1200 cut the -/// candidate count from 48 to 32 but plan time only from 211us to 181us; skipping the work up -/// front is what actually saves it. -fn candidates_inner>( - planner: &mut P, - len: usize, - planning: bool, -) -> Vec> { +/// Inner recipes come from the fixed planner, so this searches one level deep. +pub fn candidates>(planner: &mut P, len: usize) -> Vec> { let mut out: Vec> = vec![planner.plan(len)]; let mut seen: Vec = vec![to_spec_string(&out[0])]; - let mut push = |spec: Arc, out: &mut Vec>, seen: &mut Vec| { + let push = |spec: Arc, out: &mut Vec>, seen: &mut Vec| { if spec.len() != len { return; } @@ -359,24 +387,15 @@ fn candidates_inner>( continue; } let right_len = len / left_len; - if planning && left_len > right_len { - continue; - } let left = planner.plan(left_len); let right = planner.plan(right_len); let coprime = gcd(left_len, right_len) == 1; let small = left_len < 33 && right_len < 33; - // Both orderings, unless pruning. The loop above already restricts `left_len` to the - // smaller half when pruning, so the surviving order is the smaller-width-first one. - let orders: Vec<(Arc, Arc)> = if planning { - vec![(Arc::clone(&left), Arc::clone(&right))] - } else { - vec![ - (Arc::clone(&left), Arc::clone(&right)), - (Arc::clone(&right), Arc::clone(&left)), - ] - }; + let orders = [ + (Arc::clone(&left), Arc::clone(&right)), + (Arc::clone(&right), Arc::clone(&left)), + ]; for (left, right) in orders { for small_flag in if small { vec![false, true] @@ -483,33 +502,7 @@ fn candidates_inner>( // pattern is a composite whose factorisation forces Rader's onto a large prime factor: // Rader's permutation work is scalar, so at f32 it does not shrink while everything around // it does. - // When planning, offer Bluestein's only where the direct route can actually be bad: some - // prime factor with no butterfly of its own, which is what forces Rader's or an awkward - // split. If every prime factor has a butterfly the decomposition is all butterflies and - // Bluestein's, which needs an inner FFT of at least 2*len - 1, cannot compete. Across the - // four 1..1000 sweeps the model picked Bluestein's at 1315 lengths and **not one** of them - // had all its prime factors covered, so this costs nothing and skips the enumeration at - // every smooth length. `candidates` still offers it everywhere, which is how that was - // checked and how it would be re-checked. - let bluesteins_worth_it = !planning || { - let butterflies = P::butterfly_lens(); - let mut n = len; - let mut uncovered = false; - let mut d = 2; - while d * d <= n { - while n % d == 0 { - uncovered |= !butterflies.contains(&d); - n /= d; - } - d += 1; - } - if n > 1 { - uncovered |= !butterflies.contains(&n); - } - uncovered - }; - - if len > 3 && bluesteins_worth_it { + if len > 3 { let min_inner = 2 * len - 1; let mut inner_lens: Vec = Vec::new(); for multiplier in [1usize, 3, 5, 7, 9, 15] { @@ -547,52 +540,6 @@ fn candidates_inner>( /// planner's pick, Radix4 and RadixN shapes, Rader's and Bluestein's) and only splits are /// dropped, most lopsided first, on the grounds that a split with a tiny side is mostly just its /// large side plus a transpose. -/// What an estimating planner would actually enumerate at `len`. -/// -/// Identical to `candidates_capped`, except that it returns the fixed planner's pick immediately -/// at lengths where there is nothing to decide. Enumeration is pure overhead there, and it is the -/// overhead that matters most: plan-time is a large fraction of plan-plus-build at small lengths -/// and a negligible one at large lengths, so the cheapest lengths are exactly where an estimating -/// planner can least afford to enumerate. -/// -/// Two classes need no decision, and both were checked against measurement rather than assumed: -/// -/// - **A length with its own butterfly.** A single hand-written kernel beats any decomposition. -/// Over lengths 8..128 on NEON the bare butterfly is fastest at all fifteen such lengths, and -/// is never beaten by a split. -/// - **A power of two.** Radix4 wins, and the fixed planner already picks the right base, which -/// is the part that is not obvious: at 1024 `r4(3,b16)` beats `r4(4,b4)` by 1.15x. Across four -/// datasets, at every power of two from 64 up the fixed planner's pick is exactly the fastest -/// measured candidate, regret 1.000. -/// -/// It also drops the wider-of-the-two-first ordering of every split. Each two-way split is -/// otherwise enumerated twice, which roughly doubles the candidate count at a highly composite -/// length for almost no information: the two orderings differ only in how `transpose_small` walks -/// the rectangle and in which inner FFT runs first. The smaller-width ordering is the better one -/// in 90 to 97% of measured pairs for the Small variants, and for the general variants the two -/// are usually indistinguishable, which makes dropping one free rather than merely cheap. -/// -/// Measured over four datasets, that keeps 58 to 63% of candidates for a geometric mean regret of -/// 1.0016 or better against the full set. The worst single case is 1.129x at length 62 on NEON -/// f32, where `gts(b31,b2)` beats `gts(b2,b31)`; both known exceptions involve b31 or b32, where -/// the parallel-pair f32 butterflies make the chunk count matter in a way none of this models. -/// The planner's own pick is always element zero, so pruning can never leave an estimating -/// planner worse than the fixed one. -/// -/// `candidates` and `candidates_capped` stay exhaustive, because scoring a planner's pick needs -/// the alternatives even where a planner would not look at them. That is how every claim above -/// was established, and re-establishing them after a kernel change needs the same breadth. -pub fn plan_candidates>( - planner: &mut P, - len: usize, - cap: usize, -) -> Vec> { - if len.is_power_of_two() || P::butterfly_lens().contains(&len) { - return vec![planner.plan(len)]; - } - cap_list(candidates_inner(planner, len, true), cap) -} - pub fn candidates_capped>( planner: &mut P, len: usize, diff --git a/src/wasm_simd/wasm_simd_planner.rs b/src/wasm_simd/wasm_simd_planner.rs index 7130f7e9..6a15badd 100644 --- a/src/wasm_simd/wasm_simd_planner.rs +++ b/src/wasm_simd/wasm_simd_planner.rs @@ -418,6 +418,12 @@ impl FftPlannerWasmSimd { self.clear_caches(); } + /// The cost model the estimating planner uses. + #[cfg(feature = "tuning")] + pub(crate) fn cost_model(&self) -> CostModel { + self.cost_model + } + /// Replace the cost model's weights, for fitting them. Clears every cache. #[cfg(feature = "tuning")] pub(crate) fn set_cost_model(&mut self, cost_model: CostModel) { diff --git a/tools/planner_tuning/src/counted.rs b/tools/planner_tuning/src/counted.rs deleted file mode 100644 index f09ae927..00000000 --- a/tools/planner_tuning/src/counted.rs +++ /dev/null @@ -1,564 +0,0 @@ -//! A cost model built by counting instructions in the source, not by measuring. -//! -//! Every leaf cost here comes from reading `src/neon/*.rs`; the derivation is written up in -//! `OP-COUNTS.md`. On top of the arithmetic count sits a coarse memory term: each pass over the -//! buffer is charged per element touched, scaled by how it walks memory (sequential, strided, or -//! permuted) and by which level of an *assumed* cache hierarchy the working set lands in. -//! -//! The point of the exercise is that nothing in here needs a machine. The only quantities that -//! are not read off the source are the handful of weights in `Params`, which set the price of a -//! memory access relative to one arithmetic instruction. - -use rustfft::tuning::Spec; - -/// Which backend's kernels to price. The decomposition each butterfly uses is the same on both, -/// but the instruction cost of the primitives is not: SSE4.1 has no FMA, so `fmadd` is a separate -/// multiply and add, and its complex multiply needs six instructions rather than four. -#[derive(Copy, Clone, Debug, PartialEq)] -pub enum Backend { - Neon, - Sse, -} - -/// Which element type. One 128-bit vector holds one complex f64 or two complex f32, so this -/// changes both the instruction counts and how many elements each memory access covers. -#[derive(Copy, Clone, Debug, PartialEq)] -pub enum Elem { - F32, - F64, -} - -impl Elem { - /// Complex numbers per 128-bit vector. - pub fn complex_per_vector(&self) -> f64 { - match self { - Elem::F32 => 2.0, - Elem::F64 => 1.0, - } - } -} - -impl Backend { - pub fn parse(name: &str) -> Option { - match name { - "neon" => Some(Backend::Neon), - "sse" => Some(Backend::Sse), - _ => None, - } - } - - /// `NeonVector::mul_complex` is 4 instructions; `SseVector::mul_complex` is 6 - /// (unpacklo, unpackhi, two muls, shuffle, addsub). - pub fn mul_complex(&self, elem: Elem) -> f64 { - match (self, elem) { - // vcombine + vneg + vmulq_laneq + vfmaq_laneq - (Backend::Neon, Elem::F64) => 4.0, - // vtrn1q + vtrn2q + vnegq + vmulq + vrev64q + vfmaq - (Backend::Neon, Elem::F32) => 6.0, - // unpacklo + unpackhi + 2 mul + shuffle + addsub - (Backend::Sse, _) => 6.0, - } - } - - /// `column_butterfly4` is four `column_butterfly2` plus one `apply_rotate90` on both. - pub fn column_butterfly4(&self) -> f64 { - 10.0 - } - - /// Architectural vector registers: 32 `v` registers on aarch64, 16 `xmm` on x86-64 SSE. - /// - /// This matters because `cross_layer` in `src/simd_radixn.rs` gathers **two** vector columns - /// before transforming either, so a radix-R layer holds 2R rows live at once, plus the - /// butterfly's own temporaries. At radix 7 that is 14 rows before temporaries, which fits - /// comfortably in 32 registers and not at all in 16. - pub fn registers(&self) -> f64 { - match self { - Backend::Neon => 32.0, - Backend::Sse => 16.0, - } - } - - /// Instructions for one `perform_fft_direct`, excluding the load and store of each element. - /// Hand-counted from `src/neon/neon_butterflies.rs` and `src/sse/sse_butterflies.rs`; the - /// derivation is in `OP-COUNTS.md`. - pub fn butterfly_compute(&self, len: usize, elem: Elem) -> Option { - // The generated prime butterflies come out as a closed form, because the generator is a - // pair of loops. Only the weight of the fmadd chain differs between the backends. - // f32 on NEON is counted from `perform_parallel_fft_direct`, which computes two FFTs - // at once, and stored here as the per-FFT figure. See OP-COUNTS.md. - if let (Backend::Neon, Elem::F32) = (self, elem) { - return Some(match len { - 1 => 0.0, 2 => 2.0, 3 => 4.0, 4 => 5.0, 5 => 16.0, 6 => 11.0, 7 => 19.5, - 8 => 19.0, 9 => 36.0, 10 => 37.0, 11 => 44.5, 12 => 31.0, 13 => 61.5, - 15 => 68.0, 16 => 66.0, 17 => 100.0, 19 => 125.5, 23 => 180.0, 24 => 113.5, - 29 => 279.5, 31 => 321.0, 32 => 178.0, - _ => return None, - }); - } - if matches!(len, 7 | 11 | 13 | 17 | 19 | 23 | 29 | 31) { - let h = ((len + 1) / 2) as f64; - return Some(match self { - Backend::Neon => (h - 1.0) * (2.0 * h + 5.0), - Backend::Sse => (h - 1.0) * (4.0 * h + 2.0), - }); - } - let v = match self { - Backend::Neon => match len { - 1 => 0, - 2 => 2, - 3 => 8, - 4 => 10, - 5 => 28, - 6 => 22, - 8 => 38, - 9 => 64, - 10 => 66, - 12 => 62, - 15 => 124, - 16 => 114, - 24 => 201, - 32 => 314, - _ => return None, - }, - // Same decompositions, but bf3 is written without FMA (10, not a re-weighted 8) and - // bf8 reaches for rotate_45/rotate_135 where NEON uses explicit multiplies. - Backend::Sse => match len { - 1 => 0, - 2 => 2, - 3 => 10, - 4 => 10, - 5 => 28, - 6 => 26, - 8 => 38, - 9 => 84, - 10 => 66, - 12 => 70, - 15 => 134, - 16 => 122, - 24 => 233, - 32 => 346, - _ => return None, - }, - }; - Some(v as f64) - } -} - -/// How a pass walks memory. -#[derive(Copy, Clone)] -pub enum Pattern { - /// Contiguous run, one cache line feeding many elements. - Sequential, - /// Fixed stride greater than a line, as in a transpose or a cross-FFT layer. - Strided, - /// Data-dependent scatter or gather: digit reversal, CRT reindexing, Rader's permutation. - Permuted, -} - -/// The weights that are not read off the source. -/// -/// Costs are in units of one arithmetic instruction. `seq` is the price of a single load or -/// store at each level of the hierarchy; the multipliers raise that for less friendly patterns. -#[derive(Copy, Clone, Debug)] -pub struct Params { - /// Complex numbers that fit in the first level of cache. - pub l1_elems: f64, - /// Complex numbers that fit in the last level of cache. - pub l2_elems: f64, - /// Cost of one load or store at L1, L2 and memory. - pub seq: [f64; 3], - pub strided_mult: f64, - pub permuted_mult: f64, - /// Cost of computing one Rader's permutation index, in arithmetic-instruction equivalents. - /// - /// `raders_algorithm.rs` recomputes `index = index * root % len` per element, where `len` is - /// a `StrengthReducedU64`, so that `%` is two 64x64->128 widening multiplies plus a shift and - /// a subtract: about 7 instructions on aarch64. The instruction count alone understates it, - /// because `index` feeds the next iteration, making the chain loop-carried and latency-bound - /// rather than throughput-bound. This weight converts the counted instructions into the - /// effective cost of that serial chain. - /// - /// **Size it from the latency, not the instruction count.** The carried chain is - /// `mul -> umulh -> mul -> sub`, which is about 10 cycles on both an M1 firestorm core and - /// Coffee Lake. Ten cycles of a core that retires 4 to 6 instructions per cycle is 40 to 60 - /// instruction slots, not the 7 instructions counted. The first default of 20 assumed a 3x - /// latency inflation and got 12 of 54 prime lengths wrong, at up to 1.64x. - /// - /// The term is linear in `len` while everything around it grows as `len log len`, so it - /// matters most at small lengths, which is why the original 33 lengths (all 1000 or larger) - /// could not pin it down. It decides the Rader's-versus-Bluestein's call. - /// - /// **It is not one value per machine, so the default is a compromise across machines.** Slots - /// per cycle depend on the core, and the optimum moves with it: 40 on the M1 and 25 to 30 on a - /// Cortex-A76. The default is chosen by the 1..1000 sweep on the M1, the Pi 5 and the - /// ThinkCentre, to be acceptable on all three rather than optimal on one: - /// - /// - f64: **30**. Fewer losses beyond 2% than 45 on every machine (33/28/46 against - /// 54/116/54), for 0.2% of geometric mean on the M1. - /// - f32: **45**. At or near the best on all three; 30 costs 2 to 4% of geometric mean. - /// - /// Why f32 wants a larger value is not understood. `complex_per_vector` already halves the - /// arithmetic per element while this chain stays per element. - /// - /// Negative means "use that per-element default"; see `CountedModel::rader_index`. - pub rader_index: f64, - /// Cost of one spilled vector per unrolled group in a RadixN cross layer, as a store plus a - /// reload. Zero disables the register-pressure term entirely. - pub spill: f64, - /// Extra cost per element per cross-FFT layer for the **generic** `SimdRadixN` driver, - /// over the hand-written `Radix4` kernel doing the same work. - /// - /// They are different code. `cross_layer` in `src/simd_radixn.rs` is generic over the radix - /// and gathers two vector columns before transforming either, so it holds 2R rows live plus - /// the butterfly's temporaries; `sse_radix4.rs` is a hardcoded 2x unroll over six twiddles. - /// 2R rows fits aarch64's 32 vector registers at every radix it supports, and does not fit - /// x86-64's 16 xmm registers, so this is expected to be near zero on NEON and positive on SSE. - /// - /// The defaults, chosen by the 1..1000 sweep on the ThinkCentre after the transpose indices - /// were precomputed (see `NEXT-STEPS.md`): - /// - /// - NEON: **0**, both element types. - /// - SSE f64: **6**. 38 losses beyond 2% against 50 at 5, 84 at 8 and 180 at 12. The - /// 33-length dump alone points at 12 to 16, but every one of those lengths is 1000 or more. - /// - SSE f32: **1**. 20 losses against 25 at both 0 and 2, and a worst case of 0.909 against - /// 0.861 at 2. - /// - /// The SSE values were 5 and 2 before that change. Removing a per-call overhead let the f32 - /// value come down, as expected; the f64 value moved up by one, not down. - /// - /// Negative means "use that per-backend default"; see `CountedModel::radixn_extra`. - pub radixn_extra: f64, - /// Override for the cost of one complex multiply. Negative means "use the backend's counted - /// value". Exists to test whether SSE's shuffle-heavy `mul_complex` costs more than its - /// instruction count suggests: three of its six instructions are shuffle-class, and on Intel - /// those all issue to a single port, whereas NEON spreads them over symmetric pipes. - pub mul_complex: f64, - /// Charge permuted passes per complex number rather than per vector. On by default. - /// - /// A gather or scatter moves one complex number at a time: digit reversal, CRT reindexing and - /// Rader's permutation all compute a destination per element, so there is no contiguous run to - /// fill a vector with. Dividing their access count by `complex_per_vector` therefore - /// under-charges them by exactly that factor, which is invisible at f64 (where the factor is - /// 1) and a factor of 2 at f32. - /// - /// Setting this false restores the original behaviour, which is what `--permuted-vector` is - /// for. That costs nothing at f64, where both f64 datasets score byte-identically either way, - /// and at f32 it forces the fitted `permuted_mult` up from 2.5 to between 4.0 and 6.0 to - /// absorb the same factor. See the f32-on-SSE section of `RESULTS.md`. - pub permuted_scalar: bool, - /// Fixed cost per row of a pass, charged to the **general** MixedRadix and GoodThomas - /// variants and not to their `Small` counterparts. - /// - /// The two pairs move the same data in the same order; they differ in implementation. - /// `MixedRadixSmall` and `GoodThomasAlgorithmSmall` call `array_utils::transpose_small`, - /// a naive strided double loop over the whole rectangle, and read their permutation from a - /// precomputed table. `MixedRadix` and `GoodThomasAlgorithm` call the `transpose` crate's - /// blocked transpose, which keeps both streams cache resident, and `GoodThomasAlgorithm` - /// additionally computes the CRT mapping on the fly with one `StrengthReducedUsize::div_rem` - /// and a branch per row rather than per element. - /// - /// The general form is therefore cheaper per element and dearer per row, which is a - /// crossover, and the measurements are the shape of one. General over small on the M1: - /// - /// ```text - /// len 22 28 45 104 496 992 - /// gt/gts 1.341 1.270 1.182 1.056 1.001 1.027 - /// mr/mrs 1.115 1.059 1.048 1.013 0.931 0.948 - /// ``` - /// - /// The advantage decays towards 1 and MixedRadix crosses under it near len 200. A per-element - /// difference alone could produce neither: it would hold roughly constant in ratio, and it - /// could never change sign. Without this term the model has only per-element costs, so it - /// prices the pair by pattern alone and picks the general form at every length. - /// - /// Every one of the twelve pairs above is called correctly for `general_row` anywhere in - /// 21 to 42; the binding constraints are Good-Thomas at 992 below and MixedRadix at 496 - /// above. 30 sits in the middle of that window. - pub general_row: f64, - /// Cost of one outer-loop iteration of `array_utils::transpose_small`, charged to the - /// **Small** MixedRadix and GoodThomas variants only. - /// - /// This is the term that makes the model prefer one ordering of a factor pair over its - /// reverse. `transpose_small` is a naive double loop: - /// - /// ```text - /// for x in 0..width { for y in 0..height { out[y + x*height] = in[x + y*width] } } - /// ``` - /// - /// The outer loop runs `width` times and the read index strides by `width`, so the cost - /// depends on which dimension is which. The general variants call the `transpose` crate, - /// which tiles the rectangle and so does not care: that contrast is the evidence, because - /// `GoodThomasAlgorithm` and `GoodThomasAlgorithmSmall` perform the *same single transpose - /// in the same orientation* and differ only in the implementation. Over reversed pairs the - /// small form measures smaller-width-faster at 90 to 97% on both machines and both element - /// types, while the general form splits about evenly and its median gap is 0.00 ns. - /// - /// Outer-loop iterations, counted from the source: - /// - /// - `GoodThomasAlgorithmSmall`: one transpose, `(width, height)`, so `width`. - /// - `MixedRadixSmall`: three, `(w,h)`, `(h,w)`, `(w,h)`, so `2*width + height`. - /// - /// Both change by exactly `width - height` when the pair is reversed, which predicts that - /// the two should show the same asymmetry per unit of `w - h` despite having different - /// absolute transpose counts. On SSE they measure 1.48 and 1.47 ns respectively. - /// - /// So the charge is `small_row * max(width - height, 0)`, not the raw iteration count. Both - /// variants differ by exactly `width - height` iterations between the two orderings, so one - /// weight covers both; charging the worse ordering that difference and the better one - /// nothing reproduces it. Only the *difference* is evidenced here, because the absolute - /// level of a Small variant against a general one is what `general_row` already carries, - /// fitted. Charging the difference rather than the count keeps three properties that matter: - /// - /// - a square pair is charged nothing, since there is no ordering to get wrong; - /// - the better ordering keeps exactly the cost it had before this term existed, so - /// `general_row` stays valid and the Small-versus-general balance is untouched; - /// - the cost never goes negative. - /// - /// Charging the raw count instead regressed SSE f64 at length 1215, where the recipe is - /// `mr(b15, mrs(b9,b9))`: the nested square pair was inflated by its 15 repetitions and the - /// whole recipe lost to an `rn(3.3.3.3,b15)` that is 1.21x slower. - /// - /// The value is one outer iteration in instruction-equivalents: about 1.48 ns on the i3 and - /// 0.7 to 1.0 ns on the M1, which at each machine's ns-per-cost-unit is 9 to 13 either way. - /// It barely matters. Scores are byte-identical for anything from 2 to 24 on every dataset, - /// because the term only ever separates two orderings that are otherwise exactly equal in - /// cost. It is a tie-break with a derivation, not a fitted weight. - pub small_row: f64, - /// Which backend's instruction costs to use. - pub backend: Backend, - /// Which element type. - pub elem: Elem, -} - -impl Default for Params { - fn default() -> Self { - // Apple M1 performance core: 128 KiB L1d, 12 MiB L2, at 16 bytes per complex f64. - Self { - l1_elems: 8192.0, - l2_elems: 786432.0, - seq: [1.0, 2.0, 6.0], - strided_mult: 1.5, - permuted_mult: 2.5, - rader_index: -1.0, - spill: 0.0, - radixn_extra: -1.0, - mul_complex: -1.0, - permuted_scalar: true, - general_row: 30.0, - small_row: 10.0, - backend: Backend::Neon, - elem: Elem::F64, - } - } -} - -pub struct CountedModel { - pub params: Params, -} - -impl CountedModel { - pub fn new(params: Params) -> Self { - Self { params } - } - - /// The complex-multiply cost actually in force. - fn mul_complex(&self) -> f64 { - if self.params.mul_complex >= 0.0 { - self.params.mul_complex - } else { - self.params.backend.mul_complex(self.params.elem) - } - } - - /// The Rader's index cost actually in force. See `Params::rader_index` for the values. - fn rader_index(&self) -> f64 { - if self.params.rader_index >= 0.0 { - self.params.rader_index - } else { - match self.params.elem { - Elem::F64 => 30.0, - Elem::F32 => 45.0, - } - } - } - - /// The RadixN driver cost actually in force. See `Params::radixn_extra` for the values. - fn radixn_extra(&self) -> f64 { - if self.params.radixn_extra >= 0.0 { - self.params.radixn_extra - } else { - match (self.params.backend, self.params.elem) { - (Backend::Neon, _) => 0.0, - (Backend::Sse, Elem::F64) => 6.0, - (Backend::Sse, Elem::F32) => 1.0, - } - } - } - - /// Cost of touching `accesses` elements (counting each load and each store once) with the - /// given pattern, when the enclosing buffer holds `ws` complex numbers. - /// Rows walked by the three passes of a width x height decomposition. - /// - /// The passes run over `height`, `width` and `height` rows respectively, so the exact total is - /// `2h + w`. The model deliberately ties `mr(A,B)` with `mr(B,A)`, so use the mean of the two - /// orderings, `1.5 * (w + h)`, rather than introduce an asymmetry here alone. - fn rows(&self, left: &Spec, right: &Spec) -> f64 { - 1.5 * (left.len() as f64 + right.len() as f64) - } - - fn mem(&self, accesses: f64, pattern: Pattern, ws: f64) -> f64 { - let p = &self.params; - // One load or store moves a whole vector, which is one complex f64 or two complex f32, - // except under a permutation, where each element's address is computed separately. - let per_access = match (pattern, p.permuted_scalar) { - (Pattern::Permuted, true) => 1.0, - _ => p.elem.complex_per_vector(), - }; - let accesses = accesses / per_access; - let level = if ws <= p.l1_elems { - 0 - } else if ws <= p.l2_elems { - 1 - } else { - 2 - }; - let mult = match pattern { - Pattern::Sequential => 1.0, - Pattern::Strided => p.strided_mult, - Pattern::Permuted => p.permuted_mult, - }; - accesses * p.seq[level] * mult - } - - /// Estimated cost of one FFT of this recipe, in arithmetic-instruction equivalents. - /// - /// `None` if a butterfly length has no counted entry, so a gap fails loudly. - pub fn cost(&self, spec: &Spec) -> Option { - self.cost_ws(spec, spec.len() as f64) - } - - /// `ws` is the working set of the whole transform, threaded down unchanged: every pass of - /// every nested algorithm walks the same top-level buffer, so that is what decides which - /// cache level the traffic is served from. - fn cost_ws(&self, spec: &Spec, ws: f64) -> Option { - let p = &self.params; - Some(match spec { - Spec::Dft(n) => { - let n = *n as f64; - 100.0 * n * n - } - Spec::Butterfly(len) => { - p.backend.butterfly_compute(*len, p.elem)? + self.mem(2.0 * *len as f64, Pattern::Sequential, ws) - } - Spec::Radix4 { k, base } => { - let len = spec.len() as f64; - let reps = len / base.len() as f64; - // One digit-reversal transpose, then the base FFTs, then k cross layers. - let mut c = self.mem(2.0 * len, Pattern::Permuted, ws); - c += reps * self.cost_ws(base, ws)?; - for _ in 0..*k { - // len/4 column_butterfly4, each with three twiddle multiplies. - c += (len / (4.0 * p.elem.complex_per_vector())) - * (p.backend.column_butterfly4() + 3.0 * self.mul_complex()); - c += self.mem(2.0 * len, Pattern::Strided, ws); - } - c - } - Spec::RadixN { radixes, base } => { - let len = spec.len() as f64; - let reps = len / base.len() as f64; - let mut c = self.mem(2.0 * len, Pattern::Permuted, ws); - c += reps * self.cost_ws(base, ws)?; - for r in radixes.iter() { - let rf = *r as f64; - // The cross-FFT layers call the very same butterfly kernels, so the counted - // table applies directly. Row 0 needs no twiddle, hence r - 1. - c += (len / rf) * (p.backend.butterfly_compute(*r, p.elem)? + (rf - 1.0) * self.mul_complex()); - c += self.mem(2.0 * len, Pattern::Strided, ws); - c += len * self.radixn_extra(); - // Register pressure: the layer keeps 2R rows live across the two-column - // unroll. Anything past the architectural register file becomes a spill and a - // reload, once per element of the group. - if p.spill > 0.0 { - let live = 2.0 * rf; - let over = (live - p.backend.registers()).max(0.0); - c += over * p.spill * (len / rf); - } - } - c - } - Spec::MixedRadix { left, right, small } => { - let len = spec.len() as f64; - // Three transposes, one full twiddle pass, two inner dimensions. - // - // Both variants transpose the same rectangle three times, but not the same way. - // `MixedRadixSmall` calls `transpose_small`, whose read index strides by `width` - // and so touches a fresh cache line per element once `width` exceeds a line: - // line-wasting, which is what `Permuted` prices. `MixedRadix` hands the job to - // the `transpose` crate, which tiles the rectangle to get that reuse back, and - // pays `general_row` per row of setup for it. - let pat = if *small { Pattern::Permuted } else { Pattern::Strided }; - let mut c = 3.0 * self.mem(2.0 * len, pat, ws); - if *small { - // transpose_small at (w,h), (h,w), (w,h) is 2*width + height outer - // iterations; reversing the pair gives 2*height + width, so the two - // orderings differ by width - height. See `small_row`. - c += p.small_row * (left.len() as f64 - right.len() as f64).max(0.0); - } else { - c += p.general_row * self.rows(left, right); - } - c += (len / p.elem.complex_per_vector()) * self.mul_complex() - + self.mem(2.0 * len, Pattern::Sequential, ws); - c += right.len() as f64 * self.cost_ws(left, ws)?; - c += left.len() as f64 * self.cost_ws(right, ws)?; - c - } - Spec::GoodThomas { left, right, small } => { - let len = spec.len() as f64; - // Two CRT reindexing passes and one transpose, but no twiddle multiplies at all: - // dropping them is the whole point of Good-Thomas, and it pays in index work. - // - // Both reindexing passes are `Permuted` in either variant. The small one gathers - // through a precomputed table; the general one walks `destination_index` forward - // by `width + 1` and wraps modulo `len`, which cycles over the whole buffer and - // is no friendlier to a cache than a table would be. The transpose splits the two - // exactly as in MixedRadix, and the general form pays the same per-row setup. - let pat = if *small { Pattern::Permuted } else { Pattern::Strided }; - let mut c = 2.0 * self.mem(2.0 * len, Pattern::Permuted, ws); - c += self.mem(2.0 * len, pat, ws); - if *small { - // One transpose_small at (width, height): `width` outer iterations, against - // `height` reversed. The same width - height difference as MixedRadixSmall, - // which is why one weight serves both. - c += p.small_row * (left.len() as f64 - right.len() as f64).max(0.0); - } else { - c += p.general_row * self.rows(left, right); - } - c += right.len() as f64 * self.cost_ws(left, ws)?; - c += left.len() as f64 * self.cost_ws(right, ws)?; - c - } - Spec::Raders { inner } => { - let len = spec.len() as f64; - // The inner FFT runs twice, and the permutation is precomputed into a u32 table, - // so it is a gather and a scatter rather than a modular multiply per element. - let mut c = 2.0 * self.cost_ws(inner, ws)?; - // Two permutation passes, each a scatter or gather whose index comes from a - // serial modular-multiply chain rather than from a table. - c += 2.0 * (self.mem(2.0 * len, Pattern::Permuted, ws) + len * self.rader_index()); - c += len * self.mul_complex() + self.mem(2.0 * len, Pattern::Sequential, ws); - c - } - Spec::Bluesteins { len, inner } => { - let outer = *len as f64; - let ilen = inner.len() as f64; - // Inner FFT twice, pointwise multiply over the padded inner length, and a - // twiddle-and-pad pass in and out over the outer length. - let mut c = 2.0 * self.cost_ws(inner, ws)?; - c += ilen * self.mul_complex() + self.mem(2.0 * ilen, Pattern::Sequential, ws); - c += 2.0 * (outer * self.mul_complex() + self.mem(2.0 * outer, Pattern::Sequential, ws)); - c - } - }) - } -} diff --git a/tools/planner_tuning/src/emit.rs b/tools/planner_tuning/src/emit.rs deleted file mode 100644 index 848b590a..00000000 --- a/tools/planner_tuning/src/emit.rs +++ /dev/null @@ -1,77 +0,0 @@ -//! Emit a fitted model as a Rust source file for the Neon planner to include. -//! -//! The generated file is data, not logic: measured primitive costs and the fitted per-element -//! overheads. Keeping it separate from the planner keeps the boundary between "measured on some -//! machine" and "written by a person" obvious, which is what ejmahler asked for on PR #38. - -use crate::model::Model; - -/// Costs are stored relative to Butterfly32 rather than in nanoseconds, so that a table fitted -/// on one machine is at least meaningfully comparable on another. The planner only ever -/// compares costs, so the unit cancels. -const REFERENCE_BUTTERFLY: usize = 32; - -pub fn emit(model: &Model, element_type: &str, planner: &str) -> String { - let unit = model - .butterfly - .get(&REFERENCE_BUTTERFLY) - .copied() - .expect("butterfly 32 must be measured, it is the normalisation reference"); - - let mut out = String::new(); - out.push_str(&format!( - "// Generated by tools/planner_tuning. Do not edit by hand.\n\ - //\n\ - // Costs for the {} planner, relative to Butterfly{} for {} elements. Regenerate with:\n\ - // planner_tuning emit --planner {} LENGTHS...\n\n", - planner, REFERENCE_BUTTERFLY, element_type, planner - )); - - let mut butterflies: Vec<(&usize, &f64)> = model.butterfly.iter().collect(); - butterflies.sort_by_key(|(len, _)| **len); - out.push_str("/// Measured cost of one butterfly, by length.\n"); - out.push_str("pub(crate) const BUTTERFLY_COST: &[(usize, f32)] = &[\n"); - for (len, cost) in butterflies { - out.push_str(&format!(" ({}, {:.5}),\n", len, cost / unit)); - } - out.push_str("];\n\n"); - - let mut shapes: Vec<(&(usize, u32), &f64)> = model.radix4.iter().collect(); - shapes.sort_by_key(|((base, k), _)| (*base, *k)); - out.push_str("/// Measured cost of one Radix4, by (base length, k).\n"); - out.push_str("pub(crate) const RADIX4_COST: &[(usize, u32, f32)] = &[\n"); - for ((base, k), cost) in shapes { - out.push_str(&format!(" ({}, {}, {:.5}),\n", base, k, cost / unit)); - } - out.push_str("];\n\n"); - - out.push_str( - "/// Fitted overhead per element, as (log2 of working set, cost), to be interpolated.\n\ - ///\n\ - /// These are curves rather than constants because the scattered reindexing in\n\ - /// GoodThomas and Rader's costs sharply more once the array outgrows L1.\n", - ); - for (name, key) in [ - ("MIXEDRADIX", "mr"), - ("MIXEDRADIX_SMALL", "mrs"), - ("GOODTHOMAS", "gt"), - ("GOODTHOMAS_SMALL", "gts"), - ("RADERS", "rad"), - ("BLUESTEINS", "bs"), - ] { - let table = model - .overhead - .get(key) - .unwrap_or_else(|| panic!("overhead '{}' was never fitted", key)); - let entries: Vec = table - .iter() - .map(|(bucket, value)| format!("({}, {:.6})", bucket, value / unit)) - .collect(); - out.push_str(&format!( - "pub(crate) const {}_OVERHEAD: &[(u32, f32)] = &[{}];\n", - name, - entries.join(", ") - )); - } - out -} diff --git a/tools/planner_tuning/src/main.rs b/tools/planner_tuning/src/main.rs index 5d817137..95e367e0 100644 --- a/tools/planner_tuning/src/main.rs +++ b/tools/planner_tuning/src/main.rs @@ -2,28 +2,28 @@ //! //! Works against any planner that implements `TunablePlanner`, selected with `--planner`. //! Recipes are built through the planner's own internals, so what is measured here is exactly -//! what the planner would construct. +//! what the planner would construct. "planner" below always means the fixed planner the +//! estimating one replaces, and "model" the estimating planner. //! //! Subcommands: -//! time SPEC... time recipes against each other; a '*reps' suffix runs a recipe -//! over a len*reps buffer, which is how inner FFTs are invoked -//! regret LEN... measure how far the planner's pick is from the best available -//! sweep LEN... | A..B time the planner's pick against the counted model's pick, as TSV -//! model TRAIN... 0 TEST... calibrate a cost model and score its picks the same way -//! residuals LEN... show how per-element overhead varies with working set -//! verify LEN... check every enumerated candidate against a direct DFT -//! emit LEN... print the fitted model as a Rust source file - -mod counted; -mod emit; -mod model; - -use model::{bucket_of, overhead_scale, Model, FIT_ORDER}; +//! time SPEC... time recipes against each other; a '*reps' suffix runs a recipe +//! over a len*reps buffer, which is how inner FFTs are invoked +//! regret LEN... how far both planners' picks are from the best candidate measured +//! sweep LEN... | A..B time the fixed planner's pick against the estimating planner's, as TSV +//! survey A..B sweep over `--count` lengths drawn at random from a range +//! verify LEN... check every enumerated candidate against a direct DFT +//! dump LEN... time every candidate and write the raw timings, for offline replay +//! score DUMP replay a dump: regret of the cost model's pick among its candidates +//! costs DUMP replay a dump: the model's cost next to every measured time +//! explain SPEC print the cost model's cost tree for one recipe +//! plantime LEN... plan time of both planners against build time +//! crossover LEN... which recipe wins once construction cost counts + use rustfft::num_complex::Complex; use rustfft::num_traits::{ToPrimitive, Zero}; use rustfft::tuning::{ - candidates_capped, parse, plan_candidates, to_spec_string, ScalarTuner, Spec, - TunablePlanner, + candidates_capped, parse, spec_cost, to_spec_string, CostModel, InstructionSet, ScalarTuner, + Spec, TunablePlanner, }; use rustfft::{Fft, FftDirection, FftNum}; use std::sync::Arc; @@ -125,173 +125,100 @@ fn request_performance_core() { #[cfg(not(target_os = "macos"))] fn request_performance_core() {} -fn median_of(mut values: Vec) -> f64 { - values.sort_by(|a, b| a.partial_cmp(b).unwrap()); - values[values.len() / 2] -} - // --------------------------------------------------------------------------- -// Calibration +// Options // --------------------------------------------------------------------------- -/// Measure every primitive the model needs: each butterfly, and each valid Radix4 shape. -/// -/// Each is timed over a buffer of at least 8192 elements, because primitives are almost always -/// invoked many times over a larger buffer rather than standalone, and a single cold call would -/// price in a startup cost they do not really pay in use. -fn calibrate_primitives>( - rounds: usize, - block_ms: f64, - max_len: usize, -) -> Model { - const REFERENCE_ELEMENTS: usize = 8192; - let mut planner = P::new(); - let mut model = Model::default(); - - let mut shapes: Vec> = Vec::new(); - let mut subjects: Vec> = Vec::new(); - - for len in P::butterfly_lens() { - let spec = Spec::Butterfly(len); - let fft = planner.build(&spec, FftDirection::Forward); - let reps = (REFERENCE_ELEMENTS / len).max(1); - shapes.push(None); - subjects.push(Subject::new(format!("b{}", len), fft, reps)); - } - - for base in P::radix4_bases() { - let mut k = 1u32; - while base * (1usize << (2 * k)) <= max_len { - let spec = Spec::Radix4 { - k, - base: Arc::new(Spec::Butterfly(base)), - }; - let len = spec.len(); - let fft = planner.build(&spec, FftDirection::Forward); - let reps = (REFERENCE_ELEMENTS / len).max(1); - shapes.push(Some((base, k))); - subjects.push(Subject::new(to_spec_string(&spec), fft, reps)); - k += 1; - } - } - - measure(&mut subjects, rounds, block_ms); +/// Overrides for the cost model's fitted weights. Anything left `None` keeps the planner's +/// default for its backend and element type. +#[derive(Default, Debug)] +struct Weights { + strided: Option, + permuted: Option, + rader_index: Option, + radixn_extra: Option, + general_row: Option, + small_row: Option, +} - for (shape, subject) in shapes.iter().zip(subjects.iter()) { - match shape { - Some(key) => { - model.radix4.insert(*key, subject.best()); - } - None => { - model.butterfly.insert(subject.fft.len(), subject.best()); +impl Weights { + fn apply(&self, mut model: CostModel) -> CostModel { + let set = |field: &mut f64, value: Option| { + if let Some(value) = value { + *field = value; } - } + }; + set(&mut model.strided, self.strided); + set(&mut model.permuted, self.permuted); + set(&mut model.rader_index, self.rader_index); + set(&mut model.radixn_extra, self.radixn_extra); + set(&mut model.general_row, self.general_row); + set(&mut model.small_row, self.small_row); + model } - model } -/// Fit one overhead curve for each composing algorithm. -/// -/// Kinds are fitted in dependency order, since the residual of a MixedRadix that contains a -/// MixedRadixSmall only means anything once the Small's own overhead is known. -fn fit_overheads>( - model: &mut Model, - lengths: &[usize], +struct Options { rounds: usize, block_ms: f64, cap: usize, - bucketed: bool, -) -> Vec<(Arc, f64)> { - let mut samples: Vec<(Arc, f64)> = Vec::new(); - for &len in lengths { - let mut planner = P::new(); - let specs = candidates_capped(&mut planner, len, cap); - let mut subjects: Vec> = specs - .iter() - .map(|spec| { - let fft = planner.build(spec, FftDirection::Forward); - Subject::new(to_spec_string(spec), fft, 1) - }) - .collect(); - measure(&mut subjects, rounds, block_ms); - for (spec, subject) in specs.iter().zip(subjects.iter()) { - samples.push((Arc::clone(spec), subject.best())); - } - } + verbose: bool, + /// Where `dump` writes its rows. + out: Option, + weights: Weights, + f32: bool, + /// For `survey`: how many lengths, and the seed that picks them. + count: usize, + seed: u64, +} - let mut fitted: Vec<&'static str> = Vec::new(); - for target in FIT_ORDER { - let usable: Vec<(u32, f64)> = samples - .iter() - .filter(|(spec, _)| spec.kind() == target && model.descendants_known(spec, &fitted)) - .filter_map(|(spec, measured)| { - model - .inner_cost(spec) - .map(|inner| (bucket_of(spec), (measured - inner) / overhead_scale(spec))) - }) - .collect(); - if usable.is_empty() { - eprintln!("warning: no calibration samples for '{}'", target); - continue; - } +/// A tuner for `P`, with the weight overrides applied to its estimating planner. +fn tuner>(opts: &Options) -> P { + let mut planner = P::new(); + if let Some(model) = planner.cost_model() { + planner.set_cost_model(opts.weights.apply(model)); + } + planner +} - // One median per log2 bucket, but only for buckets with enough samples to mean anything. - // Sparse buckets are dropped and filled in by interpolation instead. - let mut table: Vec<(u32, f64)> = Vec::new(); - if bucketed { - let mut by_bucket: std::collections::BTreeMap> = Default::default(); - for (bucket, residual) in usable.iter() { - by_bucket.entry(*bucket).or_default().push(*residual); - } - const MIN_PER_BUCKET: usize = 3; - table = by_bucket - .iter() - .filter(|(_, values)| values.len() >= MIN_PER_BUCKET) - .map(|(bucket, values)| (*bucket, median_of(values.clone()))) - .collect(); +/// The cost model for pricing specs offline, for a planner named on the command line or in a +/// dump header. wasm_simd borrows NEON's counts, as the planner itself does. +fn offline_model(planner_label: &str, opts: &Options) -> CostModel { + let instruction_set = match planner_label { + "sse" => InstructionSet::Sse, + "neon" | "wasm_simd" => InstructionSet::Neon, + other => { + eprintln!( + "the cost model has no instruction counts for planner '{}'", + other + ); + std::process::exit(2); } + }; + opts.weights.apply(CostModel::new( + instruction_set, + if opts.f32 { 2 } else { 1 }, + )) +} - // Too little data to describe a curve, or curves not wanted, so use one constant. - if table.len() < 2 { - table = vec![(0, median_of(usable.iter().map(|(_, r)| *r).collect()))]; - } +/// The planner a dump was taken with, from its header. +fn dump_planner(text: &str) -> String { + text.lines() + .find_map(|line| line.strip_prefix("# planner\t")) + .unwrap_or("neon") + .to_string() +} - let rendered: Vec = table - .iter() - .map(|(bucket, value)| format!("{}:{:.2}", 1usize << bucket, value)) - .collect(); - println!( - " {:<5} {:>3} buckets from {:>4} samples {}", - target, - table.len(), - usable.len(), - rendered.join(" ") - ); - model.overhead.insert(target, table); - fitted.push(target); - } - samples +fn percentile(sorted: &[f64], fraction: f64) -> f64 { + sorted[((sorted.len() as f64 * fraction) as usize).min(sorted.len() - 1)] } // --------------------------------------------------------------------------- // Subcommands // --------------------------------------------------------------------------- -struct Options { - rounds: usize, - block_ms: f64, - cap: usize, - verbose: bool, - bucketed: bool, - /// Where `dump` writes its rows. - out: Option, - /// Weights for the counted model. - params: counted::Params, - backend_explicit: bool, -} - fn cmd_time>(specs: &[String], opts: &Options) { - let mut planner = P::new(); + let mut planner = tuner::(opts); let mut subjects: Vec> = specs .iter() .map(|text| { @@ -338,18 +265,132 @@ fn cmd_time>(specs: &[String], opts: &Options) { } } +/// Error budget for a correct FFT of this length in this precision. +/// +/// Two terms, because two things are inexact. +/// +/// The recipe's own error grows like `eps * sqrt(log2 len)`: an FFT accumulates rounding over its +/// `log2 len` passes, not over its elements. The margin of 20 leaves room for recipes that compose +/// several algorithms, and still catches a real defect, which shows up as orders of magnitude +/// rather than as a factor of two. +/// +/// The reference's error grows like `eps * sqrt(len)`, because a naive DFT really does sum `len` +/// terms per output. The reference runs in f64, so for f32 recipes this term is negligible and the +/// budget is tight. For f64 recipes the reference is no better than the thing it judges, and this +/// term dominates: at len 100000 it is 2.8e-13 against the recipes' 1.8e-14. +fn tolerance(len: usize) -> f64 { + let eps = if std::mem::size_of::() == 4 { + f32::EPSILON as f64 + } else { + f64::EPSILON + }; + let len = len as f64; + 20.0 * eps * len.log2().max(1.0).sqrt() + 4.0 * f64::EPSILON * len.sqrt() +} + +/// Measure every candidate at every length and write the raw timings to a file. +/// +/// This exists so that a cost model can be scored and refitted offline, against a frozen +/// dataset, instead of needing the machine for every iteration. Rows are +/// `lenspecnspass`, where pass is 1 for the wide sweep and 2 for the careful +/// re-timing the fastest few get. +fn cmd_dump>(lengths: &[usize], opts: &Options) { + use std::io::Write; + const FINALISTS: usize = 8; + + let path = opts.out.clone().unwrap_or_else(|| "dump.tsv".to_string()); + let mut f = + std::io::BufWriter::new(std::fs::File::create(&path).expect("cannot create output")); + writeln!(f, "# planner\t{}", P::label()).unwrap(); + writeln!( + f, + "# rounds\t{}\tblock_ms\t{}\tcap\t{}", + opts.rounds, opts.block_ms, opts.cap + ) + .unwrap(); + writeln!(f, "len\tspec\tns\tpass\tplanner_pick").unwrap(); + + for &len in lengths { + let mut planner = tuner::(opts); + let picked = to_spec_string(&planner.plan(len)); + let specs = candidates_capped(&mut planner, len, opts.cap); + let mut subjects: Vec> = specs + .iter() + .map(|spec| { + let fft = planner.build(spec, FftDirection::Forward); + Subject::new(to_spec_string(spec), fft, 1) + }) + .collect(); + measure(&mut subjects, opts.rounds, opts.block_ms); + + // Re-time the fastest few for longer. The minimum of many noisy draws is biased low, so + // without this the best candidate looks faster than it is and every regret is flattered. + let mut order: Vec = (0..subjects.len()).collect(); + order.sort_by(|a, b| { + subjects[*a] + .best() + .partial_cmp(&subjects[*b].best()) + .unwrap() + }); + let finalists: Vec = order.into_iter().take(FINALISTS).collect(); + let mut finals: Vec> = finalists + .iter() + .map(|&i| { + let fft = planner.build(&specs[i], FftDirection::Forward); + Subject::new(subjects[i].name.clone(), fft, 1) + }) + .collect(); + measure(&mut finals, opts.rounds * 4, opts.block_ms); + + for (i, s) in subjects.iter().enumerate() { + let is_pick = if s.name == picked { "1" } else { "0" }; + writeln!(f, "{}\t{}\t{:.3}\t1\t{}", len, s.name, s.best(), is_pick).unwrap(); + let _ = i; + } + for s in finals.iter() { + let is_pick = if s.name == picked { "1" } else { "0" }; + writeln!(f, "{}\t{}\t{:.3}\t2\t{}", len, s.name, s.best(), is_pick).unwrap(); + } + println!("{:>9} {} candidates", len, subjects.len()); + } + println!("wrote {}", path); +} + +/// Candidates for `len` with the estimating planner's pick guaranteed among them, and its index. +/// +/// The exhaustive enumeration takes its inner recipes from the fixed planner, while the +/// estimating planner optimises its inners too, so its pick is often not in the list. +fn candidates_with_model_pick>( + planner: &mut P, + len: usize, + cap: usize, +) -> (Vec>, Option) { + let mut specs = candidates_capped(planner, len, cap); + let model_index = + planner + .estimate(len) + .map(|pick| match specs.iter().position(|spec| **spec == *pick) { + Some(index) => index, + None => { + specs.push(pick); + specs.len() - 1 + } + }); + (specs, model_index) +} + fn cmd_regret>(lengths: &[usize], opts: &Options) { println!( - "{:>9} {:>8} {:>10} {:>10} {}", - "len", "regret", "planner ns", "best ns", "best recipe (when it differs)" + "{:>9} {:>8} {:>8} {:>10} {:>10} {}", + "len", "planner", "model", "planner ns", "best ns", "best recipe (when neither picked it)" ); - let mut regrets: Vec<(f64, usize, String, String)> = Vec::new(); + let mut planner_regrets = Vec::new(); + let mut model_regrets = Vec::new(); for &len in lengths { - let mut planner = P::new(); - let specs = candidates_capped(&mut planner, len, opts.cap); - let planner_spec = to_spec_string(&specs[0]); + let mut planner = tuner::(opts); + let (specs, model_index) = candidates_with_model_pick(&mut planner, len, opts.cap); let mut subjects: Vec> = specs .iter() @@ -358,35 +399,33 @@ fn cmd_regret>(lengths: &[usize], opts: &Options Subject::new(to_spec_string(spec), fft, 1) }) .collect(); - measure(&mut subjects, opts.rounds, opts.block_ms); - let planner_time = subjects[0].best(); - let best_index = subjects - .iter() - .enumerate() - .min_by(|a, b| a.1.best().partial_cmp(&b.1.best()).unwrap()) - .map(|(i, _)| i) + let best_index = (0..subjects.len()) + .min_by(|&a, &b| subjects[a].best().total_cmp(&subjects[b].best())) .unwrap(); let best_time = subjects[best_index].best(); - let best_spec = subjects[best_index].name.clone(); - let regret = planner_time / best_time; + let planner_regret = subjects[0].best() / best_time; + let model_regret = model_index.map(|i| subjects[i].best() / best_time); println!( - "{:>9} {:>7.3}x {:>10.0} {:>10.0} {}", + "{:>9} {:>7.3}x {:>8} {:>10.0} {:>10.0} {}", len, - regret, - planner_time, + planner_regret, + model_regret + .map(|r| format!("{:.3}x", r)) + .unwrap_or_else(|| "-".into()), + subjects[0].best(), best_time, - if best_index == 0 { - "= planner".to_string() + if best_index == 0 || Some(best_index) == model_index { + String::new() } else { - best_spec.clone() + subjects[best_index].name.clone() } ); if opts.verbose { let mut ranked: Vec<&Subject> = subjects.iter().collect(); - ranked.sort_by(|a, b| a.best().partial_cmp(&b.best()).unwrap()); + ranked.sort_by(|a, b| a.best().total_cmp(&b.best())); for subject in ranked.iter().take(6) { println!( " {:>6.3}x {}", @@ -394,93 +433,89 @@ fn cmd_regret>(lengths: &[usize], opts: &Options subject.name ); } + println!(" planner: {}", subjects[0].name); + if let Some(i) = model_index { + println!(" model: {}", subjects[i].name); + } println!(" ({} candidates measured)", subjects.len()); } - regrets.push((regret, len, planner_spec, best_spec)); + planner_regrets.push(planner_regret); + if let Some(r) = model_regret { + model_regrets.push(r); + } } - regrets.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); - let n = regrets.len(); - let mean = regrets.iter().map(|r| r.0).sum::() / n as f64; - println!("\n--- regret: planner's pick divided by the best recipe measured ---"); - println!(" lengths {}", n); - println!(" mean {:.4}", mean); - println!(" median {:.4}", regrets[n / 2].0); - println!(" p90 {:.4}", regrets[(n * 9) / 10].0); - println!(" worst {:.4}", regrets[n - 1].0); - let losing = regrets.iter().filter(|r| r.0 > 1.02).count(); - println!(" more than 2% off the best: {} of {} lengths", losing, n); - println!("\nworst offenders:"); - for (regret, len, planner_spec, best_spec) in regrets.iter().rev().take(10) { - println!(" {:>8} {:.3}x", len, regret); - println!(" planner: {}", planner_spec); - println!(" best: {}", best_spec); + println!("\n--- regret: pick divided by the best recipe measured ---"); + for (label, mut values) in [("planner", planner_regrets), ("model", model_regrets)] { + if values.is_empty() { + continue; + } + values.sort_by(f64::total_cmp); + let mean = values.iter().sum::() / values.len() as f64; + println!( + " {:<8} n={:<4} mean {:.4} median {:.4} p90 {:.4} worst {:.4} more than 2% off: {}", + label, + values.len(), + mean, + percentile(&values, 0.5), + percentile(&values, 0.9), + values[values.len() - 1], + values.iter().filter(|&&r| r > 1.02).count() + ); } } -/// Time the shipping planner's pick against the counted model's pick, at every length in a range. +/// Time the fixed planner's pick against the estimating planner's, at every length given. /// -/// Unlike `regret`, this builds and times only two recipes per length rather than the whole -/// candidate set, which is what makes a thousand-length sweep affordable. When both agree, the -/// recipe is timed once and reported for both, so an agreement shows as exactly 1.000 rather than -/// as timing noise. +/// Builds and times only those two recipes per length, which is what makes a thousand lengths +/// affordable. When both agree, the recipe is timed once and reported for both, so an agreement +/// shows as exactly 1.000 rather than as timing noise. /// -/// Output is TSV on stdout, one row per length, ready to plot. +/// Output is TSV on stdout, one row per length, ready for `plot_sweep.py`, followed by a summary +/// in `#` comment lines. The `cands` column is kept for the plotting script and is always `-`. fn cmd_sweep>(lengths: &[usize], opts: &Options) { - let model = counted::CountedModel::new(opts.params); + let mut planner = tuner::(opts); + let Some(model) = planner.cost_model() else { + eprintln!( + "the {} planner does not estimate, so there is nothing to sweep", + P::label() + ); + std::process::exit(2); + }; println!("# planner\t{}", P::label()); - println!("# params\t{:?}", opts.params); - println!("# rounds\t{}\tblock_ms\t{}\tcap\t{}", opts.rounds, opts.block_ms, opts.cap); + println!("# elem\t{}", if opts.f32 { "f32" } else { "f64" }); + println!("# model\t{:?}", model); + println!("# rounds\t{}\tblock_ms\t{}", opts.rounds, opts.block_ms); println!( "len\tcands\tagree\tplanner_ns\tmodel_ns\tratio\tplanner_norm\tmodel_norm\tplanner_spec\tmodel_spec" ); + // Estimating over fixed, so below 1 means the estimating planner is faster. + let mut changes = Vec::new(); for &len in lengths { - let mut planner = P::new(); - let specs = plan_candidates(&mut planner, len, opts.cap); - if specs.is_empty() { - eprintln!("len {}: no candidates", len); - continue; - } - - let model_index = specs - .iter() - .enumerate() - .filter_map(|(i, spec)| model.cost(spec).map(|c| (i, c))) - .min_by(|a, b| a.1.partial_cmp(&b.1).unwrap()) - .map(|(i, _)| i) - .unwrap_or(0); - - let agree = model_index == 0; - let planner_spec = to_spec_string(&specs[0]); - let model_spec = to_spec_string(&specs[model_index]); - - let mut subjects: Vec> = if agree { - vec![Subject::new( - planner_spec.clone(), - planner.build(&specs[0], FftDirection::Forward), + let planner_spec = planner.plan(len); + let model_spec = planner.estimate(len).unwrap(); + let agree = planner_spec == model_spec; + + let mut subjects = vec![Subject::new( + to_spec_string(&planner_spec), + planner.build(&planner_spec, FftDirection::Forward), + 1, + )]; + if !agree { + subjects.push(Subject::new( + to_spec_string(&model_spec), + planner.build(&model_spec, FftDirection::Forward), 1, - )] - } else { - vec![ - Subject::new( - planner_spec.clone(), - planner.build(&specs[0], FftDirection::Forward), - 1, - ), - Subject::new( - model_spec.clone(), - planner.build(&specs[model_index], FftDirection::Forward), - 1, - ), - ] - }; + )); + } measure(&mut subjects, opts.rounds, opts.block_ms); let planner_ns = subjects[0].best(); - let model_ns = if agree { planner_ns } else { subjects[1].best() }; + let model_ns = subjects[subjects.len() - 1].best(); + changes.push(model_ns / planner_ns); // n log2 n, the work a radix-2 FFT of this length would do. Undefined at len 1, where // there is no work to normalise by. @@ -488,48 +523,67 @@ fn cmd_sweep>(lengths: &[usize], opts: &Options) let norm = |ns: f64| if nlogn > 0.0 { ns / nlogn } else { f64::NAN }; println!( - "{}\t{}\t{}\t{:.2}\t{:.2}\t{:.4}\t{:.5}\t{:.5}\t{}\t{}", + "{}\t-\t{}\t{:.2}\t{:.2}\t{:.4}\t{:.5}\t{:.5}\t{}\t{}", len, - specs.len(), if agree { 1 } else { 0 }, planner_ns, model_ns, planner_ns / model_ns, norm(planner_ns), norm(model_ns), - planner_spec, - model_spec + subjects[0].name, + subjects[subjects.len() - 1].name ); } + + if changes.is_empty() { + return; + } + let agreed = changes.iter().filter(|&&c| c == 1.0).count(); + let geomean = (changes.iter().map(|c| c.ln()).sum::() / changes.len() as f64).exp(); + changes.sort_by(f64::total_cmp); + println!("# summary\truntime of the estimating planner's pick over the fixed planner's"); + println!("# lengths\t{}\tsame recipe\t{}", changes.len(), agreed); + println!("# geomean\t{:.4}", geomean); + for (label, fraction) in [ + ("p10", 0.1), + ("p25", 0.25), + ("p50", 0.5), + ("p75", 0.75), + ("p90", 0.9), + ] { + println!("# {}\t{:.4}", label, percentile(&changes, fraction)); + } + println!( + "# best\t{:.4}\tworst\t{:.4}", + changes[0], + changes[changes.len() - 1] + ); } -/// Check that every enumerated candidate actually computes a correct FFT. -/// -/// Timing a recipe says nothing about whether it is valid, and an invalid one would happily -/// produce fast wrong answers. Each candidate is compared against a direct DFT. -/// Error budget for a correct FFT of this length in this precision. -/// -/// Two terms, because two things are inexact. -/// -/// The recipe's own error grows like `eps * sqrt(log2 len)`: an FFT accumulates rounding over its -/// `log2 len` passes, not over its elements. The margin of 20 leaves room for recipes that compose -/// several algorithms, and still catches a real defect, which shows up as orders of magnitude -/// rather than as a factor of two. -/// -/// The reference's error grows like `eps * sqrt(len)`, because a naive DFT really does sum `len` -/// terms per output. The reference runs in f64, so for f32 recipes this term is negligible and the -/// budget is tight. For f64 recipes the reference is no better than the thing it judges, and this -/// term dominates: at len 100000 it is 2.8e-13 against the recipes' 1.8e-14. -fn tolerance(len: usize) -> f64 { - let eps = if std::mem::size_of::() == 4 { - f32::EPSILON as f64 - } else { - f64::EPSILON +/// `count` distinct lengths drawn uniformly from `lo..=hi`, sorted, from a fixed seed. +fn random_lengths(lo: usize, hi: usize, count: usize, seed: u64) -> Vec { + // splitmix64, so the tool needs no dependencies and a seed reproduces a survey anywhere. + let mut state = seed; + let mut next = || { + state = state.wrapping_add(0x9e3779b97f4a7c15); + let mut z = state; + z = (z ^ (z >> 30)).wrapping_mul(0xbf58476d1ce4e5b9); + z = (z ^ (z >> 27)).wrapping_mul(0x94d049bb133111eb); + z ^ (z >> 31) }; - let len = len as f64; - 20.0 * eps * len.log2().max(1.0).sqrt() + 4.0 * f64::EPSILON * len.sqrt() + let span = (hi - lo + 1) as u64; + let count = count.min(span as usize); + let mut lengths = std::collections::BTreeSet::new(); + while lengths.len() < count { + lengths.insert(lo + (next() % span) as usize); + } + lengths.into_iter().collect() } +/// Check that every enumerated candidate, and the estimating planner's pick, computes a correct +/// FFT. Timing a recipe says nothing about whether it is valid, and an invalid one would happily +/// produce fast wrong answers. Each candidate is compared against a direct DFT. fn cmd_verify>(lengths: &[usize], opts: &Options) { let mut worst_overall: f64 = 0.0; let mut failures = 0usize; @@ -558,8 +612,8 @@ fn cmd_verify>(lengths: &[usize], reference_fft.process_with_scratch(&mut reference, &mut reference_scratch); let reference_norm: f64 = reference.iter().map(|c| c.norm_sqr()).sum::().sqrt(); - let mut planner = P::new(); - let specs = candidates_capped(&mut planner, len, opts.cap); + let mut planner = tuner::(opts); + let (specs, _) = candidates_with_model_pick(&mut planner, len, opts.cap); let mut worst_here: f64 = 0.0; let mut worst_spec = String::new(); @@ -573,10 +627,8 @@ fn cmd_verify>(lengths: &[usize], .iter() .zip(reference.iter()) .map(|(a, b)| { - let d = Complex::new( - a.re.to_f64().unwrap() - b.re, - a.im.to_f64().unwrap() - b.im, - ); + let d = + Complex::new(a.re.to_f64().unwrap() - b.re, a.im.to_f64().unwrap() - b.im); d.norm_sqr() }) .sum::() @@ -613,290 +665,160 @@ fn cmd_verify>(lengths: &[usize], } } -/// Show how each algorithm's per-element overhead varies with size. -fn cmd_residuals>(lengths: &[usize], opts: &Options) { - let max_len = lengths.iter().copied().max().unwrap_or(1024) * 4; - eprintln!("measuring primitives..."); - let mut model = calibrate_primitives::(opts.rounds, opts.block_ms, max_len); - eprintln!("fitting overheads..."); - let samples = fit_overheads::( - &mut model, - lengths, - opts.rounds, - opts.block_ms, - opts.cap, - opts.bucketed, - ); - - let mut binned: std::collections::BTreeMap< - &'static str, - std::collections::BTreeMap>, - > = Default::default(); - for (spec, measured) in samples.iter() { - if !FIT_ORDER.contains(&spec.kind()) { +/// Rows of a dump file: len -> [(spec, best ns, is the fixed planner's pick)]. +/// +/// Pass 2, the careful re-timing, overwrites pass 1. The inner list preserves the dump's own +/// order, which is the order `candidates_capped` enumerated in, so that replay breaks cost ties +/// the same way enumeration would. Keying by spec string instead would sort `gts(b10,b3)` ahead +/// of `gts(b3,b10)` and silently reverse every width/height tie. +fn read_dump(text: &str) -> std::collections::BTreeMap> { + let mut data: std::collections::BTreeMap> = Default::default(); + for line in text.lines() { + if line.starts_with('#') || line.starts_with("len\t") { continue; } - if let Some(inner) = model.inner_cost(spec) { - let residual = (measured - inner) / overhead_scale(spec); - binned - .entry(spec.kind()) - .or_default() - .entry(bucket_of(spec)) - .or_default() - .push(residual); + let f: Vec<&str> = line.split('\t').collect(); + if f.len() < 5 { + continue; } - } - - println!( - "{:<5} {:>8} {:>8} {:>9} {:>7}", - "kind", "len>=", "median", "p25..p75", "n" - ); - for (kind, buckets) in binned { - for (bucket, values) in buckets { - if values.len() < 3 { - continue; + let (len, spec, ns, pass, pick) = ( + f[0].parse::().unwrap(), + f[1].to_string(), + f[2].parse::().unwrap(), + f[3].parse::().unwrap(), + f[4] == "1", + ); + let rows = data.entry(len).or_default(); + match rows.iter_mut().find(|row| row.0 == spec) { + Some(row) if pass != 1 => { + row.1 = ns; + row.2 = pick; } - let mut sorted = values.clone(); - sorted.sort_by(|a, b| a.partial_cmp(b).unwrap()); - let n = sorted.len(); - println!( - "{:<5} {:>8} {:>8.3} {:>9} {:>7}", - kind, - 1usize << bucket, - sorted[n / 2], - format!("{:.2}..{:.2}", sorted[n / 4], sorted[(n * 3) / 4]), - n - ); + Some(_) => {} + None => rows.push((spec, ns, pick)), } - println!(); } + data } -/// Measure every candidate at every length and write the raw timings to a file. +/// Score the cost model against a dump file. No measurement, no planner, no machine. /// -/// This exists so that a cost model can be scored and refitted offline, against a frozen -/// dataset, instead of needing the machine for every iteration. Rows are -/// `lenspecnspass`, where pass is 1 for the wide sweep and 2 for the careful -/// re-timing the fastest few get. -fn cmd_dump>(lengths: &[usize], opts: &Options) { - use std::io::Write; - const FINALISTS: usize = 8; +/// This replays a one-level choice among the dumped candidates, whose inner recipes are the fixed +/// planner's. It is how the weights were fitted, and it is not the same as the estimating +/// planner's pick, which also optimises the inner recipes; `sweep` measures that. +fn cmd_score(path: &str, opts: &Options) { + let text = std::fs::read_to_string(path).expect("cannot read dump"); + let model = offline_model(&dump_planner(&text), opts); + println!("model: {:?}", model); + println!( + "{:>9} {:>9} {:>9} {}", + "len", "model", "planner", "model's pick (when it is not the best)" + ); - let path = opts.out.clone().unwrap_or_else(|| "dump.tsv".to_string()); - let mut f = std::io::BufWriter::new(std::fs::File::create(&path).expect("cannot create output")); - writeln!(f, "# planner\t{}", P::label()).unwrap(); - writeln!(f, "# rounds\t{}\tblock_ms\t{}\tcap\t{}", opts.rounds, opts.block_ms, opts.cap).unwrap(); - writeln!(f, "len\tspec\tns\tpass\tplanner_pick").unwrap(); + let (mut model_regrets, mut planner_regrets) = (Vec::new(), Vec::new()); + for (len, rows) in read_dump(&text) { + let best = rows.iter().map(|row| row.1).fold(f64::INFINITY, f64::min); + let planner_regret = rows.iter().find(|row| row.2).map(|row| row.1 / best); - for &len in lengths { - let mut planner = P::new(); - let picked = to_spec_string(&planner.plan(len)); - let specs = candidates_capped(&mut planner, len, opts.cap); - let mut subjects: Vec> = specs - .iter() - .map(|spec| { - let fft = planner.build(spec, FftDirection::Forward); - Subject::new(to_spec_string(spec), fft, 1) - }) - .collect(); - measure(&mut subjects, opts.rounds, opts.block_ms); - - // Re-time the fastest few for longer. The minimum of many noisy draws is biased low, so - // without this the best candidate looks faster than it is and every regret is flattered. - let mut order: Vec = (0..subjects.len()).collect(); - order.sort_by(|a, b| subjects[*a].best().partial_cmp(&subjects[*b].best()).unwrap()); - let finalists: Vec = order.into_iter().take(FINALISTS).collect(); - let mut finals: Vec> = finalists + let pick = rows .iter() - .map(|&i| { - let fft = planner.build(&specs[i], FftDirection::Forward); - Subject::new(subjects[i].name.clone(), fft, 1) + .filter_map(|(text, ns, _)| { + let spec = parse(text).ok()?; + Some((text, spec_cost(&model, &spec)?, *ns)) }) - .collect(); - measure(&mut finals, opts.rounds * 4, opts.block_ms); - - for (i, s) in subjects.iter().enumerate() { - let is_pick = if s.name == picked { "1" } else { "0" }; - writeln!(f, "{}\t{}\t{:.3}\t1\t{}", len, s.name, s.best(), is_pick).unwrap(); - let _ = i; - } - for s in finals.iter() { - let is_pick = if s.name == picked { "1" } else { "0" }; - writeln!(f, "{}\t{}\t{:.3}\t2\t{}", len, s.name, s.best(), is_pick).unwrap(); - } - println!("{:>9} {} candidates", len, subjects.len()); - } - println!("wrote {}", path); -} - -/// Score the counted model against a dump file. No measurement, no planner, no machine. -fn cmd_score(path: &str, opts: &Options) { - use std::collections::BTreeMap; - let text = std::fs::read_to_string(path).expect("cannot read dump"); - - // len -> [(spec, best ns, is planner pick)]. Pass 2 overwrites pass 1. - // - // The inner container must preserve the dump's own order, which is the order - // `candidates_capped` enumerated in. The planner takes the first minimum in that order, so - // replay has to break cost ties the same way or it does not model the planner. Keying by - // spec string instead sorts `gts(b10,b3)` ahead of `gts(b3,b10)`, which is the opposite of - // the enumeration order and silently reverses every width/height tie. - let mut data: BTreeMap> = BTreeMap::new(); - for line in text.lines() { - if line.starts_with('#') || line.starts_with("len\t") { + .min_by(|a, b| a.1.total_cmp(&b.1)); + let Some((pick_name, _, pick_ns)) = pick else { + println!("{:>9} no candidate priced", len); continue; - } - let f: Vec<&str> = line.split('\t').collect(); - if f.len() < 5 { - continue; - } - let (len, spec, ns, pass, pick) = ( - f[0].parse::().unwrap(), - f[1].to_string(), - f[2].parse::().unwrap(), - f[3].parse::().unwrap(), - f[4] == "1", - ); - let e: &mut Vec<(String, f64, bool)> = data.entry(len).or_default(); - match e.iter_mut().find(|r| r.0 == spec) { - // pass 2 is the careful re-timing, so it wins - Some(row) if pass != 1 => { - row.1 = ns; - row.2 = pick; - } - Some(_) => {} - None => e.push((spec, ns, pick)), - } - } - - // Take the backend from the dump header unless it was given explicitly, so an SSE dump is - // never priced with NEON instruction costs by accident. - let mut params = opts.params; - if !opts.backend_explicit { - if let Some(label) = text - .lines() - .find_map(|l| l.strip_prefix("# planner\t")) - .and_then(counted::Backend::parse) - { - params.backend = label; - } - } - let model = counted::CountedModel::new(params); - println!("params: {:?}", params); - println!( - "{:>9} {:>9} {:>9} {}", - "len", "counted", "planner", "counted model's pick (when it is not the best)" - ); - - let (mut mr, mut pr) = (Vec::new(), Vec::new()); - for (&len, rows) in &data { - let best = rows.iter().map(|v| v.1).fold(f64::INFINITY, f64::min); - let planner_ns = rows.iter().find(|v| v.2).map(|v| v.1); - - let mut scored: Vec<(String, f64, f64)> = Vec::new(); - for (spec_text, ns, _) in rows { - let spec = match parse(spec_text) { - Ok(s) => s, - Err(_) => continue, - }; - if let Some(c) = model.cost(&spec) { - scored.push((spec_text.clone(), c, *ns)); - } - } - let pick = scored - .iter() - .min_by(|a, b| a.1.partial_cmp(&b.1).unwrap()); - let (pick_name, pick_ns) = match pick { - Some((n, _, ns)) => (n.clone(), *ns), - None => { - println!("{:>9} no candidate priced", len); - continue; - } }; - let m = pick_ns / best; - mr.push(m); - let p = planner_ns.map(|n| n / best); - if let Some(p) = p { - pr.push(p); + let model_regret = pick_ns / best; + model_regrets.push(model_regret); + if let Some(p) = planner_regret { + planner_regrets.push(p); } - let shown = if m <= 1.0001 { "= best".to_string() } else { pick_name }; println!( "{:>9} {:>8.3}x {:>8} {}", len, - m, - p.map(|p| format!("{:.3}x", p)).unwrap_or_else(|| "-".into()), - shown + model_regret, + planner_regret + .map(|p| format!("{:.3}x", p)) + .unwrap_or_else(|| "-".into()), + if model_regret <= 1.0001 { + "= best" + } else { + pick_name + } ); } - let stat = |v: &mut Vec, label: &str| { - v.sort_by(|a, b| a.partial_cmp(b).unwrap()); - let mean = v.iter().sum::() / v.len() as f64; - let median = v[v.len() / 2]; - let p90 = v[((v.len() as f64 * 0.9) as usize).min(v.len() - 1)]; - let worst = *v.last().unwrap(); + println!("\n--- regret, pick divided by best measured ---"); + for (label, mut values) in [("model", model_regrets), ("planner", planner_regrets)] { + if values.is_empty() { + continue; + } + values.sort_by(f64::total_cmp); println!( " {:<9} n={:<4} mean {:.4} median {:.4} p90 {:.4} worst {:.4}", label, - v.len(), - mean, - median, - p90, - worst + values.len(), + values.iter().sum::() / values.len() as f64, + percentile(&values, 0.5), + percentile(&values, 0.9), + values[values.len() - 1] ); - }; - println!("\n--- regret, pick divided by best measured ---"); - stat(&mut mr, "counted"); - stat(&mut pr, "planner"); + } } -/// Print the counted model's cost tree for one recipe, so the recursion can be checked by eye. -fn cmd_explain(spec_text: &str, opts: &Options) { +/// Print the model's cost for every measured candidate, for offline analysis. +/// +/// Columns: len, spec, measured ns, model cost. Pure replay, no machine needed. +fn cmd_costs(path: &str, opts: &Options) { + let text = std::fs::read_to_string(path).expect("cannot read dump"); + let model = offline_model(&dump_planner(&text), opts); + println!("len\tspec\tns\tcost\tplanner_pick"); + for (len, rows) in read_dump(&text) { + for (spec_text, ns, pick) in rows { + let cost = parse(&spec_text) + .ok() + .and_then(|spec| spec_cost(&model, &spec)) + .map(|c| format!("{:.1}", c)) + .unwrap_or_else(|| "NA".into()); + println!( + "{}\t{}\t{:.3}\t{}\t{}", + len, + spec_text, + ns, + cost, + if pick { 1 } else { 0 } + ); + } + } +} + +/// Print the model's cost tree for one recipe, so the recursion can be checked by eye. +fn cmd_explain(spec_text: &str, planner_label: &str, opts: &Options) { let spec = parse(spec_text).expect("could not parse spec"); - let model = counted::CountedModel::new(opts.params); - let root_ws = spec.len() as f64; - - fn walk( - model: &counted::CountedModel, - spec: &Spec, - root_ws: f64, - mult: f64, - depth: usize, - out: &mut Vec, - ) { - let own = model.cost(spec).unwrap_or(f64::NAN); - // cost of the children alone, at the multiplicity the parent runs them - let (kids, child_total): (Vec<(&Spec, f64)>, f64) = match spec { - Spec::MixedRadix { left, right, .. } | Spec::GoodThomas { left, right, .. } => { - let v = vec![ - (left.as_ref(), right.len() as f64), - (right.as_ref(), left.len() as f64), - ]; - let t = v - .iter() - .map(|(c, m)| m * model.cost(c).unwrap_or(0.0)) - .sum(); - (v, t) - } + let model = offline_model(planner_label, opts); + + fn walk(model: &CostModel, spec: &Spec, mult: f64, depth: usize, out: &mut Vec) { + let own = spec_cost(model, spec).unwrap_or(f64::NAN); + // Each child, at the multiplicity the parent runs it. + let kids: Vec<(&Spec, f64)> = match spec { + Spec::MixedRadix { left, right, .. } | Spec::GoodThomas { left, right, .. } => vec![ + (left.as_ref(), right.len() as f64), + (right.as_ref(), left.len() as f64), + ], Spec::RadixN { radixes, base } => { - let m = radixes.iter().product::() as f64; - ( - vec![(base.as_ref(), m)], - m * model.cost(base).unwrap_or(0.0), - ) + vec![(base.as_ref(), radixes.iter().product::() as f64)] } - Spec::Radix4 { k, base } => { - let m = (1u64 << (2 * k)) as f64; - ( - vec![(base.as_ref(), m)], - m * model.cost(base).unwrap_or(0.0), - ) - } - Spec::Raders { inner } | Spec::Bluesteins { inner, .. } => { - (vec![(inner.as_ref(), 2.0)], 2.0 * model.cost(inner).unwrap_or(0.0)) - } - _ => (vec![], 0.0), + Spec::Radix4 { k, base } => vec![(base.as_ref(), (1u64 << (2 * k)) as f64)], + Spec::Raders { inner } | Spec::Bluesteins { inner, .. } => vec![(inner.as_ref(), 2.0)], + _ => vec![], }; + let child_total: f64 = kids + .iter() + .map(|(child, m)| m * spec_cost(model, child).unwrap_or(0.0)) + .sum(); out.push(format!( "{:indent$}{:<34} len {:>7} x{:<8.0} cost {:>14.0} own {:>12.0}", "", @@ -908,116 +830,116 @@ fn cmd_explain(spec_text: &str, opts: &Options) { indent = depth * 2 )); for (child, m) in kids { - walk(model, child, root_ws, mult * m, depth + 1, out); + walk(model, child, mult * m, depth + 1, out); } } let mut out = Vec::new(); - walk(&model, &spec, root_ws, 1.0, 0, &mut out); - println!("params: {:?}", opts.params); - println!("{:<36} {:>11} {:<9} {:>19} {:>16}", "recipe", "len", "times", "total cost", "own cost"); + walk(&model, &spec, 1.0, 0, &mut out); + println!("model: {:?}", model); + println!( + "{:<36} {:>11} {:<9} {:>19} {:>16}", + "recipe", "len", "times", "total cost", "own cost" + ); for line in out { println!("{}", line); } } -/// Print the model's cost for every measured candidate, for offline analysis. +/// Compare the cost of *choosing* a recipe against the cost of *building* it. /// -/// Columns: len, spec, measured ns, model cost. Pure replay, no machine needed. -fn cmd_costs(path: &str, opts: &Options) { - use std::collections::BTreeMap; - let text = std::fs::read_to_string(path).expect("cannot read dump"); - let mut params = opts.params; - if !opts.backend_explicit { - if let Some(b) = text - .lines() - .find_map(|l| l.strip_prefix("# planner\t")) - .and_then(counted::Backend::parse) - { - params.backend = b; - } - } - let model = counted::CountedModel::new(params); - - let mut data: BTreeMap> = BTreeMap::new(); - for line in text.lines() { - if line.starts_with('#') || line.starts_with("len\t") { - continue; - } - let f: Vec<&str> = line.split('\t').collect(); - if f.len() < 5 { - continue; - } - let (len, spec, ns, pass, pick) = ( - f[0].parse::().unwrap(), - f[1].to_string(), - f[2].parse::().unwrap(), - f[3].parse::().unwrap(), - f[4] == "1", - ); - let e = data.entry(len).or_default(); - match e.get(&spec) { - Some(_) if pass == 1 => {} - _ => { - e.insert(spec, (ns, pick)); +/// Planning only produces a `Recipe`. Turning that into an `Arc` is a separate and much +/// larger job: every algorithm precomputes twiddles, and Rader's and Bluestein's both run a full +/// inner FFT inside their constructors. So the question that decides whether the estimating +/// planner is affordable is not how much slower it is than the fixed planner, but how much it +/// adds to plan-plus-build, which is what a caller actually pays before the first transform. +/// +/// Every timing uses a fresh planner, so these are cold-start figures: nothing is served from a +/// cache, including the inner lengths a session planning several lengths would share. +fn cmd_plantime>(lengths: &[usize], opts: &Options) { + println!( + "{:>7} {:>12} {:>14} {:>12} {:>9} {:>11}", + "len", "plan fixed", "plan estimate", "build", "build/plan", "extra vs" + ); + println!( + "{:>7} {:>12} {:>14} {:>12} {:>9} {:>11}", + "", "ns", "ns", "ns", "", "plan+build" + ); + let (mut tot_a, mut tot_b, mut tot_c) = (0.0, 0.0, 0.0); + for &len in lengths { + let reps = 50; + let time_plan = |plan: &dyn Fn(&mut P) -> Arc| { + let mut planners: Vec

= (0..reps).map(|_| tuner::(opts)).collect(); + let start = Instant::now(); + for planner in planners.iter_mut() { + std::hint::black_box(plan(planner)); } + start.elapsed().as_secs_f64() * 1e9 / reps as f64 + }; + let a = time_plan(&|planner| planner.plan(len)); + let b = time_plan(&|planner| planner.estimate(len).unwrap()); + + // Build the recipe the estimating planner chose. `build` uses a fresh planner each time. + let mut planner = tuner::(opts); + let spec = planner.estimate(len).unwrap(); + let build_reps = if len > 4096 { + 5 + } else if len > 256 { + 20 + } else { + 100 + }; + let start = Instant::now(); + for _ in 0..build_reps { + std::hint::black_box(planner.build(&spec, FftDirection::Forward)); } - } + let c = start.elapsed().as_secs_f64() * 1e9 / build_reps as f64; - println!("len\tspec\tns\tcost\tplanner_pick"); - for (len, rows) in &data { - for (spec_text, (ns, pick)) in rows { - let cost = parse(spec_text) - .ok() - .and_then(|s| model.cost(&s)) - .map(|c| format!("{:.1}", c)) - .unwrap_or_else(|| "NA".into()); - println!( - "{}\t{}\t{:.3}\t{}\t{}", - len, - spec_text, - ns, - cost, - if *pick { 1 } else { 0 } - ); - } + tot_a += a; + tot_b += b; + tot_c += c; + println!( + "{:>7} {:>12.0} {:>14.0} {:>12.0} {:>8.0}x {:>10.1}%", + len, + a, + b, + c, + c / b, + 100.0 * (b - a) / (a + c) + ); } + println!( + "\n totals: plan fixed {:.0} ns, plan estimate {:.0} ns ({:.1}x), build {:.0} ns", + tot_a, + tot_b, + tot_b / tot_a, + tot_c + ); + println!( + " the estimating planner adds {:.2}% to plan-plus-build", + 100.0 * (tot_b - tot_a) / (tot_a + tot_c) + ); } -/// Compare the plan-time cost of the fixed planner against enumerate-and-price. -/// -/// The fixed planner answers from a few integer operations. A cost model has to enumerate the -/// candidate set and price every member, which is real work the fixed planner never does. This -/// is the one axis where the fixed planner is unambiguously ahead, so it should be measured. -/// Compare the cost of *choosing* a recipe against the cost of *building* it. -/// -/// Planning only produces a `Recipe`. Turning that into an `Arc` is a separate and much -/// larger job: every algorithm precomputes twiddles, and Rader's and Bluestein's both run a full -/// inner FFT inside their constructors. So the question that decides whether an estimating -/// planner is affordable is not how much slower it is than the fixed planner, but how much it -/// adds to plan-plus-build, which is what a caller actually pays before the first transform. -/// How the best recipe changes once construction cost is counted, which is the question a -/// quick-and-dirty planner for one-shot transforms exists to answer. +/// How the best recipe changes once construction cost is counted. /// /// For every candidate this measures build time and execution time, then reports which recipe -/// minimises `build + k * execute` at several values of `k`. If the same recipe wins at every `k` -/// there is nothing for a construction-aware planner to choose, and the idea is dead. Plan time is -/// deliberately excluded: it is the same constant for every candidate at one length, so it cannot -/// change which recipe wins. +/// minimises `build + k * execute` at several values of `k`, and how many executions the +/// estimating planner's pick needs to repay any extra build cost. Plan time is deliberately +/// excluded: it is the same for every candidate at one length, so it cannot change which wins. fn cmd_crossover>(lengths: &[usize], opts: &Options) { - let model = counted::CountedModel::new(opts.params); println!( "{:>7} {:>6} {:>6} {:>11} {:>11} {:>12} {}", "len", "cands", "k", "build ns", "exec ns", "total ns", "recipe" ); for &len in lengths { - let mut planner = P::new(); - let specs = plan_candidates(&mut planner, len, opts.cap); - if specs.is_empty() { - eprintln!("len {}: no candidates", len); - continue; - } + let mut planner = tuner::(opts); + let (specs, model_index) = candidates_with_model_pick(&mut planner, len, opts.cap); + let Some(mi) = model_index else { + eprintln!("the {} planner does not estimate", P::label()); + std::process::exit(2); + }; let mut subjects: Vec> = specs .iter() @@ -1029,60 +951,58 @@ fn cmd_crossover>(lengths: &[usize], opts: &Opti measure(&mut subjects, opts.rounds, opts.block_ms); let exec: Vec = subjects.iter().map(|s| s.best()).collect(); - // A fresh planner per repetition, so no inner FFT is served from a cache. - let build_reps = if len > 4096 { 5 } else if len > 256 { 20 } else { 100 }; + // `build` uses a fresh planner each time, so no inner FFT is served from a cache. + let build_reps = if len > 4096 { + 5 + } else if len > 256 { + 20 + } else { + 100 + }; let build: Vec = specs .iter() .map(|spec| { - let t = Instant::now(); + let start = Instant::now(); for _ in 0..build_reps { - let mut pl = P::new(); - std::hint::black_box(pl.build(spec, FftDirection::Forward)); + std::hint::black_box(planner.build(spec, FftDirection::Forward)); } - t.elapsed().as_secs_f64() * 1e9 / build_reps as f64 + start.elapsed().as_secs_f64() * 1e9 / build_reps as f64 }) .collect(); let pick = |k: f64| -> usize { (0..specs.len()) - .min_by(|&a, &b| { - (build[a] + k * exec[a]) - .partial_cmp(&(build[b] + k * exec[b])) - .unwrap() - }) + .min_by(|&a, &b| (build[a] + k * exec[a]).total_cmp(&(build[b] + k * exec[b]))) .unwrap() }; - let mut first = true; - for k in [1.0, 10.0, 100.0, 1000.0] { + for (row, k) in [1.0, 10.0, 100.0, 1000.0].into_iter().enumerate() { let i = pick(k); println!( "{:>7} {:>6} {:>6} {:>11.0} {:>11.1} {:>12.0} {}", - if first { len.to_string() } else { String::new() }, - if first { specs.len().to_string() } else { String::new() }, + if row == 0 { + len.to_string() + } else { + String::new() + }, + if row == 0 { + specs.len().to_string() + } else { + String::new() + }, k as usize, build[i], exec[i], build[i] + k * exec[i], subjects[i].name ); - first = false; } - // What the counted model picks, which optimises execution alone. - let mi = specs - .iter() - .enumerate() - .filter_map(|(i, sp)| model.cost(sp).map(|c| (i, c))) - .min_by(|a, b| a.1.partial_cmp(&b.1).unwrap()) - .map(|(i, _)| i) - .unwrap_or(0); let one = pick(1.0); println!( "{:>7} {:>6} {:>6} {:>11.0} {:>11.1} {:>12} {}", "", "", "model", build[mi], exec[mi], "", subjects[mi].name ); - // Crossover: how many executions before the model's pick repays its extra build cost. if one != mi && exec[mi] < exec[one] { let k = (build[mi] - build[one]) / (exec[one] - exec[mi]); println!( @@ -1099,363 +1019,156 @@ fn cmd_crossover>(lengths: &[usize], opts: &Opti } } -fn cmd_plantime>(lengths: &[usize], opts: &Options) { - let model = counted::CountedModel::new(opts.params); - println!( - "{:>7} {:>6} {:>12} {:>14} {:>12} {:>9} {:>11}", - "len", "cands", "plan fixed", "plan+price", "build", "build/plan", "extra vs" - ); - println!( - "{:>7} {:>6} {:>12} {:>14} {:>12} {:>9} {:>11}", - "", "", "ns", "ns", "ns", "", "plan+build" - ); - let (mut tot_a, mut tot_b, mut tot_c) = (0.0, 0.0, 0.0); - for &len in lengths { - // fixed planner: design only, with a fresh planner each time so nothing is cached - let reps = 200; - let t0 = Instant::now(); - for _ in 0..reps { - let mut pl = P::new(); - std::hint::black_box(pl.plan(len)); - } - let a = t0.elapsed().as_secs_f64() * 1e9 / reps as f64; - - let mut pl = P::new(); - let n = plan_candidates(&mut pl, len, opts.cap).len(); - let t1 = Instant::now(); - for _ in 0..reps { - let mut pl = P::new(); - let specs = plan_candidates(&mut pl, len, opts.cap); - let best = specs - .iter() - .filter_map(|sp| model.cost(sp).map(|c| (c, sp))) - .min_by(|x, y| x.0.partial_cmp(&y.0).unwrap()); - std::hint::black_box(best); - } - let b = t1.elapsed().as_secs_f64() * 1e9 / reps as f64; - - // Build the recipe the fixed planner chose. Construction allocates and precomputes, so - // it is far slower than planning; use fewer repetitions and a fresh planner each time so - // nothing is served from a cache. - let mut pl = P::new(); - let spec = pl.plan(len); - let build_reps = if len > 4096 { 5 } else if len > 256 { 20 } else { 100 }; - let t2 = Instant::now(); - for _ in 0..build_reps { - let mut pl = P::new(); - std::hint::black_box(pl.build(&spec, FftDirection::Forward)); - } - let c = t2.elapsed().as_secs_f64() * 1e9 / build_reps as f64; - - tot_a += a; - tot_b += b; - tot_c += c; - println!( - "{:>7} {:>6} {:>12.0} {:>14.0} {:>12.0} {:>8.0}x {:>10.1}%", - len, n, a, b, c, c / a, 100.0 * (b - a) / (a + c) - ); - } - println!( - "\n totals: plan fixed {:.0} ns, plan+price {:.0} ns ({:.1}x), build {:.0} ns", - tot_a, tot_b, tot_b / tot_a, tot_c - ); - println!( - " building is {:.0}x planning; the estimating planner adds {:.2}% to plan-plus-build", - tot_c / tot_a, - 100.0 * (tot_b - tot_a) / (tot_a + tot_c) - ); -} - -fn cmd_model>(train: &[usize], test: &[usize], opts: &Options) { - let max_len = test.iter().chain(train.iter()).copied().max().unwrap_or(1024) * 4; - - println!("planner: {}", P::label()); - println!("measuring primitives..."); - let mut model = calibrate_primitives::(opts.rounds, opts.block_ms, max_len); - - println!("fitting overheads on {} training lengths...", train.len()); - fit_overheads::( - &mut model, - train, - opts.rounds, - opts.block_ms, - opts.cap, - opts.bucketed, - ); - println!("\n{}", model.describe()); - - println!( - "{:>9} {:>9} {:>9} {}", - "len", "model", "planner", "model's pick (when it is not the best)" - ); - let mut model_regrets = Vec::new(); - let mut planner_regrets = Vec::new(); - - for &len in test { - let mut planner = P::new(); - let specs = candidates_capped(&mut planner, len, opts.cap); - let mut subjects: Vec> = specs - .iter() - .map(|spec| { - let fft = planner.build(spec, FftDirection::Forward); - Subject::new(to_spec_string(spec), fft, 1) - }) - .collect(); - measure(&mut subjects, opts.rounds, opts.block_ms); - - // Fastest of the first pass, used only to pick which recipes deserve a careful re-timing. - let best_index = subjects - .iter() - .enumerate() - .min_by(|a, b| a.1.best().partial_cmp(&b.1.best()).unwrap()) - .map(|(i, _)| i) - .unwrap(); - - let model_index = specs - .iter() - .enumerate() - .filter_map(|(i, spec)| model.cost(spec).map(|c| (i, c))) - .min_by(|a, b| a.1.partial_cmp(&b.1).unwrap()) - .map(|(i, _)| i); - - let model_spec = match model_index { - Some(i) => subjects[i].name.clone(), - None => "".to_string(), - }; - - // Second pass. The winner of a wide comparison is biased fast: with sub-percent noise, - // the minimum of many draws lands below the true minimum, which puts a floor under any - // regret measured against it. Re-timing just the recipes of interest, for longer, - // removes most of that bias. - let finalists: Vec = { - let mut picked = vec![0usize, best_index]; - if let Some(i) = model_index { - picked.push(i); - } - picked.sort_unstable(); - picked.dedup(); - picked - }; - let mut finals: Vec> = finalists - .iter() - .map(|&i| { - let fft = planner.build(&specs[i], FftDirection::Forward); - Subject::new(to_spec_string(&specs[i]), fft, 1) - }) - .collect(); - measure(&mut finals, opts.rounds * 4, opts.block_ms); - - let time_of = - |index: usize| -> f64 { finals[finalists.iter().position(|&i| i == index).unwrap()].best() }; - let planner_time = time_of(0); - let best_time = finalists - .iter() - .map(|&i| time_of(i)) - .fold(f64::INFINITY, f64::min); - let model_time = match model_index { - Some(i) => time_of(i), - None => f64::NAN, - }; - - let model_regret = model_time / best_time; - let planner_regret = planner_time / best_time; - model_regrets.push(model_regret); - planner_regrets.push(planner_regret); - - println!( - "{:>9} {:>8.3}x {:>8.3}x {}", - len, - model_regret, - planner_regret, - if model_regret <= 1.001 { - "= best".to_string() - } else { - model_spec - } - ); - } - - for (label, mut values) in [("model ", model_regrets), ("planner", planner_regrets)] { - values.sort_by(|a, b| a.partial_cmp(b).unwrap()); - let n = values.len(); - let mean = values.iter().sum::() / n as f64; - println!( - "{}: mean {:.4} median {:.4} p90 {:.4} worst {:.4}", - label, - mean, - values[n / 2], - values[(n * 9) / 10], - values[n - 1] - ); - } -} - -fn cmd_emit>(lengths: &[usize], opts: &Options, element: &str) { - let max_len = lengths.iter().copied().max().unwrap_or(1024) * 4; - eprintln!("measuring primitives..."); - let mut model = calibrate_primitives::(opts.rounds, opts.block_ms, max_len); - eprintln!("fitting overheads on {} lengths...", lengths.len()); - fit_overheads::( - &mut model, - lengths, - opts.rounds, - opts.block_ms, - opts.cap, - opts.bucketed, - ); - print!("{}", emit::emit(&model, element, P::label())); -} - // --------------------------------------------------------------------------- enum Command { Time(Vec), Regret(Vec), - Model(Vec, Vec), - Residuals(Vec), + Sweep(Vec), Verify(Vec), - Emit(Vec), Dump(Vec), Score(String), Explain(String), Costs(String), Plantime(Vec), Crossover(Vec), - Sweep(Vec), } -fn run>( - command: &Command, - opts: &Options, - element: &str, -) { +fn run>(command: &Command, opts: &Options) { match command { Command::Time(specs) => cmd_time::(specs, opts), Command::Regret(lengths) => cmd_regret::(lengths, opts), Command::Sweep(lengths) => cmd_sweep::(lengths, opts), - Command::Model(train, test) => cmd_model::(train, test, opts), - Command::Residuals(lengths) => cmd_residuals::(lengths, opts), Command::Verify(lengths) => cmd_verify::(lengths, opts), - Command::Emit(lengths) => cmd_emit::(lengths, opts, element), Command::Dump(lengths) => cmd_dump::(lengths, opts), + Command::Plantime(lengths) => cmd_plantime::(lengths, opts), + Command::Crossover(lengths) => cmd_crossover::(lengths, opts), Command::Score(path) => cmd_score(path, opts), - Command::Explain(spec) => cmd_explain(spec, opts), Command::Costs(path) => cmd_costs(path, opts), - Command::Plantime(l) => cmd_plantime::(l, opts), - Command::Crossover(l) => cmd_crossover::(l, opts), + Command::Explain(spec) => cmd_explain(spec, P::label(), opts), } } -fn dispatch(planner: &str, command: &Command, opts: &Options, el: &str) { +fn dispatch(planner: &str, command: &Command, opts: &Options) { match planner { - "scalar" => run::>(command, opts, el), + "scalar" => run::>(command, opts), // The manifest gives rustfft the SIMD feature matching the target, so architecture // alone decides which of these exists. A tool-crate `feature = ...` cfg would refer to // the tool's own features and always be false. #[cfg(target_arch = "aarch64")] - "neon" => run::>(command, opts, el), + "neon" => run::>(command, opts), #[cfg(target_arch = "x86_64")] - "sse" => run::>(command, opts, el), + "sse" => run::>(command, opts), #[cfg(target_arch = "wasm32")] - "wasm_simd" => run::>(command, opts, el), + "wasm_simd" => run::>(command, opts), other => { - eprintln!( - "unknown or unavailable planner '{}' on this build; try 'scalar'", - other - ); + eprintln!("unknown or unavailable planner '{}' on this build", other); std::process::exit(2); } } } +/// The SIMD planner this build has, which is the default. +fn native_planner() -> &'static str { + if cfg!(target_arch = "aarch64") { + "neon" + } else if cfg!(target_arch = "x86_64") { + "sse" + } else if cfg!(target_arch = "wasm32") { + "wasm_simd" + } else { + "scalar" + } +} + +fn usage() -> ! { + eprintln!("usage: planner_tuning [options] ARGS..."); + eprintln!("commands: time SPEC... | regret LEN... | sweep LEN...|A..B | survey A..B"); + eprintln!(" verify LEN... | dump LEN... | score DUMP | costs DUMP | explain SPEC"); + eprintln!(" plantime LEN... | crossover LEN..."); + eprintln!( + " --planner NAME scalar, neon, sse or wasm_simd (default: this build's SIMD one)" + ); + eprintln!(" --f32 f32 instead of f64"); + eprintln!(" --rounds N timing rounds per subject (default 9)"); + eprintln!(" --block-ms MS wall-clock time per timed block (default 10)"); + eprintln!(" --cap N max candidates per length when enumerating (default 48)"); + eprintln!(" --count N survey: how many lengths (default 300)"); + eprintln!(" --seed N survey: seed for picking lengths (default 1)"); + eprintln!(" --out FILE dump: where to write (default dump.tsv)"); + eprintln!(" --verbose regret: list the top candidates per length"); + eprintln!(" cost model weights, overriding the planner's defaults:"); + eprintln!(" --strided X --permuted X --rader-index X --radixn-extra X"); + eprintln!(" --general-row X --small-row X"); + std::process::exit(2); +} + fn main() { let args: Vec = std::env::args().skip(1).collect(); if args.is_empty() { - eprintln!("usage: planner_tuning [options] ARGS..."); - eprintln!("commands: time SPEC... | regret LEN... | model TRAIN... 0 TEST..."); - eprintln!(" residuals LEN... | verify LEN... | emit LEN..."); - eprintln!(" --planner NAME scalar (default), neon, sse"); - eprintln!(" --rounds N timing rounds per subject (default 9)"); - eprintln!(" --block-ms MS wall-clock time per timed block (default 10)"); - eprintln!(" --cap N max candidates per length (default 48)"); - eprintln!(" --f32 measure f32 instead of f64"); - eprintln!(" --permuted-vector charge permuted passes per vector, not per element"); - eprintln!(" --bucketed fit overhead curves over working set, not constants"); - eprintln!(" --verbose for 'regret', list the top candidates per length"); - std::process::exit(2); + usage(); } let command_name = args[0].clone(); - let mut planner = "scalar".to_string(); + let mut planner = native_planner().to_string(); let mut opts = Options { rounds: 9, block_ms: 10.0, cap: 48, verbose: false, - bucketed: false, out: None, - params: counted::Params::default(), - backend_explicit: false, + weights: Weights::default(), + f32: false, + count: 300, + seed: 1, }; - let mut f32_mode = false; let mut rest: Vec = Vec::new(); let mut i = 1; + let value = |i: &mut usize, name: &str| -> String { + *i += 1; + args.get(*i) + .unwrap_or_else(|| { + eprintln!("{} wants a value", name); + std::process::exit(2); + }) + .clone() + }; + let number = |text: String, name: &str| -> f64 { + text.parse().unwrap_or_else(|_| { + eprintln!("{} wants a number", name); + std::process::exit(2); + }) + }; while i < args.len() { - match args[i].as_str() { - "--planner" => { - i += 1; - planner = args[i].clone(); - } - "--rounds" => { - i += 1; - opts.rounds = args[i].parse().expect("--rounds wants a number"); - } - "--block-ms" => { - i += 1; - opts.block_ms = args[i].parse().expect("--block-ms wants a number"); - } - "--cap" => { - i += 1; - opts.cap = args[i].parse().expect("--cap wants a number"); - } - "--f32" => f32_mode = true, - "--bucketed" => opts.bucketed = true, - "--out" => { - i += 1; - opts.out = Some(args[i].clone()); - } - "--seq-l1" => { i += 1; opts.params.seq[0] = args[i].parse().unwrap(); } - "--seq-l2" => { i += 1; opts.params.seq[1] = args[i].parse().unwrap(); } - "--seq-dram" => { i += 1; opts.params.seq[2] = args[i].parse().unwrap(); } - "--strided" => { i += 1; opts.params.strided_mult = args[i].parse().unwrap(); } - "--permuted" => { i += 1; opts.params.permuted_mult = args[i].parse().unwrap(); } - "--rader-index" => { i += 1; opts.params.rader_index = args[i].parse().unwrap(); } - "--radixn-extra" => { i += 1; opts.params.radixn_extra = args[i].parse().unwrap(); } - "--mul-complex" => { i += 1; opts.params.mul_complex = args[i].parse().unwrap(); } - "--spill" => { i += 1; opts.params.spill = args[i].parse().unwrap(); } - "--general-row" => { i += 1; opts.params.general_row = args[i].parse().unwrap(); } - "--small-row" => { i += 1; opts.params.small_row = args[i].parse().unwrap(); } - "--permuted-vector" => opts.params.permuted_scalar = false, - "--f64" => opts.params.elem = counted::Elem::F64, - "--backend" => { i += 1; opts.params.backend = counted::Backend::parse(&args[i]).expect("--backend wants neon or sse"); opts.backend_explicit = true; } - "--l1-elems" => { i += 1; opts.params.l1_elems = args[i].parse().unwrap(); } - "--l2-elems" => { i += 1; opts.params.l2_elems = args[i].parse().unwrap(); } + let arg = args[i].clone(); + match arg.as_str() { + "--planner" => planner = value(&mut i, &arg), + "--rounds" => opts.rounds = number(value(&mut i, &arg), &arg) as usize, + "--block-ms" => opts.block_ms = number(value(&mut i, &arg), &arg), + "--cap" => opts.cap = number(value(&mut i, &arg), &arg) as usize, + "--count" => opts.count = number(value(&mut i, &arg), &arg) as usize, + "--seed" => opts.seed = number(value(&mut i, &arg), &arg) as u64, + "--out" => opts.out = Some(value(&mut i, &arg)), + "--f32" => opts.f32 = true, "--verbose" => opts.verbose = true, + "--strided" => opts.weights.strided = Some(number(value(&mut i, &arg), &arg)), + "--permuted" => opts.weights.permuted = Some(number(value(&mut i, &arg), &arg)), + "--rader-index" => opts.weights.rader_index = Some(number(value(&mut i, &arg), &arg)), + "--radixn-extra" => opts.weights.radixn_extra = Some(number(value(&mut i, &arg), &arg)), + "--general-row" => opts.weights.general_row = Some(number(value(&mut i, &arg), &arg)), + "--small-row" => opts.weights.small_row = Some(number(value(&mut i, &arg), &arg)), + other if other.starts_with("--") => { + eprintln!("unknown option '{}'", other); + usage(); + } other => rest.push(other.to_string()), } i += 1; } - let numbers = |values: &[String]| -> Vec { - values - .iter() - .map(|s| s.parse().expect("lengths must be numbers")) - .collect() - }; - - // `sweep` takes a thousand lengths, so accept "A..B" as well as a list. - let range_or_numbers = |values: &[String]| -> Vec { + // Lengths may be a list, or `A..B` ranges, which is how a sweep takes a thousand of them. + let lengths = |values: &[String]| -> Vec { let mut out = Vec::new(); for value in values { match value.split_once("..") { @@ -1469,41 +1182,44 @@ fn main() { } out }; + let single = + |values: &[String]| -> String { values.first().cloned().unwrap_or_else(|| usage()) }; let command = match command_name.as_str() { "time" => Command::Time(rest.clone()), - "regret" => Command::Regret(numbers(&rest)), - "sweep" => Command::Sweep(range_or_numbers(&rest)), - "residuals" => Command::Residuals(numbers(&rest)), - "verify" => Command::Verify(numbers(&rest)), - "emit" => Command::Emit(numbers(&rest)), - "dump" => Command::Dump(numbers(&rest)), - "score" => Command::Score(rest[0].clone()), - "explain" => Command::Explain(rest[0].clone()), - "costs" => Command::Costs(rest[0].clone()), - "plantime" => Command::Plantime(numbers(&rest)), - "crossover" => Command::Crossover(numbers(&rest)), - "model" => { - let lengths = numbers(&rest); - let split = lengths - .iter() - .position(|&l| l == 0) - .expect("model wants TRAIN... 0 TEST..."); - let (train, test) = lengths.split_at(split); - Command::Model(train.to_vec(), test[1..].to_vec()) + "regret" => Command::Regret(lengths(&rest)), + "sweep" => Command::Sweep(lengths(&rest)), + "survey" => { + let range = single(&rest); + let (lo, hi) = range.split_once("..").unwrap_or_else(|| usage()); + let (lo, hi) = ( + lo.parse().expect("bad range start"), + hi.parse().expect("bad range end"), + ); + println!( + "# survey\t{} lengths from {}..{}, seed {}", + opts.count, lo, hi, opts.seed + ); + Command::Sweep(random_lengths(lo, hi, opts.count, opts.seed)) } + "verify" => Command::Verify(lengths(&rest)), + "dump" => Command::Dump(lengths(&rest)), + "score" => Command::Score(single(&rest)), + "explain" => Command::Explain(single(&rest)), + "costs" => Command::Costs(single(&rest)), + "plantime" => Command::Plantime(lengths(&rest)), + "crossover" => Command::Crossover(lengths(&rest)), other => { eprintln!("unknown command '{}'", other); - std::process::exit(2); + usage(); } }; request_performance_core(); - if f32_mode { - opts.params.elem = counted::Elem::F32; - dispatch::(&planner, &command, &opts, "f32"); + if opts.f32 { + dispatch::(&planner, &command, &opts); } else { - dispatch::(&planner, &command, &opts, "f64"); + dispatch::(&planner, &command, &opts); } } diff --git a/tools/planner_tuning/src/model.rs b/tools/planner_tuning/src/model.rs deleted file mode 100644 index 6246e104..00000000 --- a/tools/planner_tuning/src/model.rs +++ /dev/null @@ -1,164 +0,0 @@ -//! A cost model for FFT recipes. -//! -//! Table-driven rather than curve-fitted. Each planner has a finite and small set of primitives -//! (a couple of dozen butterflies, a few dozen valid Radix4 shapes), so their costs are simply -//! measured and stored. That removes the extrapolation error a fitted closed form introduces, -//! which is what made the 2021 scalar attempt mis-rank a direct Radix4 against a split one. -//! -//! Only the composing algorithms need fitted numbers, and each needs one: the cost per element -//! of the transposes and twiddle multiplies they add on top of their inner FFTs. - -use rustfft::tuning::Spec; -use std::collections::HashMap; - -/// What an algorithm's per-element overhead scales with. -/// -/// Usually its own length. Two exceptions: Bluestein's pointwise multiply and zero-padding run -/// over the padded inner length, which can be nearly four times the outer length; and RadixN -/// makes one pass per factor, so its overhead scales with length times the number of levels. -pub fn overhead_scale(spec: &Spec) -> f64 { - match spec { - Spec::Bluesteins { inner, .. } => inner.len() as f64, - Spec::RadixN { radixes, .. } => (spec.len() * radixes.len()) as f64, - other => other.len() as f64, - } -} - -/// Which log2 bucket a spec's overhead belongs in, bucketed on the same quantity the overhead is -/// charged per. -pub fn bucket_of(spec: &Spec) -> u32 { - overhead_scale(spec).log2() as u32 -} - -#[derive(Default, Clone)] -pub struct Model { - /// Measured nanoseconds for one FFT, by butterfly length. - pub butterfly: HashMap, - /// Measured nanoseconds for one FFT, by (base length, k). - pub radix4: HashMap<(usize, u32), f64>, - /// Fitted nanoseconds per element of overhead, by algorithm, as a curve over log2 of the - /// working set. Entries are sorted by bucket. A single entry means a flat constant. - pub overhead: HashMap<&'static str, Vec<(u32, f64)>>, -} - -impl Model { - /// Overhead per element at a working set of `scale` elements, linearly interpolated between - /// measured buckets and clamped outside the measured range. - pub fn overhead_at(&self, kind: &str, scale: f64) -> Option { - let table = self.overhead.get(kind)?; - match table.len() { - 0 => None, - 1 => Some(table[0].1), - _ => { - let x = scale.log2(); - if x <= table[0].0 as f64 { - return Some(table[0].1); - } - if x >= table[table.len() - 1].0 as f64 { - return Some(table[table.len() - 1].1); - } - for pair in table.windows(2) { - let (lo_bucket, lo_value) = pair[0]; - let (hi_bucket, hi_value) = pair[1]; - if x <= hi_bucket as f64 { - let span = (hi_bucket - lo_bucket) as f64; - let t = if span > 0.0 { - (x - lo_bucket as f64) / span - } else { - 0.0 - }; - return Some(lo_value + t * (hi_value - lo_value)); - } - } - Some(table[table.len() - 1].1) - } - } - } - - /// Estimated nanoseconds for one FFT of this recipe. - /// - /// `None` if the recipe needs a primitive that was never measured or an overhead that was - /// never fitted, so a gap is a loud failure rather than a silently wrong ranking. - pub fn cost(&self, spec: &Spec) -> Option { - Some(match spec { - Spec::Dft(n) => { - // Only reached for degenerate sizes; a quadratic keeps it last. - let n = *n as f64; - 100.0 * n * n - } - Spec::Butterfly(len) => *self.butterfly.get(len)?, - Spec::Radix4 { k, base } => *self.radix4.get(&(base.len(), *k))?, - Spec::RadixN { radixes, base } => { - let scale = overhead_scale(spec); - radixes.iter().product::() as f64 * self.cost(base)? - + self.overhead_at("rn", scale)? * scale - } - Spec::MixedRadix { left, right, .. } | Spec::GoodThomas { left, right, .. } => { - let scale = overhead_scale(spec); - right.len() as f64 * self.cost(left)? - + left.len() as f64 * self.cost(right)? - + self.overhead_at(spec.kind(), scale)? * scale - } - // Rader's runs its inner FFT twice per transform, once forward and once to invert - // the convolution, exactly like Bluestein's. - Spec::Raders { inner } => { - let scale = overhead_scale(spec); - 2.0 * self.cost(inner)? + self.overhead_at("rad", scale)? * scale - } - Spec::Bluesteins { inner, .. } => { - let scale = overhead_scale(spec); - 2.0 * self.cost(inner)? + self.overhead_at("bs", scale)? * scale - } - }) - } - - /// The sum of the inner FFT costs only, with no overhead for `spec` itself. Subtracting this - /// from a measurement is what isolates one algorithm's overhead. - pub fn inner_cost(&self, spec: &Spec) -> Option { - Some(match spec { - Spec::MixedRadix { left, right, .. } | Spec::GoodThomas { left, right, .. } => { - right.len() as f64 * self.cost(left)? + left.len() as f64 * self.cost(right)? - } - Spec::RadixN { radixes, base } => { - radixes.iter().product::() as f64 * self.cost(base)? - } - Spec::Raders { inner } => 2.0 * self.cost(inner)?, - Spec::Bluesteins { inner, .. } => 2.0 * self.cost(inner)?, - other => self.cost(other)?, - }) - } - - /// True if every strict descendant is a primitive or an already-fitted kind. - /// - /// Overheads are fitted in dependency order, because the residual for a MixedRadix - /// containing a MixedRadixSmall only means anything once the Small's overhead is known. - pub fn descendants_known(&self, spec: &Spec, fitted: &[&'static str]) -> bool { - spec.children().iter().all(|child| { - let k = child.kind(); - let ok = matches!(k, "butterfly" | "r4" | "dft") || fitted.contains(&k); - ok && self.descendants_known(child, fitted) - }) - } - - pub fn describe(&self) -> String { - let mut kinds: Vec<(&&str, &Vec<(u32, f64)>)> = self.overhead.iter().collect(); - kinds.sort_by_key(|(k, _)| **k); - let mut out = format!( - "{} butterflies, {} radix4 shapes measured\n", - self.butterfly.len(), - self.radix4.len() - ); - out.push_str("overhead, ns per element, by working set:\n"); - for (kind, table) in kinds { - let rendered: Vec = table - .iter() - .map(|(bucket, value)| format!("{}:{:.2}", 1usize << bucket, value)) - .collect(); - out.push_str(&format!(" {:<5} {}\n", kind, rendered.join(" "))); - } - out - } -} - -/// The order overheads must be fitted in, so that each kind's inners are already known. -pub const FIT_ORDER: [&str; 7] = ["mrs", "gts", "mr", "gt", "rn", "rad", "bs"]; From 6bbfc8054b4878a69903fe5e9f6a8d7a2651fc79 Mon Sep 17 00:00:00 2001 From: Henrik Date: Wed, 16 Sep 2026 22:57:00 +0200 Subject: [PATCH 15/22] Rewrite the tuning harness README for the library-driven commands --- tools/planner_tuning/README.md | 169 ++++++++++++++------------------- 1 file changed, 73 insertions(+), 96 deletions(-) diff --git a/tools/planner_tuning/README.md b/tools/planner_tuning/README.md index c4260269..2f974c0d 100644 --- a/tools/planner_tuning/README.md +++ b/tools/planner_tuning/README.md @@ -1,11 +1,12 @@ # planner_tuning -A measurement harness for RustFFT's planners. Not part of the library and not shipped: it exists -to answer "which recipe is actually fastest at this length, and does the planner pick it?" +A measurement harness for RustFFT's SIMD planners. Not part of the library and not shipped: it +exists to fit the estimating planner's cost model, and to check how its picks compare against the +fixed planner it replaces and against the fastest recipe available. -Everything here works against a `TunablePlanner`, so the same commands run on the scalar, NEON, -SSE and wasm planners. Recipes are built through the planner's own internals, so what gets timed -is exactly what the planner would construct. +It drives the planners through the library's non-default `tuning` feature, so recipes are built +by the planners' own code and priced by the library's own cost model. What gets timed is exactly +what a planner would construct. ## Vocabulary @@ -14,135 +15,111 @@ butterfly 11 against Rader's over a MixedRadixSmall of 6 and 10. The prefixes ar `r4` Radix4, `rn` RadixN, `mr`/`mrs` MixedRadix and its Small variant, `gt`/`gts` Good-Thomas and its Small variant, `rad` Rader's, `bs` Bluestein's. +**planner** in the output is the fixed planner, which the SIMD planners keep behind the `tuning` +feature while the estimating planner is a draft. **model** is the estimating planner. + **Regret** is a pick divided by the best measured candidate at that length: 1.000 is optimal, 1.05 is five percent off. It is a *lower bound* on the distance from optimal, because the candidate set is finite. -## The two workflows +## Comparing the two planners -### Measure once, replay forever +`sweep` times the fixed planner's pick against the estimating planner's at every length given. +It builds only those two recipes per length, which is what makes a thousand lengths affordable. +Where the two agree the recipe is timed once and reported for both, so agreement reads as exactly +1.000 rather than as noise. -`dump` times every candidate at each length and writes a TSV. `score`, `costs` and `explain` then -replay that file offline, so iterating on the cost model needs no machine after the first run. -This is how the weights were fitted. +`survey` is the same over lengths drawn at random from a range, which is how to cover sizes up +to a million. Both end with a summary of the estimating planner's runtime over the fixed +planner's: geometric mean and percentiles, so below 1 means faster. ```sh cargo build --release -./target/release/planner_tuning dump --planner neon --rounds 7 --cap 48 \ - --out dump_neon_f64.tsv 1000 1050 1200 1296 -./target/release/planner_tuning score dump_neon_f64.tsv # regret vs best measured -./target/release/planner_tuning costs dump_neon_f64.tsv # per-candidate cost and ns -./target/release/planner_tuning explain 'rad(rn(7.6,b24))' # cost tree for one recipe +./target/release/planner_tuning sweep 1..1000 > sweep_neon_f64.tsv +./target/release/planner_tuning sweep --f32 1..1000 > sweep_neon_f32.tsv +./target/release/planner_tuning survey --count 300 --seed 1 1..1000000 > survey_neon_f64.tsv ``` -### Sweep a whole range +The planner defaults to the SIMD one this build has: neon on aarch64, sse on x86-64. Lengths may +be a list or an `A..B` range. -`sweep` times the planner's pick against the cost model's pick at every length in a range. Unlike -`dump` it builds only those two recipes per length, which is what makes a thousand lengths -affordable. Where the two agree the recipe is timed once and reported for both, so agreement -reads as exactly 1.000 rather than as noise. +## Fitting weights: measure once, replay forever -This is the instrument that matters. Both cost-model defects fixed so far were invisible to the -33- and 44-length tuning sets and were found only by sweeping 1..1000. +`dump` times every candidate at each length and writes a TSV. `score`, `costs` and `explain` then +replay that file offline, so iterating on the cost model needs no machine after the first run. ```sh -./target/release/planner_tuning sweep --planner neon 1..1000 > sweep_neon_f64.tsv -./target/release/planner_tuning sweep --planner neon --f32 1..1000 > sweep_neon_f32.tsv +./target/release/planner_tuning dump --rounds 7 --out dump_neon_f64.tsv 1000 1050 1200 1296 +./target/release/planner_tuning score dump_neon_f64.tsv # regret vs best measured +./target/release/planner_tuning costs dump_neon_f64.tsv # per-candidate cost and ns +./target/release/planner_tuning explain 'rad(rn(7.6,b24))' # cost tree for one recipe ``` -Lengths may be a list or an `A..B` range. +The candidates in a dump are enumerated exhaustively, one level deep, with inner recipes from the +fixed planner. So `score` replays a one-level choice, which is how the weights were fitted, and +not quite what the estimating planner does: it also optimises the inner recipes. `sweep` and +`survey` measure the real thing. + +Every weight can be overridden on the command line, for `score` as well as for `sweep`: +`--strided`, `--permuted`, `--rader-index`, `--radixn-extra`, `--general-row`, `--small-row`. +The defaults per backend and element type are in `src/simd/simd_estimate.rs`. + +**Pick weights by the worst machine, not the best.** Weights ship as compile-time constants per +backend and element type, so a value fitted to one machine is a regression on another. When a +weight's optimum differs between machines, sweep it on all of them and take the value whose worst +machine looks best. + +## Other commands + +- `regret LEN...` times every candidate plus the estimating planner's pick, and reports how far + both planners are from the best. `--verbose` lists the top candidates. +- `verify LEN...` checks every candidate, and the estimating planner's pick, against a direct + DFT. Run it after any change to candidate enumeration. +- `plantime LEN...` compares the cost of planning with each planner against the cost of building + the recipe. These are cold-start figures: a planner reused across lengths shares inner lengths. +- `crossover LEN...` asks which recipe would win if construction cost counted, and how many + executions the estimating planner's pick needs to repay any extra build cost. +- `time SPEC...` times recipes against each other. A `*reps` suffix runs a recipe over a + `len*reps` buffer, which is how an inner FFT is invoked. ## Plotting -`plot_sweep.py` draws the two views worth having: the population normalised by `N log2 N`, and -the per-length ratio. It prints the same figures it draws, so it doubles as the reporting tool. -It takes several files at once, which is how cost-model stages and machines get compared. - -matplotlib is the only third-party dependency in this directory, so it lives in a venv: +`plot_sweep.py` draws the population normalised by `N log2 N`, and the per-length ratio, for one +or several sweep files. It prints the same figures it draws. matplotlib is the only dependency: ```sh python3 -m venv .venv && .venv/bin/pip install matplotlib .venv/bin/python plot_sweep.py sweep_neon_f64.tsv # one run -.venv/bin/python plot_sweep.py sweep_*_f64_*.tsv # compare stages or machines -.venv/bin/python plot_sweep.py sweep_neon_f64.tsv --lo 4 --hi 128 # zoom +.venv/bin/python plot_sweep.py sweep_*_f64_*.tsv # compare runs or machines .venv/bin/python plot_sweep.py sweep_neon_f64.tsv --out fig.png # write instead of show ``` -## Plan time +## Other backends -`plantime` compares the cost of *choosing* a recipe against the cost of *building* it. Planning -only produces a `Recipe`; turning that into an `Arc` allocates and precomputes, and -Rader's and Bluestein's both run a full inner FFT inside their constructors. Building is 2x to -900x planning and grows much faster with length, so the estimating planner's overhead is a large -fraction of plan-plus-build at small lengths and a small one at large lengths. +On x86-64 the tool crate selects the `sse` feature, so a plain `cargo build` gives the SSE +planner. Do not build the library with its default features for this: the planner would pick +AVX and measure the wrong backend. -```sh -./target/release/planner_tuning plantime --planner neon --cap 48 128 1024 1200 10007 -``` - -Single plan-time measurements are noisy, up to 1.6x apart at length 1200. Take medians of three. -The candidate counts it prints are exact. - -`crossover` goes further and asks which recipe would win if construction cost counted: it builds and -times every candidate, then reports the minimiser of `build + k * execute` at several `k`, plus how -many executions the cost model's pick needs to repay its extra build cost. Measured answer is three -to six, so this mostly documents why a construction-aware planner is not worth building. - -```sh -./target/release/planner_tuning crossover --planner neon --cap 24 1260 1009 2018 -``` - -## Two candidate sets, deliberately - -- `candidates` / `candidates_capped` are **exhaustive**: every split in both orders, every - algorithm that can express it, Bluestein's at every length. Used by `dump`, `verify` and - `regret`, because scoring a planner's pick needs the alternatives even where no planner would - look at them. -- `plan_candidates` is **what an estimating planner would really enumerate**. Used by `sweep` and - `plantime`. It skips work that measurement has shown carries no decision: butterfly lengths and - powers of two return the fixed planner's pick immediately, only the smaller-width-first ordering - of each split is emitted, and Bluestein's is offered only where some prime factor has no - butterfly of its own. - -Keep the first set exhaustive. It is how each of those shortcuts was justified, and re-justifying -them after a kernel change needs the same breadth. - -## Running on another machine - -The campaign is SSH-driven. The remote checkouts are rsync copies rather than git worktrees, so -their `.git` points at a path that does not exist there; sync source only and leave the data. +wasm_simd builds for `wasm32-wasip1`, not `wasm32-unknown-unknown`, which has no clock. Run it +under node with `run_wasm.mjs`, and always with long blocks (`--block-ms 100`): V8 starts on its +baseline compiler and tiers up later, so short blocks time unoptimised code. ```sh -rsync -az --delete src/ user@host:~/repos/RustFFT-testing/src/ -rsync -az --delete tools/planner_tuning/src/ \ - user@host:~/repos/RustFFT-testing/tools/planner_tuning/src/ -ssh user@host 'cd ~/repos/RustFFT-testing/tools/planner_tuning && cargo build --release' +RUSTFLAGS="-C target-feature=+simd128" cargo build --release --target wasm32-wasip1 +node run_wasm.mjs sweep --planner wasm_simd --block-ms 100 1..200 ``` -On x86 the tuning crate selects the `sse` feature automatically, so a plain `cargo build` gives -the SSE planner. - ## Traps -- **`sweep` and `plantime` do not infer the backend.** `score` and `costs` read it from the dump - header, but `sweep` takes the model parameters as given, so an SSE run needs `--backend sse` - explicitly. Without it the model prices SSE recipes with NEON instruction counts and the whole - run looks plausible and is meaningless. -- **`--f32` is needed on `score` as well as on `dump`.** The dump header records the planner, so - the backend is recovered, but it does not record the element type. +- **`--f32` is needed on `score` and `costs` as well as on `dump`.** The dump header records the + planner, so the backend is recovered, but it does not record the element type. - **Do not compile while measuring.** Build first, then run. -- **zsh does not word-split unquoted variables**, so `$LENGTHS` arrives as one argument and the - tool panics. Inline the numbers or use `${=VAR}`. -- **Run `verify` after any change to candidate enumeration.** It checks every enumerated candidate - against a direct DFT, which catches an illegal spec such as a Bluestein's inner shorter than - `2n - 1`. +- **zsh does not word-split unquoted variables**, so `$LENGTHS` arrives as one argument. Inline + the numbers or use `${=VAR}`. - Measurement output (`*.tsv`, `*.txt`, `*.log`) and `.venv` are gitignored. ## The documents -- `COST-MODEL.md` is how the estimating planner works: what a cost is, how one is computed, which - numbers are counted and which are fitted, and what has to be redone when a kernel changes. Start - here. -- `RESULTS.md` is the evidence: how the cost model was built and what it scores. +- `COST-MODEL.md` is how the cost model works: what a cost is, how one is computed, which numbers + are counted and which are fitted, and what has to be redone when a kernel changes. - `OP-COUNTS.md` is where the instruction counts come from. -- `NEXT-STEPS.md` is the live plan, and the place new findings get written down. From 3418e8871e6f30adc3f01b36259214a19322d475 Mon Sep 17 00:00:00 2001 From: Henrik Date: Wed, 16 Sep 2026 23:04:06 +0200 Subject: [PATCH 16/22] Refit rader_index after the Rader's permutation became a table A survey of 300 random lengths up to 1M on the M1 moves the estimating planner's runtime over the fixed planner's from geomean 0.950, p90 1.163 to 0.883, p90 1.015 in f64, and from 1.016, p90 1.260 to 0.924, p90 1.062 in f32. --- src/simd/simd_estimate.rs | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/simd/simd_estimate.rs b/src/simd/simd_estimate.rs index d2ba74f1..df627d6d 100644 --- a/src/simd/simd_estimate.rs +++ b/src/simd/simd_estimate.rs @@ -225,10 +225,13 @@ pub struct CostModel { pub strided: f64, /// A gather or scatter against a sequential pass. pub permuted: f64, - /// One element of Rader's permutation passes, on top of the gather or scatter itself. + /// One element of Rader's permutation passes, on top of the gather or scatter itself: the + /// load of its index from the precomputed `u32` table. /// - /// This was fitted as the latency of a loop-carried modular-multiply chain, which ejmahler#178 - /// has since replaced with a precomputed table. It needs refitting. + /// Before ejmahler#178 the index came from a loop-carried modular-multiply chain, latency + /// bound, and this was 30 (f64) and 45 (f32). With the table, a survey of 300 random lengths + /// up to a million on an M1 scores 2 and 8 about the same, and both far better than the old + /// values, which made Rader's look expensive enough to trade for a large Bluestein's. pub rader_index: f64, /// Extra cost per element per cross-FFT layer of the generic `SimdRadixN` driver, over the /// hand-written `Radix4` kernel doing the same work. Expected near zero on NEON's 32 vector @@ -256,7 +259,7 @@ impl CostModel { complex_per_vector, strided, permuted: 2.5, - rader_index: if f32 { 45.0 } else { 30.0 }, + rader_index: 2.0, radixn_extra, general_row: 30.0, small_row: 10.0, From 283fd6ad407b4b5a62cbcc46d4442b3b385054c1 Mon Sep 17 00:00:00 2001 From: Henrik Date: Fri, 18 Sep 2026 18:17:55 +0200 Subject: [PATCH 17/22] Charge a transpose that no longer fits in cache The cost model priced a MixedRadix transpose of a large buffer the same as a RadixN cross layer, which gathers its rows from inside the chunk it is already working on and so keeps its locality at any size. Only the transposes are charged, and only above a working set of 256 KiB, which is the smallest last-level cache worth planning for: no pick below length 16385 moves, so the weights fitted by sweeping 1 to 1000 are untouched. On 300 random lengths up to 1M on an M1, against the fixed planner: f32 worst goes from 1.485 to 1.178 and f64 p90 from 1.015 to 1.000, at unchanged geometric means of 0.93 and 0.88. --- src/simd/simd_estimate.rs | 81 ++++++++++++++++++++++++-------- tools/planner_tuning/src/main.rs | 43 ++++++++++++++++- 2 files changed, 103 insertions(+), 21 deletions(-) diff --git a/src/simd/simd_estimate.rs b/src/simd/simd_estimate.rs index df627d6d..e66bdb1e 100644 --- a/src/simd/simd_estimate.rs +++ b/src/simd/simd_estimate.rs @@ -243,6 +243,24 @@ pub struct CostModel { /// One outer iteration of `transpose_small`, charged to the `Small` variants for the width /// by which the pair is wider than it is tall. A tie-break between the two orderings. pub small_row: f64, + /// Complex numbers that still fit in cache. A transpose of more than this many is charged + /// `dram` per access on top of its pattern. Infinite prices every transpose the same. + /// + /// A node is charged at the size of the buffer it walks itself, never at the size of the + /// transform it sits inside, so a recipe's cost stays a function of its own subtree and can + /// be cached by length. + /// + /// The default is 256 KiB worth of complex numbers, which is the smallest last-level cache + /// worth planning for rather than any particular machine's. A step is what the hardware + /// does, and unlike a cost that grows smoothly with length it leaves every shorter length + /// priced exactly as before, so the weights fitted by sweeping 1 to 1000 stay valid: no pick + /// below length 16385 changes at this threshold. Lowering it to 128 KiB starts moving picks + /// from 8194, raising it to 2 MiB moves none below 20000 and is worse in the f32 tail. + pub cache_elems: f64, + /// What one access of a transpose costs once the buffer no longer fits in cache, relative to + /// the same access in cache. Anything from 2 to 5 scores the same, so this is an order of + /// magnitude rather than a fitted value. + pub dram: f64, } impl CostModel { @@ -263,6 +281,8 @@ impl CostModel { radixn_extra, general_row: 30.0, small_row: 10.0, + cache_elems: 256.0 * 1024.0 / 16.0 * complex_per_vector as f64, + dram: 2.0, } } @@ -271,14 +291,33 @@ impl CostModel { Self::new(instruction_set, complex_per_vector::()) } - /// Cost of touching `accesses` elements, counting each load and each store once. - fn mem(&self, accesses: f64, pattern: Pattern) -> f64 { - match pattern { + /// Cost of touching `accesses` elements, counting each load and each store once, in a pass + /// over a buffer of `ws` complex numbers. + fn mem(&self, accesses: f64, pattern: Pattern, ws: f64) -> f64 { + let cost = match pattern { // A gather or scatter computes an address per complex number, so it cannot fill a // vector. Pattern::Permuted => accesses * self.permuted, Pattern::Strided => accesses / self.complex_per_vector as f64 * self.strided, Pattern::Sequential => accesses / self.complex_per_vector as f64, + }; + let _ = ws; + cost + } + + /// Cost of a pass that transposes the whole buffer, which is the one access pattern that + /// really falls out of cache. + /// + /// A RadixN or Radix4 cross layer looks strided but gathers its rows from inside the chunk it + /// is already working on, so it keeps its locality at any size. A MixedRadix or GoodThomas + /// transpose walks the whole rectangle, so once that no longer fits in cache every element + /// costs a fresh line. + fn transpose_mem(&self, accesses: f64, pattern: Pattern, ws: f64) -> f64 { + let cost = self.mem(accesses, pattern, ws); + if ws > self.cache_elems { + cost * self.dram + } else { + cost } } @@ -300,7 +339,8 @@ impl CostModel { let cpv = self.complex_per_vector as f64; Some(match shape { Shape::Butterfly(len) => { - self.butterfly(*len)? + self.mem(2.0 * *len as f64, Pattern::Sequential) + self.butterfly(*len)? + + self.mem(2.0 * *len as f64, Pattern::Sequential, *len as f64) } Shape::Radix4 { k, base_len } => { let len = (*base_len << (2 * k)) as f64; @@ -308,18 +348,18 @@ impl CostModel { // One digit-reversal transpose, then the base FFTs, then k cross layers of // len/4 column_butterfly4, each four butterfly2 plus a rotate90 (10 instructions) // and three twiddle multiplies. - let mut c = self.mem(2.0 * len, Pattern::Permuted); + let mut c = self.mem(2.0 * len, Pattern::Permuted, len); c += reps * child_cost(*base_len); c += *k as f64 * ((len / (4.0 * cpv)) * (10.0 + 3.0 * self.mul_complex()) - + self.mem(2.0 * len, Pattern::Strided)); + + self.mem(2.0 * len, Pattern::Strided, len)); c } Shape::RadixN { factors, base_len } => { let radix_product: usize = factors.iter().map(|f| f.radix()).product(); let len = (*base_len * radix_product) as f64; let reps = len / *base_len as f64; - let mut c = self.mem(2.0 * len, Pattern::Permuted); + let mut c = self.mem(2.0 * len, Pattern::Permuted, len); c += reps * child_cost(*base_len); for factor in factors.iter() { let radix = factor.radix(); @@ -327,7 +367,7 @@ impl CostModel { // table applies directly. Row 0 needs no twiddle, hence radix - 1. c += (len / radix as f64) * (self.butterfly(radix)? + (radix as f64 - 1.0) * self.mul_complex()); - c += self.mem(2.0 * len, Pattern::Strided); + c += self.mem(2.0 * len, Pattern::Strided, len); c += len * self.radixn_extra; } c @@ -345,13 +385,14 @@ impl CostModel { // `MixedRadix` hands the job to the `transpose` crate, which tiles the rectangle // to get cache reuse back and pays `general_row` per row for it. let mut c = if *small { - 3.0 * self.mem(2.0 * len, Pattern::Permuted) + 3.0 * self.transpose_mem(2.0 * len, Pattern::Permuted, len) + self.small_row * left_len.saturating_sub(*right_len) as f64 } else { - 3.0 * self.mem(2.0 * len, Pattern::Strided) + 3.0 * self.transpose_mem(2.0 * len, Pattern::Strided, len) + self.general_row * 1.5 * (left_len + right_len) as f64 }; - c += (len / cpv) * self.mul_complex() + self.mem(2.0 * len, Pattern::Sequential); + c += (len / cpv) * self.mul_complex() + + self.mem(2.0 * len, Pattern::Sequential, len); c += *right_len as f64 * child_cost(*left_len); c += *left_len as f64 * child_cost(*right_len); c @@ -365,12 +406,12 @@ impl CostModel { // Two CRT reindexing passes and one transpose, but no twiddle multiplies. Both // reindexing passes are permuted in either variant; the transpose splits the two // exactly as in MixedRadix. - let mut c = 2.0 * self.mem(2.0 * len, Pattern::Permuted); + let mut c = 2.0 * self.transpose_mem(2.0 * len, Pattern::Permuted, len); c += if *small { - self.mem(2.0 * len, Pattern::Permuted) + self.transpose_mem(2.0 * len, Pattern::Permuted, len) + self.small_row * left_len.saturating_sub(*right_len) as f64 } else { - self.mem(2.0 * len, Pattern::Strided) + self.transpose_mem(2.0 * len, Pattern::Strided, len) + self.general_row * 1.5 * (left_len + right_len) as f64 }; c += *right_len as f64 * child_cost(*left_len); @@ -381,8 +422,9 @@ impl CostModel { let len_f = *len as f64; // The inner FFT runs twice, around a permuting gather and scatter. let mut c = 2.0 * child_cost(len - 1); - c += 2.0 * (self.mem(2.0 * len_f, Pattern::Permuted) + len_f * self.rader_index); - c += len_f * self.mul_complex() + self.mem(2.0 * len_f, Pattern::Sequential); + c += 2.0 + * (self.mem(2.0 * len_f, Pattern::Permuted, len_f) + len_f * self.rader_index); + c += len_f * self.mul_complex() + self.mem(2.0 * len_f, Pattern::Sequential, len_f); c } Shape::Bluesteins { len, inner_len } => { @@ -391,9 +433,10 @@ impl CostModel { // Inner FFT twice, pointwise multiply over the padded inner length, and a // twiddle-and-pad pass in and out over the outer length. let mut c = 2.0 * child_cost(*inner_len); - c += inner * self.mul_complex() + self.mem(2.0 * inner, Pattern::Sequential); - c += - 2.0 * (outer * self.mul_complex() + self.mem(2.0 * outer, Pattern::Sequential)); + c += inner * self.mul_complex() + self.mem(2.0 * inner, Pattern::Sequential, inner); + c += 2.0 + * (outer * self.mul_complex() + + self.mem(2.0 * outer, Pattern::Sequential, outer)); c } }) diff --git a/tools/planner_tuning/src/main.rs b/tools/planner_tuning/src/main.rs index 95e367e0..85fe649a 100644 --- a/tools/planner_tuning/src/main.rs +++ b/tools/planner_tuning/src/main.rs @@ -139,6 +139,9 @@ struct Weights { radixn_extra: Option, general_row: Option, small_row: Option, + /// Cache size in KiB, converted to complex numbers for whichever element type is in use. + cache_kib: Option, + dram: Option, } impl Weights { @@ -154,6 +157,14 @@ impl Weights { set(&mut model.radixn_extra, self.radixn_extra); set(&mut model.general_row, self.general_row); set(&mut model.small_row, self.small_row); + // A complex f64 is 16 bytes and a complex f32 is 8, so the same cache holds twice as + // many of the latter. + set( + &mut model.cache_elems, + self.cache_kib + .map(|kib| kib * 1024.0 / 16.0 * model.complex_per_vector as f64), + ); + set(&mut model.dram, self.dram); model } } @@ -379,6 +390,28 @@ fn candidates_with_model_pick>( (specs, model_index) } +/// Print what each planner picks at every length, without timing anything. +/// +/// Fast enough for tens of thousands of lengths, so two cost models can be diffed by their +/// decisions before spending a machine on measuring the ones that differ. +fn cmd_picks>(lengths: &[usize], opts: &Options) { + let mut planner = tuner::(opts); + println!("len\tagree\tplanner_spec\tmodel_spec"); + for &len in lengths { + let planner_spec = planner.plan(len); + let model_spec = planner + .estimate(len) + .unwrap_or_else(|| Arc::clone(&planner_spec)); + println!( + "{}\t{}\t{}\t{}", + len, + if planner_spec == model_spec { 1 } else { 0 }, + to_spec_string(&planner_spec), + to_spec_string(&model_spec) + ); + } +} + fn cmd_regret>(lengths: &[usize], opts: &Options) { println!( "{:>9} {:>8} {:>8} {:>10} {:>10} {}", @@ -1022,6 +1055,7 @@ fn cmd_crossover>(lengths: &[usize], opts: &Opti // --------------------------------------------------------------------------- enum Command { + Picks(Vec), Time(Vec), Regret(Vec), Sweep(Vec), @@ -1036,6 +1070,7 @@ enum Command { fn run>(command: &Command, opts: &Options) { match command { + Command::Picks(lengths) => cmd_picks::(lengths, opts), Command::Time(specs) => cmd_time::(specs, opts), Command::Regret(lengths) => cmd_regret::(lengths, opts), Command::Sweep(lengths) => cmd_sweep::(lengths, opts), @@ -1083,7 +1118,8 @@ fn native_planner() -> &'static str { fn usage() -> ! { eprintln!("usage: planner_tuning [options] ARGS..."); - eprintln!("commands: time SPEC... | regret LEN... | sweep LEN...|A..B | survey A..B"); + eprintln!("commands: picks LEN... | time SPEC... | regret LEN... | sweep LEN...|A..B"); + eprintln!(" survey A..B"); eprintln!(" verify LEN... | dump LEN... | score DUMP | costs DUMP | explain SPEC"); eprintln!(" plantime LEN... | crossover LEN..."); eprintln!( @@ -1099,7 +1135,7 @@ fn usage() -> ! { eprintln!(" --verbose regret: list the top candidates per length"); eprintln!(" cost model weights, overriding the planner's defaults:"); eprintln!(" --strided X --permuted X --rader-index X --radixn-extra X"); - eprintln!(" --general-row X --small-row X"); + eprintln!(" --general-row X --small-row X --cache-kib X --dram X"); std::process::exit(2); } @@ -1158,6 +1194,8 @@ fn main() { "--radixn-extra" => opts.weights.radixn_extra = Some(number(value(&mut i, &arg), &arg)), "--general-row" => opts.weights.general_row = Some(number(value(&mut i, &arg), &arg)), "--small-row" => opts.weights.small_row = Some(number(value(&mut i, &arg), &arg)), + "--cache-kib" => opts.weights.cache_kib = Some(number(value(&mut i, &arg), &arg)), + "--dram" => opts.weights.dram = Some(number(value(&mut i, &arg), &arg)), other if other.starts_with("--") => { eprintln!("unknown option '{}'", other); usage(); @@ -1186,6 +1224,7 @@ fn main() { |values: &[String]| -> String { values.first().cloned().unwrap_or_else(|| usage()) }; let command = match command_name.as_str() { + "picks" => Command::Picks(lengths(&rest)), "time" => Command::Time(rest.clone()), "regret" => Command::Regret(lengths(&rest)), "sweep" => Command::Sweep(lengths(&rest)), From 50edf18bf40f5c766b50f500fda56c684728eb55 Mon Sep 17 00:00:00 2001 From: Henrik Date: Fri, 18 Sep 2026 18:56:02 +0200 Subject: [PATCH 18/22] Charge memory that no longer fits in cache, and RadixN per call Two terms, fitted on a Raspberry Pi 5 and checked on an M1, since they are the first weights in this model whose optimum differs by machine. Above a working set of 256 KiB an access costs dram_pass, and a transpose's costs dram. They move together: penalising ordinary passes alone makes a MixedRadix wrapped around a smaller radix recipe look good, and those measure worse on both machines. On the Pi this removes losses of 4.19x in f64 and 3.81x in f32, where the model had been computing a whole transform as one Bluestein's with an inner FFT far larger than cache. radix_call charges what a RadixN or Radix4 execution costs regardless of length: the call, the scratch split, the layer setup and the virtual call into the base FFT. Nothing charged it on NEON, where radixn_extra is zero, so lengths 14 and 21 took a RadixN measuring 1.32x slower. --- src/simd/simd_estimate.rs | 50 +++++++++-- tools/planner_tuning/src/main.rs | 150 ++++++++++++++++++++++++++++--- 2 files changed, 181 insertions(+), 19 deletions(-) diff --git a/src/simd/simd_estimate.rs b/src/simd/simd_estimate.rs index e66bdb1e..6adb62a5 100644 --- a/src/simd/simd_estimate.rs +++ b/src/simd/simd_estimate.rs @@ -243,6 +243,15 @@ pub struct CostModel { /// One outer iteration of `transpose_small`, charged to the `Small` variants for the width /// by which the pair is wider than it is tall. A tie-break between the two orderings. pub small_row: f64, + /// Charged once per `SimdRadixN` or `Radix4` execution, for the work that does not scale + /// with the length: the call itself, splitting the scratch, walking the factor list, setting + /// up each layer's twiddle slice, and the virtual call into the base FFT. + /// + /// It only matters at short lengths, where it is the difference between a generic driver and + /// a table-driven `GoodThomasAlgorithmSmall`: at 100, which is about 8 ns or 25 cycles on an + /// M1, lengths 14 and 21 stop preferring a RadixN that measures 1.32x slower, while the wins + /// at 40, 150, 241 and 321 are untouched. Nothing above about 500 notices it. + pub radix_call: f64, /// Complex numbers that still fit in cache. A transpose of more than this many is charged /// `dram` per access on top of its pattern. Infinite prices every transpose the same. /// @@ -258,9 +267,27 @@ pub struct CostModel { /// from 8194, raising it to 2 MiB moves none below 20000 and is worse in the f32 tail. pub cache_elems: f64, /// What one access of a transpose costs once the buffer no longer fits in cache, relative to - /// the same access in cache. Anything from 2 to 5 scores the same, so this is an order of - /// magnitude rather than a fitted value. + /// the same access in cache. + /// + /// This one and `dram_pass` have to move together. Charging ordinary passes without charging + /// transposes just as hard makes a MixedRadix wrapped around a smaller radix recipe look + /// good, because its inner FFTs are then cache resident, and those recipes measure worse on + /// both machines: at `dram_pass` 2 with this at 2, the M1 has 7 losses beyond 5% in f64 and + /// 9 in f32, which raising this to 4 removes entirely. 4 and 6 score the same. pub dram: f64, + /// The same, for every other pass: a RadixN or Radix4 cross layer, a Rader's permutation, or + /// Bluestein's padded multiply. + /// + /// Separate from `dram` because these keep far more locality than a transpose does: a cross + /// layer gathers its rows from inside the chunk it is already working on. It decides how much + /// the model dislikes doing the work in one large FFT rather than in several small ones, + /// which is the difference between Bluestein's over the whole length and a split, and that + /// call is worth 2x to 4x on a machine with a small cache. + /// + /// Fitted on a Raspberry Pi 5, the machine that has the least cache to spare, over the + /// lengths where the two disagree. At 2 its worst case goes from 4.19 to 1.24 in f64 and + /// from 3.81 to 2.71 in f32, and lengths where the model already won are unaffected. + pub dram_pass: f64, } impl CostModel { @@ -281,8 +308,10 @@ impl CostModel { radixn_extra, general_row: 30.0, small_row: 10.0, + radix_call: 100.0, cache_elems: 256.0 * 1024.0 / 16.0 * complex_per_vector as f64, - dram: 2.0, + dram: 4.0, + dram_pass: 2.0, } } @@ -301,8 +330,11 @@ impl CostModel { Pattern::Strided => accesses / self.complex_per_vector as f64 * self.strided, Pattern::Sequential => accesses / self.complex_per_vector as f64, }; - let _ = ws; - cost + if ws > self.cache_elems { + cost * self.dram_pass + } else { + cost + } } /// Cost of a pass that transposes the whole buffer, which is the one access pattern that @@ -315,7 +347,8 @@ impl CostModel { fn transpose_mem(&self, accesses: f64, pattern: Pattern, ws: f64) -> f64 { let cost = self.mem(accesses, pattern, ws); if ws > self.cache_elems { - cost * self.dram + // `mem` already applied `dram_pass`, so scale up to `dram` in total. + cost * self.dram / self.dram_pass } else { cost } @@ -348,7 +381,8 @@ impl CostModel { // One digit-reversal transpose, then the base FFTs, then k cross layers of // len/4 column_butterfly4, each four butterfly2 plus a rotate90 (10 instructions) // and three twiddle multiplies. - let mut c = self.mem(2.0 * len, Pattern::Permuted, len); + // One digit-reversal transpose, then the base FFTs, then k cross layers. + let mut c = self.mem(2.0 * len, Pattern::Permuted, len) + self.radix_call; c += reps * child_cost(*base_len); c += *k as f64 * ((len / (4.0 * cpv)) * (10.0 + 3.0 * self.mul_complex()) @@ -359,7 +393,7 @@ impl CostModel { let radix_product: usize = factors.iter().map(|f| f.radix()).product(); let len = (*base_len * radix_product) as f64; let reps = len / *base_len as f64; - let mut c = self.mem(2.0 * len, Pattern::Permuted, len); + let mut c = self.mem(2.0 * len, Pattern::Permuted, len) + self.radix_call; c += reps * child_cost(*base_len); for factor in factors.iter() { let radix = factor.radix(); diff --git a/tools/planner_tuning/src/main.rs b/tools/planner_tuning/src/main.rs index 85fe649a..55baca3f 100644 --- a/tools/planner_tuning/src/main.rs +++ b/tools/planner_tuning/src/main.rs @@ -140,8 +140,10 @@ struct Weights { general_row: Option, small_row: Option, /// Cache size in KiB, converted to complex numbers for whichever element type is in use. + radix_call: Option, cache_kib: Option, dram: Option, + dram_pass: Option, } impl Weights { @@ -164,7 +166,9 @@ impl Weights { self.cache_kib .map(|kib| kib * 1024.0 / 16.0 * model.complex_per_vector as f64), ); + set(&mut model.radix_call, self.radix_call); set(&mut model.dram, self.dram); + set(&mut model.dram_pass, self.dram_pass); model } } @@ -178,9 +182,11 @@ struct Options { out: Option, weights: Weights, f32: bool, - /// For `survey`: how many lengths, and the seed that picks them. + /// For `survey`, and for `picks` with `--mix`: how many lengths, and the seed. count: usize, seed: u64, + /// Turn a range into the same spread of lengths `survey` uses, rather than every length. + mix: bool, } /// A tuner for `P`, with the weight overrides applied to its estimating planner. @@ -594,22 +600,129 @@ fn cmd_sweep>(lengths: &[usize], opts: &Options) ); } -/// `count` distinct lengths drawn uniformly from `lo..=hi`, sorted, from a fixed seed. -fn random_lengths(lo: usize, hi: usize, count: usize, seed: u64) -> Vec { +/// A spread of `count` distinct lengths from `lo..=hi`, sorted, reproducible from `seed`. +/// +/// Drawn in equal parts from seven strata, because what a planner does depends on how a length +/// factors, and uniform random integers are almost all "a big prime factor times something +/// small". Log-uniform within each stratum, so every decade up to a million is represented +/// rather than only the largest. +/// +/// The strata are: anything at all; primes, which force Rader's or Bluestein's; 5-smooth and +/// 7-smooth lengths, which every algorithm can decompose; powers of two and three times a power +/// of two, which are Radix4's; a large prime factor times a smooth one, which is the case that +/// forces an awkward split; a product of two middling primes; and short lengths under 2000, +/// which is where the weights were fitted and where plan time matters most. +fn mixed_lengths(lo: usize, hi: usize, count: usize, seed: u64) -> Vec { // splitmix64, so the tool needs no dependencies and a seed reproduces a survey anywhere. let mut state = seed; - let mut next = || { + let mut next = move || { state = state.wrapping_add(0x9e3779b97f4a7c15); let mut z = state; z = (z ^ (z >> 30)).wrapping_mul(0xbf58476d1ce4e5b9); z = (z ^ (z >> 27)).wrapping_mul(0x94d049bb133111eb); z ^ (z >> 31) }; - let span = (hi - lo + 1) as u64; - let count = count.min(span as usize); + let is_prime = |n: usize| { + if n < 2 { + return false; + } + if n % 2 == 0 { + return n == 2; + } + let mut d = 3; + while d * d <= n { + if n % d == 0 { + return false; + } + d += 2; + } + true + }; + // Log-uniform in [lo, hi]. + let log_uniform = |draw: u64, lo: usize, hi: usize| -> usize { + let (lo_l, hi_l) = ((lo.max(1) as f64).ln(), (hi as f64).ln()); + let t = (draw >> 11) as f64 / (1u64 << 53) as f64; + ((lo_l + t * (hi_l - lo_l)).exp() as usize).clamp(lo.max(1), hi) + }; + let mut lengths = std::collections::BTreeSet::new(); - while lengths.len() < count { - lengths.insert(lo + (next() % span) as usize); + let strata = 7; + let per_stratum = count.div_ceil(strata); + for stratum in 0..strata { + let mut taken = 0; + // Bounded, so a stratum with few members in range cannot spin forever. + for _ in 0..per_stratum * 200 { + if taken >= per_stratum || lengths.len() >= count { + break; + } + let draw = next(); + let candidate = match stratum { + // anything + 0 => log_uniform(draw, lo, hi), + // a prime: the next one at or above a log-uniform draw + 1 => { + let mut n = log_uniform(draw, lo.max(2), hi); + while n <= hi && !is_prime(n) { + n += 1; + } + n + } + // 5-smooth or 7-smooth + 2 => { + let mut n = 1usize; + let target = log_uniform(draw, lo.max(8), hi); + let factors = if draw & 1 == 0 { + [2, 3, 5, 5] + } else { + [2, 3, 5, 7] + }; + while n < target { + let f = factors[(next() % 4) as usize]; + if n.saturating_mul(f) > hi { + break; + } + n *= f; + } + n + } + // a power of two, or three times one + 3 => { + let n = log_uniform(draw, lo.max(2), hi); + let p2 = 1usize << (usize::BITS - 1 - n.leading_zeros()).min(20); + if draw & 1 == 0 || p2 * 3 > hi { + p2 + } else { + p2 * 3 + } + } + // a large prime factor times a smooth one, which forces an awkward split + 4 => { + let mut p = log_uniform(draw, 40, (hi / 4).max(41)); + while !is_prime(p) { + p += 1; + } + let smooth = [2usize, 3, 4, 5, 6, 8, 9, 12, 16][(next() % 9) as usize]; + p.saturating_mul(smooth) + } + // two middling primes multiplied together + 5 => { + let mut a = log_uniform(draw, 11, 4000); + while !is_prime(a) { + a += 1; + } + let mut b = log_uniform(next(), 11, (hi / a).max(12)); + while !is_prime(b) { + b += 1; + } + a.saturating_mul(b) + } + // short lengths, where the weights were fitted + _ => log_uniform(draw, lo, hi.min(2000)), + }; + if candidate >= lo.max(1) && candidate <= hi && lengths.insert(candidate) { + taken += 1; + } + } } lengths.into_iter().collect() } @@ -1129,13 +1242,14 @@ fn usage() -> ! { eprintln!(" --rounds N timing rounds per subject (default 9)"); eprintln!(" --block-ms MS wall-clock time per timed block (default 10)"); eprintln!(" --cap N max candidates per length when enumerating (default 48)"); - eprintln!(" --count N survey: how many lengths (default 300)"); + eprintln!(" --count N survey, and picks --mix: how many lengths (default 300)"); + eprintln!(" --mix picks: sample a range like survey does, not every length"); eprintln!(" --seed N survey: seed for picking lengths (default 1)"); eprintln!(" --out FILE dump: where to write (default dump.tsv)"); eprintln!(" --verbose regret: list the top candidates per length"); eprintln!(" cost model weights, overriding the planner's defaults:"); eprintln!(" --strided X --permuted X --rader-index X --radixn-extra X"); - eprintln!(" --general-row X --small-row X --cache-kib X --dram X"); + eprintln!(" --general-row X --small-row X --cache-kib X --dram X --dram-pass X"); std::process::exit(2); } @@ -1157,6 +1271,7 @@ fn main() { f32: false, count: 300, seed: 1, + mix: false, }; let mut rest: Vec = Vec::new(); @@ -1185,6 +1300,7 @@ fn main() { "--cap" => opts.cap = number(value(&mut i, &arg), &arg) as usize, "--count" => opts.count = number(value(&mut i, &arg), &arg) as usize, "--seed" => opts.seed = number(value(&mut i, &arg), &arg) as u64, + "--mix" => opts.mix = true, "--out" => opts.out = Some(value(&mut i, &arg)), "--f32" => opts.f32 = true, "--verbose" => opts.verbose = true, @@ -1194,8 +1310,10 @@ fn main() { "--radixn-extra" => opts.weights.radixn_extra = Some(number(value(&mut i, &arg), &arg)), "--general-row" => opts.weights.general_row = Some(number(value(&mut i, &arg), &arg)), "--small-row" => opts.weights.small_row = Some(number(value(&mut i, &arg), &arg)), + "--radix-call" => opts.weights.radix_call = Some(number(value(&mut i, &arg), &arg)), "--cache-kib" => opts.weights.cache_kib = Some(number(value(&mut i, &arg), &arg)), "--dram" => opts.weights.dram = Some(number(value(&mut i, &arg), &arg)), + "--dram-pass" => opts.weights.dram_pass = Some(number(value(&mut i, &arg), &arg)), other if other.starts_with("--") => { eprintln!("unknown option '{}'", other); usage(); @@ -1224,6 +1342,16 @@ fn main() { |values: &[String]| -> String { values.first().cloned().unwrap_or_else(|| usage()) }; let command = match command_name.as_str() { + "picks" if opts.mix => { + let range = single(&rest); + let (lo, hi) = range.split_once("..").unwrap_or_else(|| usage()); + Command::Picks(mixed_lengths( + lo.parse().expect("bad range start"), + hi.parse().expect("bad range end"), + opts.count, + opts.seed, + )) + } "picks" => Command::Picks(lengths(&rest)), "time" => Command::Time(rest.clone()), "regret" => Command::Regret(lengths(&rest)), @@ -1239,7 +1367,7 @@ fn main() { "# survey\t{} lengths from {}..{}, seed {}", opts.count, lo, hi, opts.seed ); - Command::Sweep(random_lengths(lo, hi, opts.count, opts.seed)) + Command::Sweep(mixed_lengths(lo, hi, opts.count, opts.seed)) } "verify" => Command::Verify(lengths(&rest)), "dump" => Command::Dump(lengths(&rest)), From f9fb9a3da0f4d67618a88e52e1f573f48d8720a9 Mon Sep 17 00:00:00 2001 From: Henrik Date: Fri, 18 Sep 2026 19:12:20 +0200 Subject: [PATCH 19/22] Raise the out-of-cache transpose weight to 6, and document the model Halves the M1's f32 losses beyond 5% over the validation set, 27 to 15, for an unchanged loss count on the Pi 5. Rewrites COST-MODEL.md for the model as it now lives in the library, and records what it scores on both NEON machines. --- src/simd/simd_estimate.rs | 8 +- tools/planner_tuning/COST-MODEL.md | 466 ++++++++++------------------- 2 files changed, 172 insertions(+), 302 deletions(-) diff --git a/src/simd/simd_estimate.rs b/src/simd/simd_estimate.rs index 6adb62a5..a847d1fb 100644 --- a/src/simd/simd_estimate.rs +++ b/src/simd/simd_estimate.rs @@ -273,7 +273,11 @@ pub struct CostModel { /// transposes just as hard makes a MixedRadix wrapped around a smaller radix recipe look /// good, because its inner FFTs are then cache resident, and those recipes measure worse on /// both machines: at `dram_pass` 2 with this at 2, the M1 has 7 losses beyond 5% in f64 and - /// 9 in f32, which raising this to 4 removes entirely. 4 and 6 score the same. + /// 9 in f32, which raising it removes entirely. + /// + /// 6 rather than 4 on the full validation set: it halves the M1's f32 losses beyond 5%, 27 to + /// 15, and leaves the Pi 5's loss count unchanged at 14. The Pi's worst case grows from 3.22 + /// to 3.51, inside the large-Bluestein class that neither value solves. pub dram: f64, /// The same, for every other pass: a RadixN or Radix4 cross layer, a Rader's permutation, or /// Bluestein's padded multiply. @@ -310,7 +314,7 @@ impl CostModel { small_row: 10.0, radix_call: 100.0, cache_elems: 256.0 * 1024.0 / 16.0 * complex_per_vector as f64, - dram: 4.0, + dram: 6.0, dram_pass: 2.0, } } diff --git a/tools/planner_tuning/COST-MODEL.md b/tools/planner_tuning/COST-MODEL.md index 070d3a94..bd539306 100644 --- a/tools/planner_tuning/COST-MODEL.md +++ b/tools/planner_tuning/COST-MODEL.md @@ -1,375 +1,241 @@ # How the estimating planner estimates -This is the "how it works" document. `OP-COUNTS.md` is where the instruction counts come from, -`RESULTS.md` is the evidence that the thing works, `NEXT-STEPS.md` is the live plan, and `README.md` -is how to run the tools. This file explains the mechanism that sits under all four: what a cost is, -how one is computed, which numbers are read off the source, which are fitted against measurements, -and what has to be redone when the kernels change. - -**Status.** There is no estimating planner in the library. `src/tuning/` is a feature-gated -description of what the planners can build, and the cost model lives in the measurement tool at -`tools/planner_tuning/src/counted.rs`. What exists today is a full working prototype driven from the -command line, and `NEXT-STEPS.md` holds the open question of whether it ships whole or only as the -one scoped decision it is most clearly right about. +This is the "how it works" document. `OP-COUNTS.md` is where the instruction counts come from and +`README.md` is how to run the tools. This file explains what a cost is, how one is computed, which +numbers are read off the source, which are fitted against measurements, and what has to be redone +when the kernels change. + +The model lives in the library, at `src/simd/simd_estimate.rs`, and is shared by the NEON, SSE and +wasm_simd planners. The scalar and AVX planners do not use it. ## 1. The loop An estimating planner does three things, and only the middle one is new: -1. **Enumerate.** Build every recipe the planner could plausibly use at this length, as a tree of - [`Spec`](../../src/tuning/mod.rs) nodes. +1. **Enumerate.** List the recipes that could compute this length, as `Shape` values: a top-level + algorithm plus the lengths of its inner FFTs. 2. **Price.** Give each candidate a number. 3. **Pick the minimum.** -Step 1 is [`plan_candidates`](../../src/tuning/mod.rs#L585). Step 2 is -[`CountedModel::cost`](src/counted.rs#L409). Step 3 is a `min_by`, visible in -[`cmd_sweep`](src/main.rs#L452). That is the whole planner. Everything else in this document is -about step 2. +All three are in `design_fft_for_len` in each planner, over `candidates` and `CostModel::cost`. +Ties go to the first candidate, which is always the fixed planner's pick. -The enumeration deliberately answers immediately at a length with its own butterfly and at a power -of two, because measurement showed there is nothing to decide at either, and it emits only the -smaller-width-first ordering of each two-way split. The reasoning and the measured cost of each of -those shortcuts is in the doc comment on `plan_candidates`. Element zero of the list is always the -current fixed planner's pick, so pruning can never leave the estimating planner behind the planner -it replaces. +Enumeration answers immediately at a length with its own butterfly and at a power of two, because +measurement showed there is nothing to decide at either, and it emits only the smaller-width-first +ordering of each two-way split. It offers Bluestein's only where some prime factor has no butterfly +of its own. The reasoning for each shortcut is in the doc comments on `has_choice` and `candidates`. -The separate exhaustive set (`candidates`, `candidates_capped`) exists because *scoring* a pick -needs alternatives that no planner would ever propose. Keep it exhaustive; it is how each shortcut -above was justified. +The exhaustive enumeration in `src/tuning/` (`candidates`, `candidates_capped`) is a separate thing, +used only by the measurement tools: *scoring* a pick needs alternatives no planner would propose. +Keep it exhaustive. It is how each shortcut above was justified, and re-justifying them after a +kernel change needs the same breadth. ## 2. What a cost is **One unit is one issued arithmetic instruction.** A cost is not nanoseconds and does not try to be: -nothing converts it to time, and nothing needs to, because only the ranking within one length is -ever used. That is why a single weight set travels across machines of different clock speeds, and it -is also the reason the absolute numbers below look large and mean nothing on their own. - -Every cost is the sum of two halves: +nothing converts it to time, because only the ranking within one length is ever used. That is why +one weight set travels across machines of different clock speeds, and why the absolute numbers below +are large and mean nothing on their own. ``` cost = counted arithmetic instructions (read off the source, see OP-COUNTS.md) - + a memory term (accesses x pattern x cache level) + + a memory term (accesses x pattern x whether it fits in cache) ``` -A pure operation count, which is the `FFTW_ESTIMATE` analogue, scores worse than the shipping -planner: mean regret 1.359 against 1.093. Adding the memory term takes the same op counts to 1.003. -The memory term is not a refinement, it is the entire result. +A pure operation count, the `FFTW_ESTIMATE` analogue, scores *worse* than the fixed planner. Adding +the memory term is what makes the model work; it is not a refinement. ### The memory term -[`mem()`](src/counted.rs#L382) is four lines and deliberately crude: +`mem()` is deliberately crude. Each pass is charged per element it touches, times a multiplier for +how it walks memory: ``` -accesses -> divided by complexes-per-vector, except under a permutation -level -> L1 if the working set fits in l1_elems, else L2 if it fits l2_elems, else DRAM -cost = accesses * seq[level] * (1.0 sequential | strided_mult | permuted_mult) +sequential 1.0 a contiguous run, one cache line feeding many elements +strided strided a fixed stride, as in a transpose or a cross-FFT layer +permuted permuted a computed address per element: digit reversal, CRT, Rader's ``` -Three things about it are load-bearing and worth stating plainly: - -1. **Sequential versus jumpy carries all of it.** Pricing every access the same degrades worst-case - regret from 1.043 to 1.246. The finer distinction between strided and permuted does much less. -2. **The cache level cannot reorder candidates, as the model is built.** Collapsing L1/L2/DRAM to one - flat cost produces byte-identical picks at every length measured, in cache and out. That is - structural rather than a finding about memory: the level is chosen from the whole transform's - working set (next section), which is the same for every candidate at one length, so it is a - common factor. **It is not true of the hardware.** At 100003 and 100049 on the Pi 5, Rader's is - about 3x faster than every Bluestein's candidate, because Bluestein's inner FFT is 262144 points - (4 MB of complex f64), inside the M1's 12 MB L2 and far outside the A76's 512 KB L2 and 2 MB L3. - No setting of the cache sizes or level weights can express that. See blind spot 4. -3. **A permuted pass is charged per complex number, not per vector.** A gather or scatter computes - an address per element and cannot fill a vector. This is invisible at f64, where the factor is 1, - and was worth a factor of 2 at f32; finding it is what took SSE f32 from 51 to 117 of 216 weight - settings clearing the 20% bar. `--permuted-vector` restores the old behaviour. - -### The working set is threaded down unchanged - -[`cost_ws`](src/counted.rs#L416) passes the *whole transform's* length to every nested algorithm, -not the nested algorithm's own length. Every pass of every inner FFT walks the top-level buffer, so -that is what decides where the traffic is served from. Pricing an inner FFT as if it ran standalone -is the specific mistake that sank the 2021 attempt. - -The exception is Bluestein's, whose inner FFT runs on a buffer of its own that is at least twice the -outer length. Threading the outer length down understates that working set, which is exactly the -Pi 5 defect above. Charging a Bluestein's node at its inner length would fix it and still keep a -node's cost a function of its own subtree, which is what length-keyed recipe memoisation needs. Not -tried yet. - -One caveat on reading `explain` output: it calls `cost()` per node, so each row is priced at its -own length as working set. Since the cache level cannot reorder candidates this almost never -changes a number, but the child rows of a very large transform are informational rather than exact -contributions. +Sequential and strided accesses are divided by the number of complex numbers in a vector; permuted +ones are not, because a gather computes an address per element and cannot fill a vector. + +Three things about it are load-bearing: + +1. **Sequential versus jumpy carries most of the result.** Pricing every access the same degrades + worst-case regret badly. +2. **A pass that no longer fits in cache costs more.** Above `cache_elems` complex numbers an access + costs `dram_pass`, and a transpose's accesses cost `dram`. See section 4, which is where the + machines disagree most. +3. **A node is charged at the size of the buffer it walks itself**, never at the size of the + transform it is nested inside. + +That last point is what makes the search affordable. A recipe's cost depends only on its own +subtree, so each planner caches the best cost per length beside its recipe cache and the search +recurses over the divisors of a length rather than over whole trees. **If a future term ever makes a +node's cost depend on its parent, that caching becomes wrong**, and the failure would be a subtly +bad inner recipe rather than anything that trips a test. ## 3. What each node costs -All of this is [`cost_ws`](src/counted.rs#L416), one match arm per `Spec` variant. `len` is the -node's own length and `ws` the whole transform's. +One match arm per `Shape` variant in `CostModel::cost`. `len` is the node's own length. | node | arithmetic | memory | |---|---|---| | `Butterfly(len)` | counted table lookup | `2*len` sequential | -| `Radix4 { k, base }` | `reps` x base, plus per layer `len/4` butterfly4 and 3 twiddles each | one `2*len` permuted digit reversal, plus `2*len` strided per layer | -| `RadixN { radixes, base }` | `reps` x base, plus per layer `len/r` x (butterfly `r` + `r-1` twiddles) | same shape as Radix4, plus `radixn_extra` per element per layer | +| `Radix4 { k, base_len }` | `reps` x base, plus per layer `len/4` butterfly4 and 3 twiddles each | one `2*len` permuted digit reversal, plus `2*len` strided per layer, plus `radix_call` | +| `RadixN { factors, base_len }` | `reps` x base, plus per layer `len/r` x (butterfly `r` + `r-1` twiddles) | same shape as Radix4, plus `radixn_extra` per element per layer | | `MixedRadix` | `h` x left + `w` x right, plus `len` twiddle multiplies | three transposes, plus one sequential pass | | `GoodThomas` | `h` x left + `w` x right, no twiddles at all | two permuted CRT passes plus one transpose | -| `Raders { inner }` | 2 x inner, `len` twiddles, `len * rader_index` | two permuted passes plus one sequential | -| `Bluesteins { len, inner }` | 2 x inner, `inner.len()` pointwise multiplies | sequential over the inner length and twice over the outer | -| `Dft(n)` | `100 * n^2` | none; a quadratic that only has to sort last | +| `Raders { len }` | 2 x inner, `len` twiddles, `len * rader_index` | two permuted passes plus one sequential | +| `Bluesteins { len, inner_len }` | 2 x inner, `inner_len` pointwise multiplies | sequential over the inner length and twice over the outer | -Two of these encode a decision the model would otherwise be unable to make: +Three of these encode a decision the model would otherwise be unable to make: - **`small` versus general.** `MixedRadixSmall` and `GoodThomasAlgorithmSmall` call `array_utils::transpose_small`, a naive strided double loop, so their transposes are priced `Permuted`. The general variants hand the job to the `transpose` crate, which tiles the rectangle - and gets the reuse back, and pay `general_row` per row for the privilege. That makes the general - form cheaper per element and dearer per row, which is a crossover, and the measurements are the - shape of one: `MixedRadix` crosses under its Small variant near length 200. Without the per-row - term the model has only per-element costs, so it would pick the general form at every length. -- **Which ordering of a split.** `small_row * max(width - height, 0)` is the only thing that - separates `mrs(A,B)` from `mrs(B,A)`. `transpose_small`'s outer loop runs `width` times, so both - Small variants change by exactly `width - height` when the pair is reversed, which is why one - weight covers both. Charging the *difference* rather than the raw count keeps a square pair free - and leaves the better ordering at exactly the cost it had before the term existed. + and gets the reuse back, and pay `general_row` per row for it. That makes the general form cheaper + per element and dearer per row, which is a crossover, and the measurements are the shape of one: + `MixedRadix` crosses under its Small variant near length 200. Without the per-row term the model + has only per-element costs and would pick the general form at every length. +- **Which ordering of a split.** `small_row * max(width - height, 0)` is the only thing separating + `mrs(A,B)` from `mrs(B,A)`. `transpose_small`'s outer loop runs `width` times, so both Small + variants change by exactly `width - height` when the pair is reversed, which is why one weight + covers both. Charging the *difference* rather than the raw count keeps a square pair free and + leaves the better ordering at exactly the cost it had before the term existed. +- **A generic driver against a hand-written one.** `radix_call` is what a RadixN or Radix4 execution + costs regardless of length, and `radixn_extra` is what its cross layers cost per element over + `Radix4` doing the same work. Without the first, short lengths take a RadixN that measures 1.32x + slower than a table-driven `GoodThomasAlgorithmSmall`. ### A worked example, which is also a self-check ``` $ ./target/release/planner_tuning explain 'mrs(b8,rad(rn(3.3,b9)))' recipe len times total cost own cost -mrs(b8,rad(rn(3.3,b9))) 656 x1 127628 13776 +mrs(b8,rad(rn(3.3,b9))) 656 x1 72812 13776 b8 8 x82 4428 4428 - rad(rn(3.3,b9)) 82 x8 109424 69536 - rn(3.3,b9) 81 x16 39888 28080 + rad(rn(3.3,b9)) 82 x8 54608 13120 + rn(3.3,b9) 81 x16 41488 29680 b9 9 x144 11808 11808 ``` The `b8` row by hand: 38 counted instructions from the NEON f64 table, plus `2*8 = 16` sequential -accesses at `seq[0] = 1.0` and a sequential multiplier of 1.0, is 54. It runs 82 times, giving 4428. +accesses at 1.0 each, is 54. It runs 82 times, giving 4428. The root's own cost by hand: three transposes at `3 * 2*656 * 2.5` permuted is 9840, no `small_row` charge because the width is the smaller side already, plus `656` twiddle multiplies at 4 instructions each is 2624, plus a `2*656` sequential pass is 1312. Total 13776. -If a change to `counted.rs` is supposed to leave some family of recipes alone, `explain` on one of -them is the fastest way to confirm it did. +If a change to the model is supposed to leave some family of recipes alone, `explain` on one of them +is the fastest way to confirm it did. ## 4. Where every number comes from This is the distinction the whole approach rests on. Most of the model is read off the source and must be maintained when the source changes; a small set of weights is fitted against measurements -and must be refitted when the machine changes. +and must be rechecked when the machines change. ### Read off the source, never fitted -| quantity | where it is read | notes | -|---|---|---| -| butterfly instruction counts | `src/neon/*.rs`, `src/sse/*.rs`, by hand | tables in `butterfly_compute`, derived in `OP-COUNTS.md` | -| generated prime butterflies | the generator's loop structure | closed form `(h-1)(2h+5)`, cannot drift when sizes are added | -| `mul_complex`, `column_butterfly4` | the vector trait impls | 4 on NEON f64, 6 on SSE and on NEON f32 | -| `registers` | the architecture | 32 `v` on aarch64, 16 `xmm` on x86-64 | -| access counts per pass | the algorithm source | loads plus stores, counted per pass | -| access *pattern* per pass | the algorithm source | which of sequential, strided, permuted applies | -| `l1_elems`, `l2_elems` | the machine's cache sizes | inert as the model is built: no pick depends on them. Defaults are the M1's | - -### Derived from the code, then confirmed against measurement - -| quantity | derivation | +| quantity | where it is read | |---|---| -| `rader_index` = 30 (f64), 45 (f32) | `raders_algorithm.rs` recomputes `index * root % len` per element, a loop-carried `mul -> umulh -> mul -> sub` chain of about 10 cycles. That is worth tens of instruction slots, not the 7 instructions counted, but how many depends on the core's instructions per cycle, so the optimum is per machine: about 40 on the M1, 25 to 30 on a Cortex-A76. The defaults are the values acceptable on all three measured machines (M1, Pi 5, ThinkCentre) by the 1..1000 sweep, not the optimum on any one. See `NEXT-STEPS.md`. | -| `small_row` = 10 | One outer iteration of `transpose_small`, which measures 1.48 ns on the i3 and 0.7 to 1.0 ns on the M1, or 9 to 13 instruction-equivalents at either machine's scale. Scores are byte-identical for anything from 2 to 24 on every dataset, because the term only ever separates two orderings that are otherwise exactly equal. It is a tie-break with a derivation, not a weight. | +| butterfly instruction counts | `src/neon/*.rs`, `src/sse/*.rs`, by hand; derived in `OP-COUNTS.md` | +| generated prime butterflies | the generator's loop structure, as a closed form that cannot drift | +| `mul_complex`, `column_butterfly4` | the vector trait impls | +| access counts and patterns per pass | the algorithm source, loads plus stores per pass | -### Fitted by comparing against measured times +### Fitted, and where each was fitted -Six numbers, and only these six. They convert a memory access into arithmetic-instruction -equivalents, which is the one thing that cannot be read off RustFFT's source because it is a -property of the machine. - -| parameter | what it prices | fitted over | +| weight | value | how it was chosen | |---|---|---| -| `seq[0]`, `seq[1]`, `seq[2]` | one access at L1, L2, DRAM | pinned at 1.0; the other two are inert, see below | -| `strided_mult` | a fixed-stride pass against a sequential one | 1.0, 1.5, 2.5, 4.0 | -| `permuted_mult` | a gather or scatter against a sequential pass | 1.5, 2.5, 4.0, 6.0 | -| `radixn_extra` | the generic `SimdRadixN` driver against the hand-written `Radix4` kernel, per element per layer | 0, 1, 2, 3, 5, 8; defaults 0 on NEON, 6 on SSE f64, 1 on SSE f32, chosen by the 1..1000 sweep | -| `general_row` | per-row setup of the blocked transpose and the on-the-fly CRT mapping | the twelve measured general-over-small ratios | - -`spill` and `mul_complex` are diagnostic overrides rather than fitted weights. Both are off by -default. They exist so that a specific hypothesis can be tested (register pressure in the RadixN -cross layer, and whether SSE's shuffle-heavy complex multiply costs more than its instruction count) -without editing the model. - -### The fitted values, per dataset - -``` -dataset seq[1] strided permuted radixn_extra mean worst -NEON f64 any 1.5 1.5 0 1.003 1.043 -NEON f32 3.0 2.5 2.5 0 1.032 1.121 -SSE f64 2.0 2.5 1.5 5 1.052 1.152 -SSE f32 1.5 2.5 2.5 3 1.034 1.121 -``` - -**Each of those four is one (machine, backend, element type) cell.** Every NEON number comes from -the M1 and every SSE number from the ThinkCentre, so on those two alone "the NEON weights" and "the -M1 weights" are the same column. Which of the two it really is decides whether this ships: - -- **If the split is per backend**, it costs nothing. The backend is chosen at compile time plus a - feature check, and the element type is a generic parameter, so each combination can carry its own - constants exactly as it already carries its own op counts. Four constant sets, all clearing the - bar. -- **If the split is per machine**, no compiled-in constant is right for anybody, and the whole - approach needs a runtime calibration step that nothing here has designed. - -**A second ARM machine, the Pi 5 (Cortex-A76), says mostly per backend.** Its 1..1000 sweeps with -the M1's NEON weights give the same picks at every length, so only the timing differs. Grouped by -which algorithm each planner chose, every class of decision lands within 1.5% of the M1 except -Rader's versus Bluestein's. That includes MixedRadix versus GoodThomas, which `RESULTS.md` had as -the one preference following the machine (wasm on the M1 within 0.004 of NEON on the M1, SSE on the -ThinkCentre 0.088 away): on the Pi those calls score 1.009 in f64 and 1.018 in f32, against 1.008 -and 1.030 on the M1. - -The exception is genuinely per machine. `rader_index` is a latency converted into instruction slots, -so it depends on the core: about 40 on the M1, 25 to 30 on the A76. It is handled without runtime -calibration, by choosing the default that is acceptable on all three machines rather than optimal -on one (`NEXT-STEPS.md` has the three-machine table). That is the rule for any future per-machine weight too: sweep it -everywhere and take the value whose worst machine looks best. - -### The floor, if one set had to serve everything - -Worth knowing because it bounds the damage. Gridding a single memory weight set jointly against all -four dumps, allowing only `radixn_extra` to differ per backend since that one is a register-file -property, the best is `strided 1.5, permuted 2.5, radixn_extra 0 on NEON and 5 on SSE`, with the -cache levels irrelevant as always: - -``` - one shared set own weights fixed planner -NEON f64 1.003 / 1.043 1.003 / 1.043 1.093 / 1.495 -NEON f32 1.032 / 1.225 1.032 / 1.121 1.171 / 1.969 -SSE f64 1.085 / 1.323 1.052 / 1.152 1.246 / 1.734 -SSE f32 1.052 / 1.250 1.034 / 1.121 1.207 / 1.927 -``` - -(mean / worst.) Three of the four then miss the 20% target, so this is not the proposal. But it -still beats the fixed planner on both statistics on all four datasets, and NEON f64 loses nothing at -all. The weights are worth getting right; getting them wrong degrades the result rather than -inverting it. - -`radixn_extra` is the one fitted weight that is not just a fudge: it is exactly 0 on NEON and -positive on SSE, which is what a register-count argument predicts, since `cross_layer` holds 2R rows -live and 2R fits 32 `v` registers at every supported radix and does not fit 16 `xmm`. It also shrinks -from 5 to 2-3 when the element type halves, because a spilled register then covers twice the -elements. Predicted 2.5, observed 2 to 3. - -**That was before the RadixN transpose fix** (415a29f), which removed per-call divides the weight had -been partly absorbing. Refitted by the 1..1000 sweep on the ThinkCentre, the defaults are now 6 for -SSE f64 and 1 for SSE f32, so the halving prediction no longer holds. Zero on NEON still does. The -table above predates the fix; `NEXT-STEPS.md` has the refit. - -Note that `Params::default()` is the NEON working set with `permuted_mult` at 2.5 rather than the -fitted 1.5. It makes no difference to that dataset, but a run that means to reproduce a table above -should pass the weights explicitly rather than trust the defaults. - -## 5. How the weights are fitted - -The procedure is a grid search against a frozen measurement, with a held-out half. - -```sh -# 1. measure once. every candidate at every length, times written to a TSV. -./target/release/planner_tuning dump --planner neon --rounds 7 --cap 48 \ - --out dump_neon_f64.tsv - -# 2. split into halves. even-indexed lengths train, odd-indexed test, on sorted length, -# so each half spans the whole size range. -python3 split.py dump_neon_f64.tsv train.tsv test.tsv - -# 3. grid the weights against the training half only. -./grid.sh train.tsv # 216 points -./sweep.sh train.tsv # the older 108-point grid - -# 4. score the winner on the test half, once. -./target/release/planner_tuning score --seq-l2 1.5 --strided 1.5 --permuted 1.5 \ - --rader-index 30 test.tsv -``` - -Steps 2 to 4 are **pure replay**. No machine is involved, no planner is built, and a full grid takes -seconds. That is the property that makes this maintainable where the old measured-table model was -not: one measurement run per machine, then unlimited model iteration anywhere. - -**Fit on the training half only.** Gridding on the full dump and then quoting a held-out number is -not a held-out number. The honest SSE f32 held-out worst case is 1.163, not the 1.105 that a -fit-on-everything run reports. - -The metric is **regret**: measured time of the chosen recipe divided by measured time of the best -enumerated candidate, so 1.000 is optimal. It is a lower bound on the distance from optimal, because -the candidate set is finite and the inner recipes inside each candidate come from the same planner. - -Two sanity properties of the fit are worth knowing before touching a weight: - -- **96 of 108 settings clear the 20% bar** on NEON f64. The result does not depend on hitting the - weights precisely, which is the main reason to think they travel. -- **The DRAM weight is inert.** 3.0, 6.0, 10.0 and 16.0 give byte-identical results, for the - structural reason in section 2. -- **A flat weight on the tuning set proves nothing about the sweep.** The Rader's weight is flat - from 15 to 120 on the original 33 lengths, because none of them is a small prime. The 1..1000 - sweep moves geometric mean by up to 4% between 30 and 45. - -## 6. Updating the model when the code changes - -This is the cost of the approach. The model is accurate *because* it tracks the source, and that -means the source moving invalidates it. The upside is that every such update is a re-count, which is -mechanical and needs no machine, rather than a re-measurement campaign. - -| what changed | what to redo | -|---|---| -| a butterfly's body | re-count it into `OP-COUNTS.md`, update the table in `butterfly_compute` | -| a butterfly length added or removed | nothing for the generated primes, they are a closed form; a table entry for a hand-written one | -| a vector primitive, for instance FMA or `vcmlaq` arriving | update `Backend::mul_complex` or `column_butterfly4`, then re-count every butterfly built from it. This is what the fcma work would trigger | -| an algorithm's pass structure, for instance `raders_precompute` landing | re-derive that node's arm in `cost_ws`. A precomputed Rader's permutation drops the per-element cost by roughly 4x and `rader_index` stops being a latency term at all | -| a transpose or index computation swapped | recheck the `Pattern` on that pass, and whether `general_row` still describes the same per-row work | -| a new algorithm in a planner | a `Spec` variant, an adapter arm, and a `cost_ws` arm | -| a new backend | a `Backend` variant, its counts, and its register file size | -| a new machine | nothing counted changes. Run `sweep 1..1000` in f64 and f32 and compare it class by class against a machine on the same backend. Where a weight's optimum moves, choose the value acceptable on every machine, never refit to the new one alone | - -After any of them, in this order: +| `strided` | 1.5, or 2.5 on SSE f32 | grid against measured dumps | +| `permuted` | 2.5 | grid against measured dumps | +| `general_row` | 30 | the twelve measured general-over-small ratios | +| `small_row` | 10 | one outer iteration of `transpose_small`, 1.48 ns on an i3, 0.7 to 1.0 on an M1. A tie-break with a derivation; anything from 2 to 24 scores the same | +| `rader_index` | 2 | one load from the permutation table ejmahler#178 added. It was 30 and 45 when that index came from a loop-carried modular multiply | +| `radixn_extra` | 0 on NEON, 6 and 1 on SSE f64 and f32 | expected near zero where 2R rows fit the register file: 32 vector registers on aarch64, 16 on x86-64. **Fitted before the memory terms below and not yet rechecked** | +| `radix_call` | 100 | sized at length 14 on an M1, about 8 ns or 25 cycles: the call, the scratch split, the layer setup and the virtual call into the base FFT | +| `cache_elems` | 256 KiB worth | the smallest last-level cache worth planning for, not any one machine's. No pick below length 16385 changes at this threshold, so weights fitted on short lengths stay valid | +| `dram`, `dram_pass` | 6 and 2 | fitted on a Raspberry Pi 5 and checked on an M1; see below | + +### The two weights where machines disagree + +`dram_pass` and `dram` are the first weights in this model whose optimum depends on the machine +rather than the backend, because they price the memory system rather than the instruction set. The +Pi 5 has 512 KB of L2 and 2 MB of L3 against the M1's 12 MB of L2, and without them the model +computes a whole transform as one Bluestein's whose inner FFT is far larger than cache: on the Pi +that costs 4.19x in f64 and 3.81x in f32, while on the M1 the same picks are fine. + +They have to move together. Charging ordinary passes without charging transposes just as hard makes +a MixedRadix wrapped around a smaller radix recipe look good, since its inner FFTs are then cache +resident, and those recipes measure worse on both machines. Over the lengths where the settings +disagree, as geomean / worst / losses beyond 5%: + +| `dram_pass` / `dram` | M1 f64 | Pi f64 | M1 f32 | Pi f32 | +|---|---|---|---|---| +| 1 (off) | 0.847 / 1.27 / 1 | 0.705 / 4.19 / 7 | 0.982 / 1.49 / 3 | 1.365 / 3.81 / 10 | +| 2 / 2 | 0.981 / 1.62 / 7 | 0.710 / 2.23 / 5 | 1.251 / 1.67 / 9 | 0.916 / 2.84 / 6 | +| 2 / 4 | 0.848 / 1.04 / 0 | 0.624 / 1.29 / 2 | 0.952 / 1.05 / 1 | 0.851 / 3.63 / 2 | +| **2 / 6** | **0.854 / 1.02 / 0** | **0.629 / 1.25 / 3** | **0.949 / 1.07 / 1** | **0.848 / 3.58 / 2** | + +Over the whole validation set, 6 beats 4 by halving the M1's f32 losses beyond 5%, 27 to 15, for +an unchanged loss count on the Pi and a worst case there that grows from 3.22 to 3.51 inside the +large-Bluestein class neither value solves. + +## 4a. What it scores + +334 lengths spread over 1 to 1,000,000 (`survey --count 350 --seed 1`), as the estimating +planner's runtime over the fixed planner's, so below 1 is faster: + +| | geomean | p10 | p50 | p90 | worst | losses >5% | wins >5% | +|---|---|---|---|---|---|---|---| +| M1 f64 | 0.917 | 0.76 | 0.98 | 1.000 | 1.14 | 5 | 147 | +| M1 f32 | 0.933 | 0.79 | 1.00 | 1.005 | 1.40 | 15 | 118 | +| Pi 5 f64 | 0.881 | 0.70 | 0.96 | 1.000 | 1.26 | 3 | 161 | +| Pi 5 f32 | 0.938 | 0.78 | 1.00 | 1.000 | 3.51 | 14 | 123 | + +SSE is not in this table: the machine was unreachable when it was taken, and `radixn_extra` there +was fitted before the memory terms existed. + +The rule for any future weight like these: sweep it on every machine and take the value whose worst +machine looks best, rather than the optimum on any one. + +## 5. What to redo when the kernels change + +In this order: 1. `verify` at a spread of lengths, first. It checks every enumerated candidate against an f64 reference DFT, which catches an illegal spec such as a Bluestein's inner shorter than `2n - 1`. Never trust a timing taken before `verify` is clean. 2. `explain` on a recipe whose cost you can predict, as in the worked example above. -3. `score` against an existing dump, which is free, to see whether the ranking moved at all. -4. `sweep --planner

1..1000` if it did. Both cost-model defects found so far were invisible to - the 33- and 44-length tuning sets and showed up only in the full sweep. -5. A fresh `dump` and refit only if the counts changed enough to move the fitted weights, which is - unusual: the weights price memory, and a kernel change usually moves arithmetic. - -## 7. Known blind spots - -Short version; `RESULTS.md` has the numbers and `NEXT-STEPS.md` has what to do about each. - -1. **Width and height are tied.** The cost function gives `mr(A,B)` and `mr(B,A)` the same cost at - all 415 reversed pairs, apart from the `small_row` tie-break, yet 133 to 179 of them measure more - than 2% apart. This is the clearest unexploited improvement and it is derivable from the code. -2. **MixedRadix versus GoodThomas is a fixed preference, not a decision.** With identical inners the - cost difference is a constant multiple of `len`, so its sign cannot vary with length. It picks - GoodThomas at 142 of 142 pairs, which is right 121 times on the M1 and 38 times on the i3. Keep - any correction small: a stride-aware rewrite aimed at this regressed both backends and was - reverted. -3. **Three machines, two per backend on ARM only.** The Pi 5 separated machine from backend on - ARM (section 4), but x86 still has only the ThinkCentre. A Zen 5 box is the instrument for that. -4. **Bluestein's working set is not priced.** The cache level is taken from the outer transform - (section 2), so a Bluestein's inner FFT two to four times the outer length is charged as if it - fitted where the outer one does. Invisible on the M1 and in any 1..1000 sweep. On the Pi 5 it - picks Bluestein's at 100003 and 100049 where Rader's is about 3x faster. -5. **Plan time.** Enumerate-and-price is 20x to 1265x the fixed planner's plan time, which is the - wrong denominator: what a caller pays is plan plus build, and building is 113x planning overall. - Against plan-plus-build, medians on the M1 at `--cap 48`: - - ``` - len fixed plan plan+price build extra on plan+build in FFT executions - 1260 0.6 us 62.5 us 8.0 us +720% ~12 - 10080 0.5 us 138 us 63 us +219% ~3 - 100800 0.5 us 284 us 537 us +53% ~0.5 - ``` - - So the cost is about twelve executions of the transform being planned at the worst measured - length, under one above 100k, and zero at butterfly lengths and powers of two where enumeration - short-circuits. Memoising inner recipes across candidates is the obvious optimisation and has - deliberately not been done. +3. `picks` before and after, which needs no machine and takes seconds over tens of thousands of + lengths. It says how far the change reaches before anything has been measured. +4. `score` against an existing dump, also free, to see whether the ranking moved. +5. `sweep 1..1000` and a `survey` up to a million if it did. Both cost-model defects found during + the spike were invisible to the 33- and 44-length tuning sets and showed up only in a sweep. + +**A dump goes stale when the algorithm it measured changes.** Every dump taken before ejmahler#178 +holds the old, slow Rader's, so it cannot judge any Rader's decision: at `rader_index` 2 the three +worst lengths in the old NEON f64 dump are all Rader's picks, which is the stale timing talking, not +the weight. + +## 6. Known blind spots + +1. **Bluestein's at large lengths in f32.** The worst remaining class. On the Pi 5, lengths like + 774209 and 232371 still take a Bluestein's whose inner FFT is several times the cache, at 3.2x + and 2.5x the fixed planner. The memory terms improved these without fixing them. +2. **Width and height are tied.** The cost function gives `mr(A,B)` and `mr(B,A)` the same cost + apart from the `small_row` tie-break, yet many reversed pairs measure more than 2% apart. This is + the clearest unexploited improvement and it is derivable from the code. +3. **MixedRadix versus GoodThomas is a fixed preference, not a decision.** With identical inners the + cost difference is a constant multiple of `len`, so its sign cannot vary with length. Keep any + correction small: a stride-aware rewrite aimed at this regressed both backends and was reverted. +4. **x86 has one machine.** The Pi 5 separated machine from backend on ARM. Nothing has done that + for SSE, and `radixn_extra` is the largest single weight fitted there. +5. **Plan time.** Enumerate-and-price costs far more than the fixed planner's plan, which is the + wrong denominator: a caller pays plan plus build, and building dominates. Against plan-plus-build + the estimating planner costs a few executions of the transform it is planning at small lengths, + under one at 100k, and nothing at butterfly lengths and powers of two where enumeration + short-circuits. These are cold-start figures; a planner reused across lengths shares inner + lengths through its caches. From 40c1513c4e7ad11337ea53be7d21bf4503727925 Mon Sep 17 00:00:00 2001 From: Henrik Date: Sat, 19 Sep 2026 02:17:26 +0200 Subject: [PATCH 20/22] Record what the f32 large-prime losses are, and why a fixed cache size cannot fix them --- tools/planner_tuning/COST-MODEL.md | 36 +++++++++++++++++++++++++++--- 1 file changed, 33 insertions(+), 3 deletions(-) diff --git a/tools/planner_tuning/COST-MODEL.md b/tools/planner_tuning/COST-MODEL.md index bd539306..1f0d6ec8 100644 --- a/tools/planner_tuning/COST-MODEL.md +++ b/tools/planner_tuning/COST-MODEL.md @@ -222,9 +222,39 @@ the weight. ## 6. Known blind spots -1. **Bluestein's at large lengths in f32.** The worst remaining class. On the Pi 5, lengths like - 774209 and 232371 still take a Bluestein's whose inner FFT is several times the cache, at 3.2x - and 2.5x the fixed planner. The memory terms improved these without fixing them. +1. **A large prime factor, in f32, on a machine with a small cache.** The worst remaining class, + and the one place a fixed cache threshold cannot be right for everyone. Pi 5 f32, grouped by + largest prime factor, over the validation set: + + | class | n | geomean | losses >5% | losses >20% | worst | + |---|---|---|---|---|---| + | every factor has a butterfly | 131 | 0.941 | 2 | 0 | 1.10 | + | largest prime 33 to 2000 | 121 | 0.918 | 2 | 1 | 1.30 | + | largest prime 2k to 20k | 44 | 0.988 | 3 | 2 | 3.51 | + | largest prime above 20000 | 38 | 0.936 | 7 | 5 | 2.58 | + + In f64 the same table has three losses beyond 5% and a worst of 1.26. Smooth lengths never lose + more than 10% on either machine. + + **The diagnosis, so it does not have to be re-derived.** At 774209 (331 x 2339) the model takes + one Bluestein's with a 1572864-point inner, 12.6 MB in f32. Timing the pieces on the Pi gives + 0.19 to 0.21 ns per unit of cost for the cache-resident parts and 0.79 to 0.98 for that inner, + so out-of-cache work is still under-priced about four times over, after `dram_pass`. The same + pick is *correct* on the M1, where 12.6 MB still fits a 12 MB L2. + + So the multiplier is not the machine-specific part, the threshold is. Fixing it at 256 KiB makes + one number cover both "just past the Pi's 2 MB L3" and "still inside the M1's L2". Raising + `dram_pass` to 3 fixes the Pi outright (774209 goes 3.40 to 1.02, 232371 2.50 to 1.09) and costs + the M1 dearly at that threshold, 0 losses to 7 in f64. Give each machine its own cache size and + the conflict disappears: at a 12 MiB threshold, `dram_pass` 3 costs the M1 one loss instead of + seven, while still fixing the Pi. + + **The fix, when it is wanted, is to read the last-level cache size at construction** and set + `cache_elems` from it: `/sys/devices/system/cpu/cpu0/cache/` on Linux, and + `hw.perflevel0.l2cachesize` on macOS, whose flat `hw.*cachesize` keys report the E-core sizes. + That is a deliberate reversal of the spike's conclusion that a shipping model needs no cache + sizes, which held only while the cache term was inert. + 2. **Width and height are tied.** The cost function gives `mr(A,B)` and `mr(B,A)` the same cost apart from the `small_row` tie-break, yet many reversed pairs measure more than 2% apart. This is the clearest unexploited improvement and it is derivable from the code. From a91c00e292bd898322bb04037ba0c3150bb3d8e0 Mon Sep 17 00:00:00 2001 From: Henrik Date: Sat, 19 Sep 2026 10:01:32 +0200 Subject: [PATCH 21/22] Make rader_index per backend: 2 on NEON, 20 on SSE NEON gathers a complex number with one lane load where SSE assembles it from scalar halves, so the same Rader's permutation costs far more per element there. Over the SSE validation set this takes f64 losses beyond 5% from 35 to 20 and f32 from 13 to 7, with wins up from 94 to 110. --- src/simd/simd_estimate.rs | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/src/simd/simd_estimate.rs b/src/simd/simd_estimate.rs index a847d1fb..dabce5ce 100644 --- a/src/simd/simd_estimate.rs +++ b/src/simd/simd_estimate.rs @@ -226,12 +226,17 @@ pub struct CostModel { /// A gather or scatter against a sequential pass. pub permuted: f64, /// One element of Rader's permutation passes, on top of the gather or scatter itself: the - /// load of its index from the precomputed `u32` table. + /// load of its index from the precomputed `u32` table, and assembling the element. /// - /// Before ejmahler#178 the index came from a loop-carried modular-multiply chain, latency - /// bound, and this was 30 (f64) and 45 (f32). With the table, a survey of 300 random lengths - /// up to a million on an M1 scores 2 and 8 about the same, and both far better than the old - /// values, which made Rader's look expensive enough to trade for a large Bluestein's. + /// Per backend, and the gap is large: 2 on NEON, 20 on SSE. NEON gathers a complex number + /// with a single lane load, while SSE has to assemble one from scalar halves, so the same + /// permutation costs far more per element there. Before ejmahler#178 the index came from a + /// loop-carried modular multiply and this was 30 and 45, dominated by that latency. + /// + /// On NEON a survey of 300 random lengths up to a million on an M1 scores 2 and 8 the same, + /// and both far better than the old values. On SSE, over the lengths where the value changes + /// a pick, 20 gives 2 losses beyond 5% in f64 against 12 at 2, and it is also the best value + /// in f32, so one number serves both precisions there. pub rader_index: f64, /// Extra cost per element per cross-FFT layer of the generic `SimdRadixN` driver, over the /// hand-written `Radix4` kernel doing the same work. Expected near zero on NEON's 32 vector @@ -308,7 +313,10 @@ impl CostModel { complex_per_vector, strided, permuted: 2.5, - rader_index: 2.0, + rader_index: match instruction_set { + InstructionSet::Neon => 2.0, + InstructionSet::Sse => 20.0, + }, radixn_extra, general_row: 30.0, small_row: 10.0, From 8d1ee426a16ca4556509111faae395142859f9c5 Mon Sep 17 00:00:00 2001 From: Henrik Date: Sat, 19 Sep 2026 10:12:41 +0200 Subject: [PATCH 22/22] Record the SSE campaign: rader_index is per backend, the rest holds Measured on the ThinkCentre, the first SSE numbers since the memory terms landed. rader_index was refitted to 2 on NEON evidence alone and is wrong there by a factor of ten. strided, radixn_extra, dram and dram_pass all survive the recheck, the last two on a machine whose cache sits between the other two. --- src/simd/simd_estimate.rs | 13 +++++++++++-- tools/planner_tuning/COST-MODEL.md | 30 ++++++++++++++++-------------- 2 files changed, 27 insertions(+), 16 deletions(-) diff --git a/src/simd/simd_estimate.rs b/src/simd/simd_estimate.rs index dabce5ce..a164eea3 100644 --- a/src/simd/simd_estimate.rs +++ b/src/simd/simd_estimate.rs @@ -239,8 +239,17 @@ pub struct CostModel { /// in f32, so one number serves both precisions there. pub rader_index: f64, /// Extra cost per element per cross-FFT layer of the generic `SimdRadixN` driver, over the - /// hand-written `Radix4` kernel doing the same work. Expected near zero on NEON's 32 vector - /// registers and positive on SSE's 16, since a layer holds two rows per radix live at once. + /// hand-written `Radix4` kernel doing the same work. + /// + /// Zero on NEON and positive on SSE, which is what a register-count argument predicts: a + /// cross layer holds two rows per radix live at once, which fits aarch64's 32 vector + /// registers at every supported radix and does not fit x86-64's 16. + /// + /// Rechecked on the ThinkCentre after the memory terms and `rader_index` changed underneath + /// it, and 6 and 1 survive. A grid over the previous run's losing lengths preferred 10 in + /// f64, but that set is selected for what the old value got wrong, and over the whole + /// validation set 10 scores 24 losses beyond 5% against 6's 20. Grid on an independent + /// sample, not on the lengths a previous run lost. pub radixn_extra: f64, /// Fixed cost per row charged to the general MixedRadix and GoodThomas, for the blocked /// transpose and on-the-fly CRT mapping that the `Small` variants do not have. diff --git a/tools/planner_tuning/COST-MODEL.md b/tools/planner_tuning/COST-MODEL.md index 1f0d6ec8..e0355341 100644 --- a/tools/planner_tuning/COST-MODEL.md +++ b/tools/planner_tuning/COST-MODEL.md @@ -153,11 +153,11 @@ and must be rechecked when the machines change. | `permuted` | 2.5 | grid against measured dumps | | `general_row` | 30 | the twelve measured general-over-small ratios | | `small_row` | 10 | one outer iteration of `transpose_small`, 1.48 ns on an i3, 0.7 to 1.0 on an M1. A tie-break with a derivation; anything from 2 to 24 scores the same | -| `rader_index` | 2 | one load from the permutation table ejmahler#178 added. It was 30 and 45 when that index came from a loop-carried modular multiply | -| `radixn_extra` | 0 on NEON, 6 and 1 on SSE f64 and f32 | expected near zero where 2R rows fit the register file: 32 vector registers on aarch64, 16 on x86-64. **Fitted before the memory terms below and not yet rechecked** | +| `rader_index` | 2 on NEON, 20 on SSE | one load from the permutation table ejmahler#178 added, plus assembling the element, which costs far more on SSE. It was 30 and 45 when that index came from a loop-carried modular multiply | +| `radixn_extra` | 0 on NEON, 6 and 1 on SSE f64 and f32 | expected near zero where 2R rows fit the register file: 32 vector registers on aarch64, 16 on x86-64. Rechecked on the ThinkCentre after the memory terms landed | | `radix_call` | 100 | sized at length 14 on an M1, about 8 ns or 25 cycles: the call, the scratch split, the layer setup and the virtual call into the base FFT | | `cache_elems` | 256 KiB worth | the smallest last-level cache worth planning for, not any one machine's. No pick below length 16385 changes at this threshold, so weights fitted on short lengths stay valid | -| `dram`, `dram_pass` | 6 and 2 | fitted on a Raspberry Pi 5 and checked on an M1; see below | +| `dram`, `dram_pass` | 6 and 2 | fitted on a Raspberry Pi 5, checked on an M1 and on an i3-8100T whose 6 MiB L3 sits between the other two; see below | ### The two weights where machines disagree @@ -186,20 +186,20 @@ large-Bluestein class neither value solves. ## 4a. What it scores 334 lengths spread over 1 to 1,000,000 (`survey --count 350 --seed 1`), as the estimating -planner's runtime over the fixed planner's, so below 1 is faster: +planner's runtime over the fixed planner's, so below 1 is faster and a loss is a length where the +fixed planner was more than 5% faster: | | geomean | p10 | p50 | p90 | worst | losses >5% | wins >5% | |---|---|---|---|---|---|---|---| -| M1 f64 | 0.917 | 0.76 | 0.98 | 1.000 | 1.14 | 5 | 147 | -| M1 f32 | 0.933 | 0.79 | 1.00 | 1.005 | 1.40 | 15 | 118 | -| Pi 5 f64 | 0.881 | 0.70 | 0.96 | 1.000 | 1.26 | 3 | 161 | -| Pi 5 f32 | 0.938 | 0.78 | 1.00 | 1.000 | 3.51 | 14 | 123 | +| M1 (NEON) f64 | 0.917 | 0.76 | 0.98 | 1.000 | 1.14 | 5 | 147 | +| M1 (NEON) f32 | 0.933 | 0.77 | 1.00 | 1.000 | 1.40 | 15 | 118 | +| Pi 5 (NEON) f64 | 0.881 | 0.70 | 0.96 | 1.000 | 1.26 | 3 | 161 | +| Pi 5 (NEON) f32 | 0.938 | 0.79 | 1.00 | 1.000 | 3.51 | 14 | 123 | +| i3-8100T (SSE) f64 | 0.983 | 0.91 | 1.00 | 1.017 | 1.41 | 20 | 64 | +| i3-8100T (SSE) f32 | 0.943 | 0.81 | 1.00 | 1.000 | 1.38 | 7 | 110 | -SSE is not in this table: the machine was unreachable when it was taken, and `radixn_extra` there -was fitted before the memory terms existed. - -The rule for any future weight like these: sweep it on every machine and take the value whose worst -machine looks best, rather than the optimum on any one. +SSE f64 is the weakest cell: it wins at 64 lengths of 334 and loses at 20, for about 2% on +average. The NEON cells win at a third to a half of all lengths. ## 5. What to redo when the kernels change @@ -262,7 +262,9 @@ the weight. cost difference is a constant multiple of `len`, so its sign cannot vary with length. Keep any correction small: a stride-aware rewrite aimed at this regressed both backends and was reverted. 4. **x86 has one machine.** The Pi 5 separated machine from backend on ARM. Nothing has done that - for SSE, and `radixn_extra` is the largest single weight fitted there. + for SSE, so an SSE weight and an i3-8100T weight are still the same column. `rader_index` shows + why that matters: it is 2 on NEON and 20 on SSE, a gap explained by how each backend gathers a + complex number, but only a second x86 machine can confirm that reading. 5. **Plan time.** Enumerate-and-price costs far more than the fixed planner's plan, which is the wrong denominator: a caller pays plan plus build, and building dominates. Against plan-plus-build the estimating planner costs a few executions of the transform it is planning at small lengths,