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 dd3dafd1..719b4c2c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -137,6 +137,15 @@ mod math_utils; mod plan; mod twiddles; +// 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; + use num_complex::Complex; use num_traits::Zero; @@ -602,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/math_utils.rs b/src/math_utils.rs index 57d1cd2f..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() @@ -185,6 +187,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 @@ -487,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/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_planner.rs b/src/neon/neon_planner.rs index b5282f1c..07f1dd8f 100644 --- a/src/neon/neon_planner.rs +++ b/src/neon/neon_planner.rs @@ -4,17 +4,25 @@ 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::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 + 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 +58,10 @@ pub enum Recipe { k: u32, base_fft: Arc, }, + RadixN { + factors: Box<[RadixFactor]>, + base_fft: Arc, + }, Butterfly1, Butterfly2, Butterfly3, @@ -74,6 +86,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, @@ -141,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 { @@ -191,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, }); } } @@ -224,22 +253,222 @@ impl FftPlannerNeon { self.plan_fft(len, FftDirection::Inverse) } - // Make a recipe for a length - fn design_fft_for_len(&mut self, len: usize) -> Arc { + // 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)) } else if let Some(recipe) = self.recipe_cache.get(&len) { 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(); + } + + /// 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) { + 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 - 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 @@ -268,6 +497,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 +684,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, @@ -636,6 +882,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, @@ -643,6 +907,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, @@ -674,7 +945,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)); @@ -685,7 +956,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); @@ -697,7 +968,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); @@ -721,8 +992,20 @@ mod unit_tests { #[test] fn test_plan_neon_mixedradix() { - // Products of several different primes should become MixedRadix - let mut planner = FftPlannerNeon::::new().unwrap(); + // Products of several primes that are all too big for a RadixN cross-FFT layer should + // become MixedRadix + 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); + 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 = fixed(FftPlannerNeon::::new().unwrap()); for pow2 in 2..5 { for pow3 in 2..5 { for pow5 in 2..5 { @@ -732,7 +1015,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 +1023,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 = 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); + 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() { + 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!( is_mixedradixsmall(&plan), @@ -757,7 +1057,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!( @@ -777,7 +1077,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/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/neon/neon_radixn.rs b/src/neon/neon_radixn.rs new file mode 100644 index 00000000..62f07f4d --- /dev/null +++ b/src/neon/neon_radixn.rs @@ -0,0 +1,47 @@ +//! The NEON side of `SimdRadixN`. +//! +//! 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 crate::simd::simd_radixn::SimdRadixN; + +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>; + +#[cfg(test)] +mod unit_tests { + use crate::simd::simd_radixn::test_bodies; + use std::arch::aarch64::{float32x4_t, float64x2_t}; + + #[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/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/plan.rs b/src/plan.rs index 38052f3d..738fd7a5 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), @@ -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,11 @@ 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 @@ -506,18 +510,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 +559,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 +570,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 = 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 diff --git a/src/simd/mod.rs b/src/simd/mod.rs new file mode 100644 index 00000000..c7a30358 --- /dev/null +++ b/src/simd/mod.rs @@ -0,0 +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..a164eea3 --- /dev/null +++ b/src/simd/simd_estimate.rs @@ -0,0 +1,689 @@ +//! 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(crate) 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: the + /// load of its index from the precomputed `u32` table, and assembling the element. + /// + /// 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. + /// + /// 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. + 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, + /// 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. + /// + /// 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. + /// + /// 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 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. + /// + /// 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 { + /// 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: match instruction_set { + InstructionSet::Neon => 2.0, + InstructionSet::Sse => 20.0, + }, + 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: 6.0, + dram_pass: 2.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, 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, + }; + 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 + /// 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 { + // `mem` already applied `dram_pass`, so scale up to `dram` in total. + cost * self.dram / self.dram_pass + } else { + cost + } + } + + 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(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) => { + 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; + 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. + // 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()) + + 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, len) + self.radix_call; + 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, len); + 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.transpose_mem(2.0 * len, Pattern::Permuted, len) + + self.small_row * left_len.saturating_sub(*right_len) as f64 + } else { + 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, len); + 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.transpose_mem(2.0 * len, Pattern::Permuted, len); + c += if *small { + self.transpose_mem(2.0 * len, Pattern::Permuted, len) + + self.small_row * left_len.saturating_sub(*right_len) as f64 + } else { + 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); + 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) + 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 } => { + 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, inner); + c += 2.0 + * (outer * self.mul_complex() + + self.mem(2.0 * outer, Pattern::Sequential, outer)); + 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(crate) 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/simd/simd_planner.rs b/src/simd/simd_planner.rs new file mode 100644 index 00000000..445bf68f --- /dev/null +++ b/src/simd/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::{split_cross_len, 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: 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/simd_radixn.rs b/src/simd/simd_radixn.rs new file mode 100644 index 00000000..f4e4da85 --- /dev/null +++ b/src/simd/simd_radixn.rs @@ -0,0 +1,644 @@ +//! 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 `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 +//! 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::{reverse_remainders, workaround_transmute_mut, TransposeFactor}; +use crate::common::{FftNum, RadixFactor}; +use crate::{Direction, Fft, FftDirection, Length}; + +use super::simd_vector::SimdVector; + +/// 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, + + // 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, + 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, + }); + } + } + + // 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. + 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, + + unroll_factor, + reversed_columns, + 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.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 + 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 + 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(); + + // 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 => { + cross_layer_chunks::(out, layer_twiddles, num_columns, |v| { + V::column_butterfly2(v) + }) + } + InternalRadixFactor::Factor3(bf) => { + cross_layer_chunks::(out, layer_twiddles, num_columns, |v| { + V::column_butterfly3(bf, v) + }) + } + InternalRadixFactor::Factor4(rotation) => { + cross_layer_chunks::(out, layer_twiddles, num_columns, |v| { + V::column_butterfly4(v, *rotation) + }) + } + InternalRadixFactor::Factor5(bf) => { + cross_layer_chunks::(out, layer_twiddles, num_columns, |v| { + V::column_butterfly5(bf, v) + }) + } + InternalRadixFactor::Factor6(bf) => { + cross_layer_chunks::(out, layer_twiddles, num_columns, |v| { + V::column_butterfly6(bf, v) + }) + } + InternalRadixFactor::Factor7(bf) => { + cross_layer_chunks::(out, 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..]; + } + } +} + +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(), + |input, output, scratch| { + self.transpose(input, output); + self.base_fft.process_with_scratch(output, scratch); + self.cross_ffts(output); + }, + ); + } + } + 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(), + |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); + }, + ); + } + } + 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 (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); + }, + ) + } + } + #[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 + } +} + +/// 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. +/// +/// 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] { + std::array::from_fn(|r| { + let v = V::load(data, idx + r * num_columns); + 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); + for i in 0..unroll_count { + let vcol = i * 2; + 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); + } + } + + // an odd vector column count leaves one behind + 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() { + 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: SimdVector, + 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: SimdVector, + V64: SimdVector, + { + 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: SimdVector, + V64: SimdVector, + { + 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: SimdVector, + V64: SimdVector, + { + 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: SimdVector, + V64: SimdVector, + { + 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: SimdVector, + 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); + } +} 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/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_planner.rs b/src/sse/sse_planner.rs index 6d2e790d..495dc77b 100644 --- a/src/sse/sse_planner.rs +++ b/src/sse/sse_planner.rs @@ -4,17 +4,25 @@ 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::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 + 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 +58,10 @@ pub enum Recipe { k: u32, base_fft: Arc, }, + RadixN { + factors: Box<[RadixFactor]>, + base_fft: Arc, + }, Butterfly1, Butterfly2, Butterfly3, @@ -74,6 +86,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, @@ -141,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 { @@ -191,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, }); } } @@ -224,22 +253,222 @@ impl FftPlannerSse { self.plan_fft(len, FftDirection::Inverse) } - // Make a recipe for a length - fn design_fft_for_len(&mut self, len: usize) -> Arc { + // 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)) } else if let Some(recipe) = self.recipe_cache.get(&len) { 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(); + } + + /// 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) { + 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 - 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 @@ -268,6 +497,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 +685,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, @@ -637,6 +883,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, @@ -644,6 +908,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, @@ -675,7 +946,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)); @@ -686,7 +957,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); @@ -698,7 +969,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); @@ -722,8 +993,20 @@ mod unit_tests { #[test] fn test_plan_sse_mixedradix() { - // Products of several different primes should become MixedRadix - let mut planner = FftPlannerSse::::new().unwrap(); + // Products of several primes that are all too big for a RadixN cross-FFT layer should + // become MixedRadix + 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); + 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 = fixed(FftPlannerSse::::new().unwrap()); for pow2 in 2..5 { for pow3 in 2..5 { for pow5 in 2..5 { @@ -733,7 +1016,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 +1024,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 = 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); + 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() { + 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!( is_mixedradixsmall(&plan), @@ -758,7 +1058,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!( @@ -778,7 +1078,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/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/sse/sse_radixn.rs b/src/sse/sse_radixn.rs new file mode 100644 index 00000000..c1eba04a --- /dev/null +++ b/src/sse/sse_radixn.rs @@ -0,0 +1,47 @@ +//! The SSE side of `SimdRadixN`. +//! +//! 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 crate::simd::simd_radixn::SimdRadixN; + +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>; + +#[cfg(test)] +mod unit_tests { + use crate::simd::simd_radixn::test_bodies; + use std::arch::x86_64::{__m128, __m128d}; + + #[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>(); + } +} 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/tuning/adapters.rs b/src/tuning/adapters.rs new file mode 100644 index 00000000..26d85119 --- /dev/null +++ b/src/tuning/adapters.rs @@ -0,0 +1,490 @@ +//! 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 { + // The fixed planner, whose picks `plan` reports. + planner: $planner, + // The estimating planner, whose picks `estimate` reports. + estimator: $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 { + let mut planner = + <$planner>::new().expect(concat!("this machine does not support ", $label)); + planner.set_estimating(false); + Self { + planner, + estimator: <$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 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); + 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/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 new file mode 100644 index 00000000..4e1555a9 --- /dev/null +++ b/src/tuning/mod.rs @@ -0,0 +1,583 @@ +//! 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::*; + +#[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 +// --------------------------------------------------------------------------- + +/// 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 + } + + /// 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) {} +} + +// --------------------------------------------------------------------------- +// 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 fixed planner's pick first. +/// +/// 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. +/// +/// 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 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; + 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; + + 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] + } 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. + 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] { + 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. +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/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_planner.rs b/src/wasm_simd/wasm_simd_planner.rs index 4b4fd0a6..6a15badd 100644 --- a/src/wasm_simd/wasm_simd_planner.rs +++ b/src/wasm_simd/wasm_simd_planner.rs @@ -4,12 +4,16 @@ use crate::algorithm::{ BluesteinsAlgorithm, Dft, GoodThomasAlgorithm, GoodThomasAlgorithmSmall, MixedRadix, MixedRadixSmall, RadersAlgorithm, }; +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}; 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 +49,10 @@ pub enum Recipe { k: u32, base_fft: Arc, }, + RadixN { + factors: Box<[RadixFactor]>, + base_fft: Arc, + }, Butterfly1, Butterfly2, Butterfly3, @@ -69,6 +77,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, @@ -136,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. @@ -174,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`. @@ -200,20 +228,221 @@ impl FftPlannerWasmSimd { } impl FftPlannerWasmSimd { - fn design_fft_for_len(&mut self, len: usize) -> Arc { + // 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)) } else if let Some(recipe) = self.recipe_cache.get(&len) { 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 } } - 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(); + } + + /// 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) { + 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 @@ -241,6 +470,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 +657,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, @@ -607,6 +854,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 { @@ -616,6 +881,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,9 +917,9 @@ 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(); + 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)); @@ -656,9 +928,9 @@ 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(); + let mut planner = fixed(FftPlannerWasmSimd::::new().unwrap()); for pow in 6..32 { let len = 1 << pow; let plan = planner.design_fft_for_len(len); @@ -668,9 +940,9 @@ 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(); + 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); @@ -693,9 +965,21 @@ mod unit_tests { } #[wasm_bindgen_test] - fn test_plan_sse_mixedradix() { - // Products of several different primes should become MixedRadix - let mut planner = FftPlannerWasmSimd::::new().unwrap(); + 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 = 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); + 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 = fixed(FftPlannerWasmSimd::::new().unwrap()); for pow2 in 2..5 { for pow3 in 2..5 { for pow5 in 2..5 { @@ -705,7 +989,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 +998,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 = 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); + 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() { + 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!( is_mixedradixsmall(&plan), @@ -729,8 +1030,8 @@ mod unit_tests { } #[wasm_bindgen_test] - fn test_plan_sse_goodthomasbutterfly() { - let mut planner = FftPlannerWasmSimd::::new().unwrap(); + fn test_plan_wasm_simd_goodthomasbutterfly() { + 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!( @@ -743,14 +1044,14 @@ 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, 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!( @@ -768,24 +1069,24 @@ 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(); + 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!( @@ -796,9 +1097,9 @@ 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 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!( 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/src/wasm_simd/wasm_simd_radixn.rs b/src/wasm_simd/wasm_simd_radixn.rs new file mode 100644 index 00000000..24aeb8f0 --- /dev/null +++ b/src/wasm_simd/wasm_simd_radixn.rs @@ -0,0 +1,48 @@ +//! The WASM SIMD side of `SimdRadixN`. +//! +//! 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 crate::simd::simd_radixn::SimdRadixN; + +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>; + +#[cfg(test)] +mod unit_tests { + 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] + 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::(); + } +} 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::*; 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 { 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..e0355341 --- /dev/null +++ b/tools/planner_tuning/COST-MODEL.md @@ -0,0 +1,273 @@ +# How the estimating planner estimates + +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.** 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.** + +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. + +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 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, 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 whether it fits in cache) +``` + +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()` is deliberately crude. Each pass is charged per element it touches, times a multiplier for +how it walks memory: + +``` +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 +``` + +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 + +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_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 { 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 | + +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 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 72812 13776 + b8 8 x82 4428 4428 + 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 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 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 rechecked when the machines change. + +### Read off the source, never fitted + +| quantity | where it is read | +|---|---| +| 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, and where each was fitted + +| weight | value | how it was chosen | +|---|---|---| +| `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 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, 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 + +`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 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 (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 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 + +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. `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. **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. +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, 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, + 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. 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..2f974c0d --- /dev/null +++ b/tools/planner_tuning/README.md @@ -0,0 +1,125 @@ +# planner_tuning + +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. + +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 + +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. + +**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. + +## Comparing the two planners + +`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. + +`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 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 +``` + +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. + +## Fitting weights: 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. + +```sh +./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 +``` + +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 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 runs or machines +.venv/bin/python plot_sweep.py sweep_neon_f64.tsv --out fig.png # write instead of show +``` + +## Other backends + +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. + +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 +RUSTFLAGS="-C target-feature=+simd128" cargo build --release --target wasm32-wasip1 +node run_wasm.mjs sweep --planner wasm_simd --block-ms 100 1..200 +``` + +## Traps + +- **`--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. Inline + the numbers or use `${=VAR}`. +- Measurement output (`*.tsv`, `*.txt`, `*.log`) and `.venv` are gitignored. + +## The documents + +- `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. 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/main.rs b/tools/planner_tuning/src/main.rs new file mode 100644 index 00000000..55baca3f --- /dev/null +++ b/tools/planner_tuning/src/main.rs @@ -0,0 +1,1392 @@ +//! 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. "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... 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, spec_cost, to_spec_string, CostModel, InstructionSet, 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() {} + +// --------------------------------------------------------------------------- +// Options +// --------------------------------------------------------------------------- + +/// 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, + /// 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 { + 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); + // 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.radix_call, self.radix_call); + set(&mut model.dram, self.dram); + set(&mut model.dram_pass, self.dram_pass); + model + } +} + +struct Options { + rounds: usize, + block_ms: f64, + cap: usize, + verbose: bool, + /// Where `dump` writes its rows. + out: Option, + weights: Weights, + f32: bool, + /// 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. +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 +} + +/// 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 }, + )) +} + +/// 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() +} + +fn percentile(sorted: &[f64], fraction: f64) -> f64 { + sorted[((sorted.len() as f64 * fraction) as usize).min(sorted.len() - 1)] +} + +// --------------------------------------------------------------------------- +// Subcommands +// --------------------------------------------------------------------------- + +fn cmd_time>(specs: &[String], opts: &Options) { + let mut planner = tuner::(opts); + 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); + } + } +} + +/// 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) +} + +/// 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} {}", + "len", "planner", "model", "planner ns", "best ns", "best recipe (when neither picked it)" + ); + + let mut planner_regrets = Vec::new(); + let mut model_regrets = Vec::new(); + + for &len in lengths { + let mut planner = tuner::(opts); + let (specs, model_index) = candidates_with_model_pick(&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); + + 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 planner_regret = subjects[0].best() / best_time; + let model_regret = model_index.map(|i| subjects[i].best() / best_time); + + println!( + "{:>9} {:>7.3}x {:>8} {:>10.0} {:>10.0} {}", + len, + planner_regret, + model_regret + .map(|r| format!("{:.3}x", r)) + .unwrap_or_else(|| "-".into()), + subjects[0].best(), + best_time, + if best_index == 0 || Some(best_index) == model_index { + String::new() + } else { + subjects[best_index].name.clone() + } + ); + if opts.verbose { + let mut ranked: Vec<&Subject> = subjects.iter().collect(); + ranked.sort_by(|a, b| a.best().total_cmp(&b.best())); + for subject in ranked.iter().take(6) { + println!( + " {:>6.3}x {}", + subject.best() / best_time, + subject.name + ); + } + println!(" planner: {}", subjects[0].name); + if let Some(i) = model_index { + println!(" model: {}", subjects[i].name); + } + println!(" ({} candidates measured)", subjects.len()); + } + + planner_regrets.push(planner_regret); + if let Some(r) = model_regret { + model_regrets.push(r); + } + } + + 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 fixed planner's pick against the estimating planner's, at every length given. +/// +/// 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 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 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!("# 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 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, + )); + } + measure(&mut subjects, opts.rounds, opts.block_ms); + + let planner_ns = subjects[0].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. + 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, + if agree { 1 } else { 0 }, + planner_ns, + model_ns, + planner_ns / model_ns, + norm(planner_ns), + norm(model_ns), + 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] + ); +} + +/// 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 = 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 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(); + 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() +} + +/// 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; + + 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 = 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(); + + 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); + } +} + +/// 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; + } + 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 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; + } + Some(_) => {} + None => rows.push((spec, ns, pick)), + } + } + data +} + +/// Score the cost model against a dump file. No measurement, no planner, no machine. +/// +/// 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 (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); + + let pick = rows + .iter() + .filter_map(|(text, ns, _)| { + let spec = parse(text).ok()?; + Some((text, spec_cost(&model, &spec)?, *ns)) + }) + .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 model_regret = pick_ns / best; + model_regrets.push(model_regret); + if let Some(p) = planner_regret { + planner_regrets.push(p); + } + println!( + "{:>9} {:>8.3}x {:>8} {}", + len, + model_regret, + planner_regret + .map(|p| format!("{:.3}x", p)) + .unwrap_or_else(|| "-".into()), + if model_regret <= 1.0001 { + "= best" + } else { + pick_name + } + ); + } + + 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, + values.len(), + values.iter().sum::() / values.len() as f64, + percentile(&values, 0.5), + percentile(&values, 0.9), + values[values.len() - 1] + ); + } +} + +/// 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 = 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 } => { + vec![(base.as_ref(), radixes.iter().product::() as f64)] + } + 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}", + "", + to_spec_string(spec), + spec.len(), + mult, + mult * own, + mult * (own - child_total), + indent = depth * 2 + )); + for (child, m) in kids { + walk(model, child, mult * m, depth + 1, out); + } + } + + let mut out = Vec::new(); + 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); + } +} + +/// 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 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; + + 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) + ); +} + +/// 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`, 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) { + println!( + "{:>7} {:>6} {:>6} {:>11} {:>11} {:>12} {}", + "len", "cands", "k", "build ns", "exec ns", "total ns", "recipe" + ); + + for &len in lengths { + 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() + .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(); + + // `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 start = Instant::now(); + for _ in 0..build_reps { + std::hint::black_box(planner.build(spec, FftDirection::Forward)); + } + 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]).total_cmp(&(build[b] + k * exec[b]))) + .unwrap() + }; + + 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 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 + ); + } + + let one = pick(1.0); + println!( + "{:>7} {:>6} {:>6} {:>11.0} {:>11.1} {:>12} {}", + "", "", "model", build[mi], exec[mi], "", subjects[mi].name + ); + 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!(); + } +} + +// --------------------------------------------------------------------------- + +enum Command { + Picks(Vec), + Time(Vec), + Regret(Vec), + Sweep(Vec), + Verify(Vec), + Dump(Vec), + Score(String), + Explain(String), + Costs(String), + Plantime(Vec), + Crossover(Vec), +} + +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), + Command::Verify(lengths) => cmd_verify::(lengths, opts), + 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::Costs(path) => cmd_costs(path, opts), + Command::Explain(spec) => cmd_explain(spec, P::label(), opts), + } +} + +fn dispatch(planner: &str, command: &Command, opts: &Options) { + match planner { + "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), + #[cfg(target_arch = "x86_64")] + "sse" => run::>(command, opts), + #[cfg(target_arch = "wasm32")] + "wasm_simd" => run::>(command, opts), + 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: 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!( + " --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, 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 --dram-pass X"); + std::process::exit(2); +} + +fn main() { + let args: Vec = std::env::args().skip(1).collect(); + if args.is_empty() { + usage(); + } + + let command_name = args[0].clone(); + let mut planner = native_planner().to_string(); + let mut opts = Options { + rounds: 9, + block_ms: 10.0, + cap: 48, + verbose: false, + out: None, + weights: Weights::default(), + f32: false, + count: 300, + seed: 1, + mix: 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() { + 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, + "--mix" => opts.mix = true, + "--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)), + "--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(); + } + other => rest.push(other.to_string()), + } + i += 1; + } + + // 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("..") { + 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 single = + |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)), + "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(mixed_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); + usage(); + } + }; + + request_performance_core(); + + if opts.f32 { + dispatch::(&planner, &command, &opts); + } else { + dispatch::(&planner, &command, &opts); + } +}