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