Skip to content

Commit

Permalink
Make planner reusable
Browse files Browse the repository at this point in the history
- Planner should be re-usable so it can be re-used for FFT's of the same
  size

- Add regression tests to make sure `fft_64`/`fft_32` gives the same
  results as `fft_64_with_opts_and_plan`/`fft_32_with_opts_and_plan`
  • Loading branch information
smu160 committed Jul 1, 2024
1 parent 2df2f00 commit 9db5739
Show file tree
Hide file tree
Showing 6 changed files with 272 additions and 43 deletions.
8 changes: 7 additions & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,14 @@ num-traits = "0.2.18"
multiversion = "0.7"

[dev-dependencies]
utilities = { path = "utilities" }
criterion = "0.5.1"
fftw = "0.8.0"
rand = "0.8.5"
utilities = { path = "utilities" }

[[bench]]
name = "bench"
harness = false

[profile.release]
codegen-units = 1
Expand Down
119 changes: 119 additions & 0 deletions benches/bench.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,119 @@
use criterion::{black_box, criterion_group, criterion_main, Criterion};
use num_traits::Float;
use phastft::{
fft_32_with_opts_and_plan, fft_64_with_opts_and_plan,
options::Options,
planner::{Direction, Planner32, Planner64},
};
use rand::{
distributions::{Distribution, Standard},
thread_rng, Rng,
};

const LENGTHS: &[usize] = &[
6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28,
];

fn generate_numbers<T: Float>(n: usize) -> (Vec<T>, Vec<T>)
where
Standard: Distribution<T>,
{
let mut rng = thread_rng();

let samples: Vec<T> = (&mut rng).sample_iter(Standard).take(2 * n).collect();

let mut reals = vec![T::zero(); n];
let mut imags = vec![T::zero(); n];

for ((z_re, z_im), rand_chunk) in reals
.iter_mut()
.zip(imags.iter_mut())
.zip(samples.chunks_exact(2))
{
*z_re = rand_chunk[0];
*z_im = rand_chunk[1];
}

(reals, imags)
}

fn benchmark_forward_f32(c: &mut Criterion) {
let options = Options::default();

for n in LENGTHS.iter() {
let len = 1 << n;
let id = format!("FFT Forward f32 {} elements", len);
c.bench_function(&id, |b| {
let (mut reals, mut imags) = generate_numbers(len);
let planner = Planner32::new(len, Direction::Forward);
b.iter(|| {
black_box(fft_32_with_opts_and_plan(
&mut reals, &mut imags, &options, &planner,
));
});
});
}
}

fn benchmark_inverse_f32(c: &mut Criterion) {
let options = Options::default();

for n in LENGTHS.iter() {
let len = 1 << n;
let id = format!("FFT Inverse f32 {} elements", len);
c.bench_function(&id, |b| {
let (mut reals, mut imags) = generate_numbers(len);
let planner = Planner32::new(len, Direction::Reverse);
b.iter(|| {
black_box(fft_32_with_opts_and_plan(
&mut reals, &mut imags, &options, &planner,
));
});
});
}
}

fn benchmark_forward_f64(c: &mut Criterion) {
let options = Options::default();

for n in LENGTHS.iter() {
let len = 1 << n;
let id = format!("FFT Forward f64 {} elements", len);
c.bench_function(&id, |b| {
let (mut reals, mut imags) = generate_numbers(len);
let planner = Planner64::new(len, Direction::Forward);
b.iter(|| {
black_box(fft_64_with_opts_and_plan(
&mut reals, &mut imags, &options, &planner,
));
});
});
}
}

fn benchmark_inverse_f64(c: &mut Criterion) {
let options = Options::default();

for n in LENGTHS.iter() {
let len = 1 << n;
let id = format!("FFT Inverse f64 {} elements", len);
c.bench_function(&id, |b| {
let (mut reals, mut imags) = generate_numbers(len);
let planner = Planner64::new(len, Direction::Reverse);
b.iter(|| {
black_box(fft_64_with_opts_and_plan(
&mut reals, &mut imags, &options, &planner,
));
});
});
}
}

criterion_group!(
benches,
benchmark_forward_f32,
benchmark_inverse_f32,
benchmark_forward_f64,
benchmark_inverse_f64
);
criterion_main!(benches);
119 changes: 90 additions & 29 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -48,35 +48,15 @@ macro_rules! impl_fft_for {
imags.len()
);

let mut planner = <$planner>::new(reals.len(), Direction::Forward);
let planner = <$planner>::new(reals.len(), direction);
assert!(
planner.num_twiddles().is_power_of_two()
&& planner.num_twiddles() == reals.len() / 2
);

let opts = Options::guess_options(reals.len());

match direction {
Direction::Reverse => {
for z_im in imags.iter_mut() {
*z_im = -*z_im;
}
}
_ => (),
}

$opts_and_plan(reals, imags, &opts, &mut planner);

