1use std::f64::consts::PI;
2
3#[derive(Clone, Copy, Debug, Default, PartialEq)]
4#[repr(C)]
5pub struct Complex {
6 pub re: f64,
7 pub im: f64,
8}
9
10impl Complex {
11 #[must_use]
12 pub const fn new(re: f64, im: f64) -> Self {
13 Self { re, im }
14 }
15
16 #[must_use]
17 pub fn mul(self, o: Self) -> Self {
18 Self::new(self.re * o.re - self.im * o.im, self.re * o.im + self.im * o.re)
19 }
20
21 #[must_use]
22 pub fn norm_sqr(self) -> f64 {
23 self.re * self.re + self.im * self.im
24 }
25}
26
27pub fn fft(buf: &mut [Complex], inverse: bool) {
28 fft_with(buf, inverse, slain_rt::cpu::avx2());
29}
30
31thread_local! {
32 static SWAPS: std::cell::RefCell<Vec<Vec<(u32, u32)>>> = const { std::cell::RefCell::new(Vec::new()) };
33 static TWIDDLES: std::cell::RefCell<[Vec<Vec<Complex>>; 2]> = const { std::cell::RefCell::new([Vec::new(), Vec::new()]) };
34}
35
36fn stage_twiddles(len: usize, inverse: bool) -> Vec<Complex> {
37 let sign = if inverse { 1.0 } else { -1.0 };
38 let ang = sign * 2.0 * PI / len as f64;
39 let step = Complex::new(ang.cos(), ang.sin());
40 let mut tw = Vec::with_capacity(len / 2);
41 let mut w = Complex::new(1.0, 0.0);
42 for _ in 0..len / 2 {
43 tw.push(w);
44 w = w.mul(step);
45 }
46 tw
47}
48
49fn bit_reverse(buf: &mut [Complex], bits: u32) {
50 match bits {
51 0..=11 => bit_reverse_pairs(buf, bits),
52 12 | 13 => bit_reverse_tiles(buf, bits, 6, false),
53 14 => bit_reverse_tiles(buf, bits, 5, false),
54 _ => bit_reverse_tiles(buf, bits, 5, true),
55 }
56}
57
58fn bit_reverse_pairs(buf: &mut [Complex], bits: u32) {
59 SWAPS.with(|cache| {
60 let mut cache = cache.borrow_mut();
61 while cache.len() <= bits as usize {
62 let b = cache.len() as u32;
63 let pairs = (0..1usize << b)
64 .filter_map(|i| {
65 let j = if b == 0 { 0 } else { i.reverse_bits() >> (usize::BITS - b) };
66 (j > i).then_some((i as u32, j as u32))
67 })
68 .collect();
69 cache.push(pairs);
70 }
71 for &(i, j) in &cache[bits as usize] {
72 buf.swap(i as usize, j as usize);
73 }
74 });
75}
76
77fn rev(x: usize, bits: u32) -> usize {
78 if bits == 0 {
79 0
80 } else {
81 x.reverse_bits() >> (usize::BITS - bits)
82 }
83}
84
85struct SyncPtr(*mut Complex);
86unsafe impl Sync for SyncPtr {}
87
88struct Tiles {
89 rb: Vec<usize>,
90 cross: Vec<(usize, usize)>,
91 own: Vec<usize>,
92}
93
94thread_local! {
95 static TILES: std::cell::RefCell<Vec<((u32, u32), std::sync::Arc<Tiles>)>> = const { std::cell::RefCell::new(Vec::new()) };
96}
97
98fn tiles(bits: u32, b: u32) -> std::sync::Arc<Tiles> {
99 TILES.with(|c| {
100 let mut c = c.borrow_mut();
101 if let Some((_, t)) = c.iter().find(|(k, _)| *k == (bits, b)) {
102 return std::sync::Arc::clone(t);
103 }
104 let m = bits - 2 * b;
105 let mids = 0..1usize << m;
106 let t = std::sync::Arc::new(Tiles {
107 rb: (0..1usize << b).map(|x| rev(x, b)).collect(),
108 cross: mids.clone().filter_map(|mid| (mid < rev(mid, m)).then_some((mid, rev(mid, m)))).collect(),
109 own: mids.filter(|&mid| mid == rev(mid, m)).collect(),
110 });
111 c.push(((bits, b), std::sync::Arc::clone(&t)));
112 t
113 })
114}
115
116fn bit_reverse_tiles(buf: &mut [Complex], bits: u32, b: u32, parallel: bool) {
117 let b = b.min(bits / 2);
118 let (m, side) = (bits - 2 * b, 1usize << b);
119 let t = tiles(bits, b);
120 let (rb, sh) = (&t.rb[..], m + b);
121 let p = SyncPtr(buf.as_mut_ptr());
122 let cross = |k: usize| {
123 let (mid, mr) = t.cross[k];
124 let ptr = &p;
125 for hi in 0..side {
126 for lo in 0..side {
127 let i = (hi << sh) | (mid << b) | lo;
128 let j = (rb[lo] << sh) | (mr << b) | rb[hi];
129 unsafe { core::ptr::swap(ptr.0.add(i), ptr.0.add(j)) };
130 }
131 }
132 };
133 let own = |k: usize| {
134 let mid = t.own[k];
135 let ptr = &p;
136 for u in 1..side {
137 for v in 0..u {
138 let i = (u << sh) | (mid << b) | rb[v];
139 let j = (v << sh) | (mid << b) | rb[u];
140 unsafe { core::ptr::swap(ptr.0.add(i), ptr.0.add(j)) };
141 }
142 }
143 };
144 if parallel {
145 slain_rt::for_range(t.cross.len(), cross);
146 slain_rt::for_range(t.own.len(), own);
147 } else {
148 (0..t.cross.len()).for_each(cross);
149 (0..t.own.len()).for_each(own);
150 }
151}
152
153fn fft_with(buf: &mut [Complex], inverse: bool, #[cfg_attr(not(target_arch = "x86_64"), expect(unused_variables))] fast: bool) {
154 let n = buf.len();
155 assert!(n.is_power_of_two(), "FFT length must be a power of two");
156 if n < 2 {
157 return;
158 }
159 let bits = n.trailing_zeros();
160 bit_reverse(buf, bits);
161 TWIDDLES.with(|cache| {
162 let mut cache = cache.borrow_mut();
163 let tables = &mut cache[usize::from(inverse)];
164 while tables.len() < bits as usize {
165 tables.push(stage_twiddles(2 << tables.len(), inverse));
166 }
167 #[cfg(target_arch = "x86_64")]
168 let first = if fast && bits >= 4 {
169 unsafe { x86::first_four(buf, [&tables[0], &tables[1], &tables[2], &tables[3]]) };
170 4
171 } else {
172 0
173 };
174 #[cfg(not(target_arch = "x86_64"))]
175 let first = 0;
176 #[cfg(target_arch = "x86_64")]
177 let first = if fast && first > 0 {
178 let mut s = first;
179 while s + 1 < bits as usize {
180 unsafe { x86::stage_pair(buf, &tables[s], &tables[s + 1]) };
181 s += 2;
182 }
183 s
184 } else {
185 first
186 };
187 for (s, tw) in tables[..bits as usize].iter().enumerate().skip(first) {
188 let len = 2usize << s;
189 #[cfg(target_arch = "x86_64")]
190 if fast && len >= 4 {
191 unsafe { x86::stage(buf, tw) };
192 continue;
193 }
194 for block in buf.chunks_exact_mut(len) {
195 let (lo, hi) = block.split_at_mut(len / 2);
196 for ((a, b), &t) in lo.iter_mut().zip(hi.iter_mut()).zip(tw) {
197 let v = b.mul(t);
198 *b = Complex::new(a.re - v.re, a.im - v.im);
199 *a = Complex::new(a.re + v.re, a.im + v.im);
200 }
201 }
202 }
203 });
204 if inverse {
205 let k = 1.0 / n as f64;
206 #[cfg(target_arch = "x86_64")]
207 if fast {
208 unsafe { x86::scale(buf, k) };
209 return;
210 }
211 for v in buf.iter_mut() {
212 v.re *= k;
213 v.im *= k;
214 }
215 }
216}
217
218#[cfg(target_arch = "x86_64")]
219mod x86 {
220 use super::Complex;
221 use core::arch::x86_64::*;
222
223 #[target_feature(enable = "avx2")]
224 pub(super) unsafe fn stage(buf: &mut [Complex], tw: &[Complex]) {
225 let half = tw.len();
226 let p = buf.as_mut_ptr().cast::<f64>();
227 let t = tw.as_ptr().cast::<f64>();
228 let mut base = 0;
229 while base < buf.len() {
230 let mut k = 0;
231 while k < half {
232 let pa = p.add(2 * (base + k));
233 let pb = p.add(2 * (base + k + half));
234 let a = _mm256_loadu_pd(pa);
235 let b = _mm256_loadu_pd(pb);
236 let w = _mm256_loadu_pd(t.add(2 * k));
237 let wr = _mm256_movedup_pd(w);
238 let wi = _mm256_permute_pd::<0b1111>(w);
239 let bs = _mm256_permute_pd::<0b0101>(b);
240 let v = _mm256_addsub_pd(_mm256_mul_pd(b, wr), _mm256_mul_pd(bs, wi));
241 _mm256_storeu_pd(pb, _mm256_sub_pd(a, v));
242 _mm256_storeu_pd(pa, _mm256_add_pd(a, v));
243 k += 2;
244 }
245 base += 2 * half;
246 }
247 }
248
249 #[inline(always)]
250 unsafe fn cmul(b: __m256d, w: __m256d) -> __m256d {
251 let wr = _mm256_movedup_pd(w);
252 let wi = _mm256_permute_pd::<0b1111>(w);
253 let bs = _mm256_permute_pd::<0b0101>(b);
254 _mm256_addsub_pd(_mm256_mul_pd(b, wr), _mm256_mul_pd(bs, wi))
255 }
256
257 #[inline(always)]
258 unsafe fn bfly(a: &mut __m256d, b: &mut __m256d, w: __m256d) {
259 let v = cmul(*b, w);
260 *b = _mm256_sub_pd(*a, v);
261 *a = _mm256_add_pd(*a, v);
262 }
263
264 #[target_feature(enable = "avx2")]
265 pub(super) unsafe fn first_four(buf: &mut [Complex], tw: [&[Complex]; 4]) {
266 let p = buf.as_mut_ptr().cast::<f64>();
267 let t2 = tw[0].as_ptr().cast::<f64>();
268 let w2 = _mm256_castpd128_pd256(_mm_loadu_pd(t2));
269 let w2 = _mm256_insertf128_pd::<1>(w2, _mm_loadu_pd(t2));
270 let w4 = _mm256_loadu_pd(tw[1].as_ptr().cast::<f64>());
271 let t8 = tw[2].as_ptr().cast::<f64>();
272 let t16 = tw[3].as_ptr().cast::<f64>();
273 let mut c = 0;
274 while c < buf.len() {
275 let q = p.add(2 * c);
276 let mut r: [__m256d; 8] = core::array::from_fn(|k| _mm256_loadu_pd(q.add(4 * k)));
277 for x in &mut r {
278 let a = _mm256_permute2f128_pd::<0x00>(*x, *x);
279 let b = _mm256_permute2f128_pd::<0x11>(*x, *x);
280 let v = cmul(b, w2);
281 *x = _mm256_blend_pd::<0b1100>(_mm256_add_pd(a, v), _mm256_sub_pd(a, v));
282 }
283 for g in 0..4 {
284 let (lo, hi) = r.split_at_mut(2 * g + 1);
285 bfly(&mut lo[2 * g], &mut hi[0], w4);
286 }
287 for g in 0..2 {
288 for j in 0..2 {
289 let w = _mm256_loadu_pd(t8.add(4 * j));
290 let (lo, hi) = r.split_at_mut(4 * g + j + 2);
291 bfly(&mut lo[4 * g + j], &mut hi[0], w);
292 }
293 }
294 for k in 0..4 {
295 let w = _mm256_loadu_pd(t16.add(4 * k));
296 let (lo, hi) = r.split_at_mut(k + 4);
297 bfly(&mut lo[k], &mut hi[0], w);
298 }
299 for (k, x) in r.iter().enumerate() {
300 _mm256_storeu_pd(q.add(4 * k), *x);
301 }
302 c += 16;
303 }
304 }
305
306 #[target_feature(enable = "avx2")]
307 pub(super) unsafe fn stage_pair(buf: &mut [Complex], tw_a: &[Complex], tw_b: &[Complex]) {
308 let h = tw_a.len();
309 let p = buf.as_mut_ptr().cast::<f64>();
310 let (ta, tb) = (tw_a.as_ptr().cast::<f64>(), tw_b.as_ptr().cast::<f64>());
311 let mut base = 0;
312 while base < buf.len() {
313 let mut k = 0;
314 while k < h {
315 let q = [base + k, base + k + h, base + k + 2 * h, base + k + 3 * h].map(|e| p.add(2 * e));
316 let [mut e0, mut e1, mut e2, mut e3] = q.map(|x| _mm256_loadu_pd(x));
317 let wa = _mm256_loadu_pd(ta.add(2 * k));
318 bfly(&mut e0, &mut e1, wa);
319 bfly(&mut e2, &mut e3, wa);
320 bfly(&mut e0, &mut e2, _mm256_loadu_pd(tb.add(2 * k)));
321 bfly(&mut e1, &mut e3, _mm256_loadu_pd(tb.add(2 * (k + h))));
322 for (x, v) in q.into_iter().zip([e0, e1, e2, e3]) {
323 _mm256_storeu_pd(x, v);
324 }
325 k += 2;
326 }
327 base += 4 * h;
328 }
329 }
330
331 #[target_feature(enable = "avx2")]
332 pub(super) unsafe fn scale(buf: &mut [Complex], k: f64) {
333 let p = buf.as_mut_ptr().cast::<f64>();
334 let kv = _mm256_set1_pd(k);
335 let n = buf.len() * 2;
336 let mut i = 0;
337 while i + 4 <= n {
338 _mm256_storeu_pd(p.add(i), _mm256_mul_pd(_mm256_loadu_pd(p.add(i)), kv));
339 i += 4;
340 }
341 while i < n {
342 *p.add(i) *= k;
343 i += 1;
344 }
345 }
346}
347
348#[must_use]
349pub fn convolve(a: &[f64], b: &[f64]) -> Vec<f64> {
350 if a.is_empty() || b.is_empty() {
351 return Vec::new();
352 }
353 let out_len = a.len() + b.len() - 1;
354 let n = out_len.next_power_of_two();
355 let load = |s: &[f64]| {
356 let mut v = vec![Complex::default(); n];
357 for (d, &x) in v.iter_mut().zip(s) {
358 d.re = x;
359 }
360 v
361 };
362 let (mut fa, mut fb) = (load(a), load(b));
363 fft(&mut fa, false);
364 fft(&mut fb, false);
365 for (x, y) in fa.iter_mut().zip(&fb) {
366 *x = x.mul(*y);
367 }
368 fft(&mut fa, true);
369 fa[..out_len].iter().map(|c| c.re).collect()
370}
371
372#[cfg(test)]
373mod tests {
374 use super::*;
375
376 fn signal(n: usize, seed: f64) -> Vec<Complex> {
377 (0..n).map(|i| Complex::new(((i as f64 + seed) * 0.37).sin() * 1e3, ((i as f64 * 1.3 + seed) * 0.11).cos() - 0.5)).collect()
378 }
379
380 fn reference(buf: &mut [Complex], inverse: bool) {
381 let n = buf.len();
382 let bits = n.trailing_zeros();
383 for i in 0..n {
384 let j = i.reverse_bits() >> (usize::BITS - bits);
385 if j > i {
386 buf.swap(i, j);
387 }
388 }
389 let mut len = 2;
390 while len <= n {
391 let tw = stage_twiddles(len, inverse);
392 for block in buf.chunks_exact_mut(len) {
393 let (lo, hi) = block.split_at_mut(len / 2);
394 for ((a, b), &t) in lo.iter_mut().zip(hi.iter_mut()).zip(&tw) {
395 let v = b.mul(t);
396 *b = Complex::new(a.re - v.re, a.im - v.im);
397 *a = Complex::new(a.re + v.re, a.im + v.im);
398 }
399 }
400 len <<= 1;
401 }
402 if inverse {
403 let k = 1.0 / n as f64;
404 for v in buf.iter_mut() {
405 v.re *= k;
406 v.im *= k;
407 }
408 }
409 }
410
411 #[test]
412 fn scalar_and_avx2_equal_the_reference_bit_for_bit() {
413 let same = |a: &[Complex], b: &[Complex]| a.iter().zip(b).all(|(x, y)| x.re.to_bits() == y.re.to_bits() && x.im.to_bits() == y.im.to_bits());
414 for bits in 1..=12 {
415 let n = 1usize << bits;
416 for inverse in [false, true] {
417 let mut want = signal(n, bits as f64);
418 let mut scalar = want.clone();
419 let mut avx = want.clone();
420 reference(&mut want, inverse);
421 fft_with(&mut scalar, inverse, false);
422 assert!(same(&want, &scalar), "n {n} inverse {inverse}: scalar differs from the reference");
423 if slain_rt::cpu::avx2() {
424 fft_with(&mut avx, inverse, true);
425 assert!(same(&want, &avx), "n {n} inverse {inverse}: AVX2 differs from the reference");
426 }
427 }
428 }
429 }
430
431 #[test]
432 fn every_reversal_equals_the_direct_permutation() {
433 for bits in 1..=17u32 {
434 let src = signal(1 << bits, f64::from(bits));
435 let mut want = src.clone();
436 for (i, &v) in src.iter().enumerate() {
437 want[rev(i, bits)] = v;
438 }
439 let mut got = src.clone();
440 bit_reverse(&mut got, bits);
441 assert!(want == got, "bit_reverse bits {bits}");
442 let mut got = src.clone();
443 bit_reverse_pairs(&mut got, bits);
444 assert!(want == got, "pairs bits {bits}");
445 for b in [3u32, 4, 5, 6] {
446 for parallel in [false, true] {
447 let mut got = src.clone();
448 bit_reverse_tiles(&mut got, bits, b, parallel);
449 assert!(want == got, "tiles bits {bits} b {b} parallel {parallel}");
450 }
451 }
452 }
453 }
454
455 #[test]
456 #[ignore]
457 fn time_fork_join() {
458 for jobs in [0usize, 1, 2, 3, 16, 64] {
459 let reps = 20000;
460 let mut best = f64::MAX;
461 for _ in 0..11 {
462 let t = std::time::Instant::now();
463 for _ in 0..reps {
464 slain_rt::for_range(crate::bench::slain(jobs), |k| {
465 crate::bench::slain(k);
466 });
467 }
468 best = best.min(t.elapsed().as_secs_f64());
469 }
470 println!("FORKJOIN jobs {jobs:3} {:.9} us per for_range", best / f64::from(reps) * 1e6);
471 }
472 }
473
474 #[test]
475 #[ignore]
476 fn time_tiled_reversal() {
477 for bits in 10u32..=16 {
478 let n = 1usize << bits;
479 let src = signal(n, 1.0);
480 let reps = ((1 << 24) / n).max(4);
481 let time = |f: &mut dyn FnMut(&mut [Complex])| {
482 let mut best = f64::MAX;
483 for _ in 0..11 {
484 let mut v = src.clone();
485 let t = std::time::Instant::now();
486 for _ in 0..reps {
487 f(crate::bench::slain(&mut v));
488 }
489 best = best.min(t.elapsed().as_secs_f64());
490 }
491 best / reps as f64 * 1e6
492 };
493 let pairs = time(&mut |v| bit_reverse_pairs(v, bits));
494 let mut line = format!("TILES n {n:6} pairs {pairs:.9} us");
495 for b in [3u32, 4, 5, 6] {
496 for parallel in [false, true] {
497 let t = time(&mut |v| bit_reverse_tiles(v, bits, b, parallel));
498 line += &format!(" b{b}{} {t:.9} us {:.9}x", if parallel { "par" } else { "" }, pairs / t);
499 }
500 }
501 println!("{line}");
502 }
503 }
504
505 #[test]
506 #[ignore]
507 fn time_bitrev_share() {
508 for n in [64usize, 256, 1024, 4096] {
509 let src = signal(n, 1.0);
510 let reps = (1 << 22) / n;
511 let time = |f: &mut dyn FnMut(&mut [Complex])| {
512 let mut best = f64::MAX;
513 for _ in 0..15 {
514 let mut v = src.clone();
515 let t = std::time::Instant::now();
516 for _ in 0..reps {
517 f(crate::bench::slain(&mut v));
518 }
519 best = best.min(t.elapsed().as_secs_f64());
520 }
521 best / reps as f64
522 };
523 let whole = time(&mut |v| fft_with(v, false, true));
524 let inline = time(&mut |v| {
525 let bits = v.len().trailing_zeros();
526 for i in 0..v.len() {
527 let j = i.reverse_bits() >> (usize::BITS - bits);
528 if j > i {
529 v.swap(i, j);
530 }
531 }
532 });
533 let cached = time(&mut |v| bit_reverse(v, v.len().trailing_zeros()));
534 println!(
535 "BITREV n {n:5} whole {:.9} us inline {:.9} us cached {:.9} us cached share {:.9}% cached/inline {:.9}x",
536 whole * 1e6,
537 inline * 1e6,
538 cached * 1e6,
539 100.0 * cached / whole,
540 inline / cached
541 );
542 }
543 }
544
545 #[test]
546 #[ignore]
547 fn time_fft() {
548 for n in [64usize, 256, 1024, 4096] {
549 let src = signal(n, 1.0);
550 let mut rate = [0f64; 3];
551 for arm in 0..3 {
552 let reps = (1 << 22) / n;
553 let mut best = f64::MAX;
554 for _ in 0..15 {
555 let mut v = src.clone();
556 let t = std::time::Instant::now();
557 for r in 0..reps {
558 let v = crate::bench::slain(&mut v);
559 match arm {
560 0 => reference(v, r % 2 == 1),
561 1 => fft_with(v, r % 2 == 1, false),
562 _ => fft_with(v, r % 2 == 1, true),
563 }
564 }
565 best = best.min(t.elapsed().as_secs_f64());
566 }
567 rate[arm] = reps as f64 / best;
568 }
569 println!(
570 "FFT n {n:5} reference {:.9} scalar+cache {:.9} avx2+cache {:.9} calls/s scalar+cache {:.9}x avx2+cache {:.9}x",
571 rate[0],
572 rate[1],
573 rate[2],
574 rate[1] / rate[0],
575 rate[2] / rate[0]
576 );
577 }
578 }
579
580 #[test]
581 fn round_trip_and_convolution() {
582 let orig: Vec<Complex> = (0..64).map(|i| Complex::new((i as f64 * 0.3).sin(), (i as f64 * 0.7).cos())).collect();
583 let mut v = orig.clone();
584 fft(&mut v, false);
585 let s = orig.iter().fold(Complex::default(), |a, b| Complex::new(a.re + b.re, a.im + b.im));
586 assert!((v[0].re - s.re).abs() < 1e-9 && (v[0].im - s.im).abs() < 1e-9);
587 fft(&mut v, true);
588 assert!(v.iter().zip(&orig).all(|(a, b)| (a.re - b.re).abs() < 1e-12 && (a.im - b.im).abs() < 1e-12));
589 let (a, b) = ([1.0, 2.0, 3.0], [0.5, -1.0]);
590 let c = convolve(&a, &b);
591 let want = [0.5, 0.0, -0.5, -3.0];
592 assert!(c.iter().zip(want).all(|(g, w)| (g - w).abs() < 1e-12));
593 }
594}