Skip to content

Commit 9db5739

Browse files
committed
Make planner reusable
- 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`
1 parent 2df2f00 commit 9db5739

File tree

6 files changed

+272
-43
lines changed

6 files changed

+272
-43
lines changed

Cargo.toml

+7-1
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,14 @@ num-traits = "0.2.18"
1515
multiversion = "0.7"
1616

1717
[dev-dependencies]
18-
utilities = { path = "utilities" }
18+
criterion = "0.5.1"
1919
fftw = "0.8.0"
20+
rand = "0.8.5"
21+
utilities = { path = "utilities" }
22+
23+
[[bench]]
24+
name = "bench"
25+
harness = false
2026

2127
[profile.release]
2228
codegen-units = 1

benches/bench.rs

+119
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
use criterion::{black_box, criterion_group, criterion_main, Criterion};
2+
use num_traits::Float;
3+
use phastft::{
4+
fft_32_with_opts_and_plan, fft_64_with_opts_and_plan,
5+
options::Options,
6+
planner::{Direction, Planner32, Planner64},
7+
};
8+
use rand::{
9+
distributions::{Distribution, Standard},
10+
thread_rng, Rng,
11+
};
12+
13+
const LENGTHS: &[usize] = &[
14+
6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28,
15+
];
16+
17+
fn generate_numbers<T: Float>(n: usize) -> (Vec<T>, Vec<T>)
18+
where
19+
Standard: Distribution<T>,
20+
{
21+
let mut rng = thread_rng();
22+
23+
let samples: Vec<T> = (&mut rng).sample_iter(Standard).take(2 * n).collect();
24+
25+
let mut reals = vec![T::zero(); n];
26+
let mut imags = vec![T::zero(); n];
27+
28+
for ((z_re, z_im), rand_chunk) in reals
29+
.iter_mut()
30+
.zip(imags.iter_mut())
31+
.zip(samples.chunks_exact(2))
32+
{
33+
*z_re = rand_chunk[0];
34+
*z_im = rand_chunk[1];
35+
}
36+
37+
(reals, imags)
38+
}
39+
40+
fn benchmark_forward_f32(c: &mut Criterion) {
41+
let options = Options::default();
42+
43+
for n in LENGTHS.iter() {
44+
let len = 1 << n;
45+
let id = format!("FFT Forward f32 {} elements", len);
46+
c.bench_function(&id, |b| {
47+
let (mut reals, mut imags) = generate_numbers(len);
48+
let planner = Planner32::new(len, Direction::Forward);
49+
b.iter(|| {
50+
black_box(fft_32_with_opts_and_plan(
51+
&mut reals, &mut imags, &options, &planner,
52+
));
53+
});
54+
});
55+
}
56+
}
57+
58+
fn benchmark_inverse_f32(c: &mut Criterion) {
59+
let options = Options::default();
60+
61+
for n in LENGTHS.iter() {
62+
let len = 1 << n;
63+
let id = format!("FFT Inverse f32 {} elements", len);
64+
c.bench_function(&id, |b| {
65+
let (mut reals, mut imags) = generate_numbers(len);
66+
let planner = Planner32::new(len, Direction::Reverse);
67+
b.iter(|| {
68+
black_box(fft_32_with_opts_and_plan(
69+
&mut reals, &mut imags, &options, &planner,
70+
));
71+
});
72+
});
73+
}
74+
}
75+
76+
fn benchmark_forward_f64(c: &mut Criterion) {
77+
let options = Options::default();
78+
79+
for n in LENGTHS.iter() {
80+
let len = 1 << n;
81+
let id = format!("FFT Forward f64 {} elements", len);
82+
c.bench_function(&id, |b| {
83+
let (mut reals, mut imags) = generate_numbers(len);
84+
let planner = Planner64::new(len, Direction::Forward);
85+
b.iter(|| {
86+
black_box(fft_64_with_opts_and_plan(
87+
&mut reals, &mut imags, &options, &planner,
88+
));
89+
});
90+
});
91+
}
92+
}
93+
94+
fn benchmark_inverse_f64(c: &mut Criterion) {
95+
let options = Options::default();
96+
97+
for n in LENGTHS.iter() {
98+
let len = 1 << n;
99+
let id = format!("FFT Inverse f64 {} elements", len);
100+
c.bench_function(&id, |b| {
101+
let (mut reals, mut imags) = generate_numbers(len);
102+
let planner = Planner64::new(len, Direction::Reverse);
103+
b.iter(|| {
104+
black_box(fft_64_with_opts_and_plan(
105+
&mut reals, &mut imags, &options, &planner,
106+
));
107+
});
108+
});
109+
}
110+
}
111+
112+
criterion_group!(
113+
benches,
114+
benchmark_forward_f32,
115+
benchmark_inverse_f32,
116+
benchmark_forward_f64,
117+
benchmark_inverse_f64
118+
);
119+
criterion_main!(benches);

src/lib.rs

+90-29
Original file line numberDiff line numberDiff line change
@@ -48,35 +48,15 @@ macro_rules! impl_fft_for {
4848
imags.len()
4949
);
5050

51-
let mut planner = <$planner>::new(reals.len(), Direction::Forward);
51+
let planner = <$planner>::new(reals.len(), direction);
5252
assert!(
5353
planner.num_twiddles().is_power_of_two()
5454
&& planner.num_twiddles() == reals.len() / 2
5555
);
5656

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

59-
match direction {
60-
Direction::Reverse => {
61-
for z_im in imags.iter_mut() {
62-
*z_im = -*z_im;
63-
}
64-
}
65-
_ => (),
66-
}
67-
68-
$opts_and_plan(reals, imags, &opts, &mut planner);
69-
70-
match direction {
71-
Direction::Reverse => {
72-
let scaling_factor = (reals.len() as $precision).recip();
73-
for (z_re, z_im) in reals.iter_mut().zip(imags.iter_mut()) {
74-
*z_re *= scaling_factor;
75-
*z_im *= -scaling_factor;
76-
}
77-
}
78-
_ => (),
79-
}
59+
$opts_and_plan(reals, imags, &opts, &planner);
8060
}
8161
};
8262
}
@@ -112,30 +92,39 @@ macro_rules! impl_fft_with_opts_and_plan_for {
11292
reals: &mut [$precision],
11393
imags: &mut [$precision],
11494
opts: &Options,
115-
planner: &mut $planner,
95+
planner: &$planner,
11696
) {
11797
assert!(reals.len() == imags.len() && reals.len().is_power_of_two());
11898
let n: usize = reals.len().ilog2() as usize;
11999

120-
let twiddles_re = &mut planner.twiddles_re;
121-
let twiddles_im = &mut planner.twiddles_im;
100+
let mut twiddles_re = planner.twiddles_re.clone();
101+
let mut twiddles_im = planner.twiddles_im.clone();
122102

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

107+
match planner.direction {
108+
Direction::Reverse => {
109+
for z_im in imags.iter_mut() {
110+
*z_im = -*z_im;
111+
}
112+
}
113+
_ => (),
114+
}
115+
127116
for t in (0..n).rev() {
128117
let dist = 1 << t;
129118
let chunk_size = dist << 1;
130119

131120
if chunk_size > 4 {
132121
if t < n - 1 {
133-
filter_twiddles(twiddles_re, twiddles_im);
122+
(twiddles_re, twiddles_im) = filter_twiddles(&twiddles_re, &twiddles_im);
134123
}
135124
if chunk_size >= $lanes * 2 {
136-
$simd_butterfly_kernel(reals, imags, twiddles_re, twiddles_im, dist);
125+
$simd_butterfly_kernel(reals, imags, &twiddles_re, &twiddles_im, dist);
137126
} else {
138-
fft_chunk_n(reals, imags, twiddles_re, twiddles_im, dist);
127+
fft_chunk_n(reals, imags, &twiddles_re, &twiddles_im, dist);
139128
}
140129
} else if chunk_size == 2 {
141130
fft_chunk_2(reals, imags);
@@ -153,6 +142,17 @@ macro_rules! impl_fft_with_opts_and_plan_for {
153142
cobra_apply(reals, n);
154143
cobra_apply(imags, n);
155144
}
145+
146+
match planner.direction {
147+
Direction::Reverse => {
148+
let scaling_factor = (reals.len() as $precision).recip();
149+
for (z_re, z_im) in reals.iter_mut().zip(imags.iter_mut()) {
150+
*z_re *= scaling_factor;
151+
*z_im *= -scaling_factor;
152+
}
153+
}
154+
_ => (),
155+
}
156156
}
157157
};
158158
}
@@ -179,7 +179,7 @@ mod tests {
179179

180180
use utilities::rustfft::num_complex::Complex;
181181
use utilities::rustfft::FftPlanner;
182-
use utilities::{assert_float_closeness, gen_random_signal};
182+
use utilities::{assert_float_closeness, gen_random_signal, gen_random_signal_f32};
183183

184184
use super::*;
185185

@@ -308,4 +308,65 @@ mod tests {
308308
}
309309
}
310310
}
311+
312+
#[test]
313+
fn fft_64_with_opts_and_plan_vs_fft_64() {
314+
let num_points = 4096;
315+
316+
let mut reals = vec![0.0; num_points];
317+
let mut imags = vec![0.0; num_points];
318+
gen_random_signal(&mut reals, &mut imags);
319+
320+
let mut re = reals.clone();
321+
let mut im = imags.clone();
322+
323+
let mut planner = Planner64::new(num_points, Direction::Forward);
324+
let opts = Options::guess_options(reals.len());
325+
fft_64_with_opts_and_plan(&mut reals, &mut imags, &opts, &mut planner);
326+
327+
fft_64(&mut re, &mut im, Direction::Forward);
328+
329+
reals
330+
.iter()
331+
.zip(imags.iter())
332+
.zip(re.iter())
333+
.zip(im.iter())
334+
.for_each(|(((r, i), z_re), z_im)| {
335+
assert_float_closeness(*r, *z_re, 1e-6);
336+
assert_float_closeness(*i, *z_im, 1e-6);
337+
});
338+
}
339+
340+
#[test]
341+
fn fft_32_with_opts_and_plan_vs_fft_64() {
342+
let dirs = [Direction::Forward, Direction::Reverse];
343+
344+
for direction in dirs {
345+
for n in 4..14 {
346+
let num_points = 1 << n;
347+
let mut reals = vec![0.0; num_points];
348+
let mut imags = vec![0.0; num_points];
349+
gen_random_signal_f32(&mut reals, &mut imags);
350+
351+
let mut re = reals.clone();
352+
let mut im = imags.clone();
353+
354+
let mut planner = Planner32::new(num_points, direction);
355+
let opts = Options::guess_options(reals.len());
356+
fft_32_with_opts_and_plan(&mut reals, &mut imags, &opts, &mut planner);
357+
358+
fft_32(&mut re, &mut im, direction);
359+
360+
reals
361+
.iter()
362+
.zip(imags.iter())
363+
.zip(re.iter())
364+
.zip(im.iter())
365+
.for_each(|(((r, i), z_re), z_im)| {
366+
assert_float_closeness(*r, *z_re, 1e-6);
367+
assert_float_closeness(*i, *z_im, 1e-6);
368+
});
369+
}
370+
}
371+
}
311372
}

src/planner.rs

+10-1
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ macro_rules! impl_planner_for {
2525
pub twiddles_re: Vec<$precision>,
2626
/// The imaginary components of the twiddle factors
2727
pub twiddles_im: Vec<$precision>,
28+
/// The direction of the FFT associated with this `Planner`
29+
pub direction: Direction,
2830
}
2931
impl $struct_name {
3032
/// Create a `Planner` for an FFT of size `num_points`.
@@ -35,12 +37,18 @@ macro_rules! impl_planner_for {
3537
/// # Panics
3638
///
3739
/// Panics if `num_points < 1` or if `num_points` is __not__ a power of 2.
38-
pub fn new(num_points: usize, _direction: Direction) -> Self {
40+
pub fn new(num_points: usize, direction: Direction) -> Self {
3941
assert!(num_points > 0 && num_points.is_power_of_two());
42+
let dir = match direction {
43+
Direction::Forward => Direction::Forward,
44+
Direction::Reverse => Direction::Reverse,
45+
};
46+
4047
if num_points <= 4 {
4148
return Self {
4249
twiddles_re: vec![],
4350
twiddles_im: vec![],
51+
direction: dir,
4452
};
4553
}
4654

@@ -57,6 +65,7 @@ macro_rules! impl_planner_for {
5765
Self {
5866
twiddles_re,
5967
twiddles_im,
68+
direction: dir,
6069
}
6170
}
6271

0 commit comments

Comments
 (0)