match direction {
Direction::Reverse => {
let scaling_factor = (reals.len() as $precision).recip();
for (z_re, z_im) in reals.iter_mut().zip(imags.iter_mut()) {
*z_re *= scaling_factor;
*z_im *= -scaling_factor;
}
}
_ => (),
}
$opts_and_plan(reals, imags, &opts, &planner);
}
};
}
Expand Down Expand Up @@ -112,30 +92,39 @@ macro_rules! impl_fft_with_opts_and_plan_for {
reals: &mut [$precision],
imags: &mut [$precision],
opts: &Options,
planner: &mut $planner,
planner: &$planner,
) {
assert!(reals.len() == imags.len() && reals.len().is_power_of_two());
let n: usize = reals.len().ilog2() as usize;

let twiddles_re = &mut planner.twiddles_re;
let twiddles_im = &mut planner.twiddles_im;
let mut twiddles_re = planner.twiddles_re.clone();
let mut twiddles_im = planner.twiddles_im.clone();

// We shouldn't be able to execute FFT if the # of twiddles isn't equal to the distance
// between pairs
assert!(twiddles_re.len() == reals.len() / 2 && twiddles_im.len() == imags.len() / 2);

match planner.direction {
Direction::Reverse => {
for z_im in imags.iter_mut() {
*z_im = -*z_im;
}
}
_ => (),
}

for t in (0..n).rev() {
let dist = 1 << t;
let chunk_size = dist << 1;

if chunk_size > 4 {
if t < n - 1 {
filter_twiddles(twiddles_re, twiddles_im);
(twiddles_re, twiddles_im) = filter_twiddles(&twiddles_re, &twiddles_im);
}
if chunk_size >= $lanes * 2 {
$simd_butterfly_kernel(reals, imags, twiddles_re, twiddles_im, dist);
$simd_butterfly_kernel(reals, imags, &twiddles_re, &twiddles_im, dist);
} else {
fft_chunk_n(reals, imags, twiddles_re, twiddles_im, dist);
fft_chunk_n(reals, imags, &twiddles_re, &twiddles_im, dist);
}
} else if chunk_size == 2 {
fft_chunk_2(reals, imags);
Expand All @@ -153,6 +142,17 @@ macro_rules! impl_fft_with_opts_and_plan_for {
cobra_apply(reals, n);
cobra_apply(imags, n);
}

match planner.direction {
Direction::Reverse => {
let scaling_factor = (reals.len() as $precision).recip();
for (z_re, z_im) in reals.iter_mut().zip(imags.iter_mut()) {
*z_re *= scaling_factor;
*z_im *= -scaling_factor;
}
}
_ => (),
}
}
};
}
Expand All @@ -179,7 +179,7 @@ mod tests {

use utilities::rustfft::num_complex::Complex;
use utilities::rustfft::FftPlanner;
use utilities::{assert_float_closeness, gen_random_signal};
use utilities::{assert_float_closeness, gen_random_signal, gen_random_signal_f32};

use super::*;

Expand Down Expand Up @@ -308,4 +308,65 @@ mod tests {
}
}
}

#[test]
fn fft_64_with_opts_and_plan_vs_fft_64() {
let num_points = 4096;

let mut reals = vec![0.0; num_points];
let mut imags = vec![0.0; num_points];
gen_random_signal(&mut reals, &mut imags);

let mut re = reals.clone();
let mut im = imags.clone();

let mut planner = Planner64::new(num_points, Direction::Forward);
let opts = Options::guess_options(reals.len());
fft_64_with_opts_and_plan(&mut reals, &mut imags, &opts, &mut planner);

fft_64(&mut re, &mut im, Direction::Forward);

reals
.iter()
.zip(imags.iter())
.zip(re.iter())
.zip(im.iter())
.for_each(|(((r, i), z_re), z_im)| {
assert_float_closeness(*r, *z_re, 1e-6);
assert_float_closeness(*i, *z_im, 1e-6);
});
}

#[test]
fn fft_32_with_opts_and_plan_vs_fft_64() {
let dirs = [Direction::Forward, Direction::Reverse];

for direction in dirs {
for n in 4..14 {
let num_points = 1 << n;
let mut reals = vec![0.0; num_points];
let mut imags = vec![0.0; num_points];
gen_random_signal_f32(&mut reals, &mut imags);

let mut re = reals.clone();
let mut im = imags.clone();

let mut planner = Planner32::new(num_points, direction);
let opts = Options::guess_options(reals.len());
fft_32_with_opts_and_plan(&mut reals, &mut imags, &opts, &mut planner);

fft_32(&mut re, &mut im, direction);

reals
.iter()
.zip(imags.iter())
.zip(re.iter())
.zip(im.iter())
.for_each(|(((r, i), z_re), z_im)| {
assert_float_closeness(*r, *z_re, 1e-6);
assert_float_closeness(*i, *z_im, 1e-6);
});
}
}
}
}
11 changes: 10 additions & 1 deletion src/planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@ macro_rules! impl_planner_for {
pub twiddles_re: Vec<$precision>,
/// The imaginary components of the twiddle factors
pub twiddles_im: Vec<$precision>,
/// The direction of the FFT associated with this `Planner`
pub direction: Direction,
}
impl $struct_name {
/// Create a `Planner` for an FFT of size `num_points`.
Expand All @@ -35,12 +37,18 @@ macro_rules! impl_planner_for {
/// # Panics
///
/// Panics if `num_points < 1` or if `num_points` is __not__ a power of 2.
pub fn new(num_points: usize, _direction: Direction) -> Self {
pub fn new(num_points: usize, direction: Direction) -> Self {
assert!(num_points > 0 && num_points.is_power_of_two());
let dir = match direction {
Direction::Forward => Direction::Forward,
Direction::Reverse => Direction::Reverse,
};

if num_points <= 4 {
return Self {
twiddles_re: vec![],
twiddles_im: vec![],
direction: dir,
};
}

Expand All @@ -57,6 +65,7 @@ macro_rules! impl_planner_for {
Self {
twiddles_re,
twiddles_im,
direction: dir,
}
}

Expand Down
Loading

0 comments on commit 9db5739

Please sign in to comment.