use std::f64::consts::PI;
#[derive(Clone, Copy, Debug, Default, PartialEq)]
#[repr(C)]
pub struct Complex {
pub re: f64,
pub im: f64,
}
impl Complex {
#[must_use]
pub const fn new(re: f64, im: f64) -> Self {
Self { re, im }
}
#[must_use]
pub fn norm_sqr(self) -> f64 {
self.re * self.re + self.im * self.im
}
}
impl core::ops::Mul for Complex {
type Output = Self;
#[inline]
fn mul(self, o: Self) -> Self {
Self::new(self.re * o.re - self.im * o.im, self.re * o.im + self.im * o.re)
}
}
impl core::ops::MulAssign for Complex {
#[inline]
fn mul_assign(&mut self, o: Self) {
*self = *self * o;
}
}
pub fn fft(buf: &mut [Complex], inverse: bool) {
fft_with(buf, inverse, slain_rt::cpu::avx2());
}
thread_local! {
static SWAPS: std::cell::RefCell<Vec<Vec<(u32, u32)>>> = const { std::cell::RefCell::new(Vec::new()) };
static TWIDDLES: std::cell::RefCell<[Vec<Vec<Complex>>; 2]> = const { std::cell::RefCell::new([Vec::new(), Vec::new()]) };
}
fn stage_twiddles(len: usize, inverse: bool) -> Vec<Complex> {
let sign = if inverse { 1.0 } else { -1.0 };
let ang = sign * 2.0 * PI / len as f64;
let step = Complex::new(ang.cos(), ang.sin());
let mut tw = Vec::with_capacity(len / 2);
let mut w = Complex::new(1.0, 0.0);
for _ in 0..len / 2 {
tw.push(w);
w *= step;
}
tw
}
fn bit_reverse(buf: &mut [Complex], bits: u32) {
match bits {
0..=11 => bit_reverse_pairs(buf, bits),
12 | 13 => bit_reverse_tiles(buf, bits, 6, false),
14 => bit_reverse_tiles(buf, bits, 5, false),
_ => bit_reverse_tiles(buf, bits, 5, true),
}
}
fn bit_reverse_pairs(buf: &mut [Complex], bits: u32) {
SWAPS.with(|cache| {
let mut cache = cache.borrow_mut();
while cache.len() <= bits as usize {
let b = cache.len() as u32;
let pairs = (0..1usize << b)
.filter_map(|i| {
let j = if b == 0 { 0 } else { i.reverse_bits() >> (usize::BITS - b) };
(j > i).then_some((i as u32, j as u32))
})
.collect();
cache.push(pairs);
}
for &(i, j) in &cache[bits as usize] {
buf.swap(i as usize, j as usize);
}
});
}
fn rev(x: usize, bits: u32) -> usize {
if bits == 0 {
0
} else {
x.reverse_bits() >> (usize::BITS - bits)
}
}
struct SyncPtr(*mut Complex);
unsafe impl Sync for SyncPtr {}
struct Tiles {
rb: Vec<usize>,
cross: Vec<(usize, usize)>,
own: Vec<usize>,
}
type TileCache = Vec<((u32, u32), std::sync::Arc<Tiles>)>;
thread_local! {
static TILES: std::cell::RefCell<TileCache> = const { std::cell::RefCell::new(Vec::new()) };
}
fn tiles(bits: u32, b: u32) -> std::sync::Arc<Tiles> {
TILES.with(|c| {
let mut c = c.borrow_mut();
if let Some((_, t)) = c.iter().find(|(k, _)| *k == (bits, b)) {
return std::sync::Arc::clone(t);
}
let m = bits - 2 * b;
let mids = 0..1usize << m;
let t = std::sync::Arc::new(Tiles {
rb: (0..1usize << b).map(|x| rev(x, b)).collect(),
cross: mids.clone().filter_map(|mid| (mid < rev(mid, m)).then_some((mid, rev(mid, m)))).collect(),
own: mids.filter(|&mid| mid == rev(mid, m)).collect(),
});
c.push(((bits, b), std::sync::Arc::clone(&t)));
t
})
}
fn bit_reverse_tiles(buf: &mut [Complex], bits: u32, b: u32, parallel: bool) {
let b = b.min(bits / 2);
let (m, side) = (bits - 2 * b, 1usize << b);
let t = tiles(bits, b);
let (rb, sh) = (&t.rb[..], m + b);
let p = SyncPtr(buf.as_mut_ptr());
let cross = |k: usize| {
let (mid, mr) = t.cross[k];
let ptr = &p;
for hi in 0..side {
for lo in 0..side {
let i = (hi << sh) | (mid << b) | lo;
let j = (rb[lo] << sh) | (mr << b) | rb[hi];
unsafe { core::ptr::swap(ptr.0.add(i), ptr.0.add(j)) };
}
}
};
let own = |k: usize| {
let mid = t.own[k];
let ptr = &p;
for u in 1..side {
for v in 0..u {
let i = (u << sh) | (mid << b) | rb[v];
let j = (v << sh) | (mid << b) | rb[u];
unsafe { core::ptr::swap(ptr.0.add(i), ptr.0.add(j)) };
}
}
};
if parallel {
slain_rt::for_range(t.cross.len(), cross);
slain_rt::for_range(t.own.len(), own);
} else {
(0..t.cross.len()).for_each(cross);
(0..t.own.len()).for_each(own);
}
}
fn fft_with(buf: &mut [Complex], inverse: bool, fast: bool) {
debug_assert!(!fast || slain_rt::cpu::avx2(), "AVX2 FFT kernels requested on a CPU without AVX2");
let n = buf.len();
assert!(n.is_power_of_two(), "FFT length must be a power of two");
if n < 2 {
return;
}
let bits = n.trailing_zeros();
bit_reverse(buf, bits);
TWIDDLES.with(|cache| {
let mut cache = cache.borrow_mut();
let tables = &mut cache[usize::from(inverse)];
while tables.len() < bits as usize {
tables.push(stage_twiddles(2 << tables.len(), inverse));
}
#[cfg(target_arch = "x86_64")]
let first = if fast && bits >= 4 {
unsafe { x86::first_four(buf, [&tables[0], &tables[1], &tables[2], &tables[3]]) };
4
} else {
0
};
#[cfg(not(target_arch = "x86_64"))]
let first = 0;
#[cfg(target_arch = "x86_64")]
let first = if fast && first > 0 {
let mut s = first;
while s + 1 < bits as usize {
unsafe { x86::stage_pair(buf, &tables[s], &tables[s + 1]) };
s += 2;
}
s
} else {
first
};
for (s, tw) in tables[..bits as usize].iter().enumerate().skip(first) {
let len = 2usize << s;
#[cfg(target_arch = "x86_64")]
if fast && len >= 4 {
unsafe { x86::stage(buf, tw) };
continue;
}
for block in buf.chunks_exact_mut(len) {
let (lo, hi) = block.split_at_mut(len / 2);
for ((a, b), &t) in lo.iter_mut().zip(hi.iter_mut()).zip(tw) {
let v = *b * t;
*b = Complex::new(a.re - v.re, a.im - v.im);
*a = Complex::new(a.re + v.re, a.im + v.im);
}
}
}
});
if inverse {
let k = 1.0 / n as f64;
#[cfg(target_arch = "x86_64")]
if fast {
unsafe { x86::scale(buf, k) };
return;
}
for v in buf.iter_mut() {
v.re *= k;
v.im *= k;
}
}
}
#[cfg(target_arch = "x86_64")]
mod x86 {
use super::Complex;
use core::arch::x86_64::*;
#[target_feature(enable = "avx2")]
pub(super) unsafe fn stage(buf: &mut [Complex], tw: &[Complex]) {
let half = tw.len();
let p = buf.as_mut_ptr().cast::<f64>();
let t = tw.as_ptr().cast::<f64>();
let mut base = 0;
while base < buf.len() {
let mut k = 0;
while k < half {
let pa = p.add(2 * (base + k));
let pb = p.add(2 * (base + k + half));
let a = _mm256_loadu_pd(pa);
let b = _mm256_loadu_pd(pb);
let w = _mm256_loadu_pd(t.add(2 * k));
let wr = _mm256_movedup_pd(w);
let wi = _mm256_permute_pd::<0b1111>(w);
let bs = _mm256_permute_pd::<0b0101>(b);
let v = _mm256_addsub_pd(_mm256_mul_pd(b, wr), _mm256_mul_pd(bs, wi));
_mm256_storeu_pd(pb, _mm256_sub_pd(a, v));
_mm256_storeu_pd(pa, _mm256_add_pd(a, v));
k += 2;
}
base += 2 * half;
}
}
#[inline(always)]
unsafe fn cmul(b: __m256d, w: __m256d) -> __m256d {
let wr = _mm256_movedup_pd(w);
let wi = _mm256_permute_pd::<0b1111>(w);
let bs = _mm256_permute_pd::<0b0101>(b);
_mm256_addsub_pd(_mm256_mul_pd(b, wr), _mm256_mul_pd(bs, wi))
}
#[inline(always)]
unsafe fn bfly(a: &mut __m256d, b: &mut __m256d, w: __m256d) {
let v = cmul(*b, w);
*b = _mm256_sub_pd(*a, v);
*a = _mm256_add_pd(*a, v);
}
#[target_feature(enable = "avx2")]
pub(super) unsafe fn first_four(buf: &mut [Complex], tw: [&[Complex]; 4]) {
let p = buf.as_mut_ptr().cast::<f64>();
let t2 = tw[0].as_ptr().cast::<f64>();
let w2 = _mm256_castpd128_pd256(_mm_loadu_pd(t2));
let w2 = _mm256_insertf128_pd::<1>(w2, _mm_loadu_pd(t2));
let w4 = _mm256_loadu_pd(tw[1].as_ptr().cast::<f64>());
let t8 = tw[2].as_ptr().cast::<f64>();
let t16 = tw[3].as_ptr().cast::<f64>();
let mut c = 0;
while c < buf.len() {
let q = p.add(2 * c);
let mut r: [__m256d; 8] = core::array::from_fn(|k| _mm256_loadu_pd(q.add(4 * k)));
for x in &mut r {
let a = _mm256_permute2f128_pd::<0x00>(*x, *x);
let b = _mm256_permute2f128_pd::<0x11>(*x, *x);
let v = cmul(b, w2);
*x = _mm256_blend_pd::<0b1100>(_mm256_add_pd(a, v), _mm256_sub_pd(a, v));
}
for g in 0..4 {
let (lo, hi) = r.split_at_mut(2 * g + 1);
bfly(&mut lo[2 * g], &mut hi[0], w4);
}
for g in 0..2 {
for j in 0..2 {
let w = _mm256_loadu_pd(t8.add(4 * j));
let (lo, hi) = r.split_at_mut(4 * g + j + 2);
bfly(&mut lo[4 * g + j], &mut hi[0], w);
}
}
for k in 0..4 {
let w = _mm256_loadu_pd(t16.add(4 * k));
let (lo, hi) = r.split_at_mut(k + 4);
bfly(&mut lo[k], &mut hi[0], w);
}
for (k, x) in r.iter().enumerate() {
_mm256_storeu_pd(q.add(4 * k), *x);
}
c += 16;
}
}
#[target_feature(enable = "avx2")]
pub(super) unsafe fn stage_pair(buf: &mut [Complex], tw_a: &[Complex], tw_b: &[Complex]) {
let h = tw_a.len();
let p = buf.as_mut_ptr().cast::<f64>();
let (ta, tb) = (tw_a.as_ptr().cast::<f64>(), tw_b.as_ptr().cast::<f64>());
let mut base = 0;
while base < buf.len() {
let mut k = 0;
while k < h {
let q = [base + k, base + k + h, base + k + 2 * h, base + k + 3 * h].map(|e| p.add(2 * e));
let [mut e0, mut e1, mut e2, mut e3] = q.map(|x| _mm256_loadu_pd(x));
let wa = _mm256_loadu_pd(ta.add(2 * k));
bfly(&mut e0, &mut e1, wa);
bfly(&mut e2, &mut e3, wa);
bfly(&mut e0, &mut e2, _mm256_loadu_pd(tb.add(2 * k)));
bfly(&mut e1, &mut e3, _mm256_loadu_pd(tb.add(2 * (k + h))));
for (x, v) in q.into_iter().zip([e0, e1, e2, e3]) {
_mm256_storeu_pd(x, v);
}
k += 2;
}
base += 4 * h;
}
}
#[target_feature(enable = "avx2")]
pub(super) unsafe fn scale(buf: &mut [Complex], k: f64) {
let p = buf.as_mut_ptr().cast::<f64>();
let kv = _mm256_set1_pd(k);
let n = buf.len() * 2;
let mut i = 0;
while i + 4 <= n {
_mm256_storeu_pd(p.add(i), _mm256_mul_pd(_mm256_loadu_pd(p.add(i)), kv));
i += 4;
}
while i < n {
*p.add(i) *= k;
i += 1;
}
}
}
#[must_use]
pub fn convolve(a: &[f64], b: &[f64]) -> Vec<f64> {
if a.is_empty() || b.is_empty() {
return Vec::new();
}
let out_len = a.len() + b.len() - 1;
let n = out_len.next_power_of_two();
let load = |s: &[f64]| {
let mut v = vec![Complex::default(); n];
for (d, &x) in v.iter_mut().zip(s) {
d.re = x;
}
v
};
let (mut fa, mut fb) = (load(a), load(b));
fft(&mut fa, false);
fft(&mut fb, false);
for (x, y) in fa.iter_mut().zip(&fb) {
*x *= *y;
}
fft(&mut fa, true);
fa[..out_len].iter().map(|c| c.re).collect()
}
#[cfg(test)]
mod tests {
use super::*;
fn signal(n: usize, seed: f64) -> Vec<Complex> {
(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()
}
fn reference(buf: &mut [Complex], inverse: bool) {
let n = buf.len();
let bits = n.trailing_zeros();
for i in 0..n {
let j = i.reverse_bits() >> (usize::BITS - bits);
if j > i {
buf.swap(i, j);
}
}
let mut len = 2;
while len <= n {
let tw = stage_twiddles(len, inverse);
for block in buf.chunks_exact_mut(len) {
let (lo, hi) = block.split_at_mut(len / 2);
for ((a, b), &t) in lo.iter_mut().zip(hi.iter_mut()).zip(&tw) {
let v = *b * t;
*b = Complex::new(a.re - v.re, a.im - v.im);
*a = Complex::new(a.re + v.re, a.im + v.im);
}
}
len <<= 1;
}
if inverse {
let k = 1.0 / n as f64;
for v in buf.iter_mut() {
v.re *= k;
v.im *= k;
}
}
}
#[test]
fn scalar_and_avx2_equal_the_reference_bit_for_bit() {
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());
for bits in 1..=12 {
let n = 1usize << bits;
for inverse in [false, true] {
let mut want = signal(n, bits as f64);
let mut scalar = want.clone();
let mut avx = want.clone();
reference(&mut want, inverse);
fft_with(&mut scalar, inverse, false);
assert!(same(&want, &scalar), "n {n} inverse {inverse}: scalar differs from the reference");
if slain_rt::cpu::avx2() {
fft_with(&mut avx, inverse, true);
assert!(same(&want, &avx), "n {n} inverse {inverse}: AVX2 differs from the reference");
}
}
}
}
#[test]
fn every_reversal_equals_the_direct_permutation() {
for bits in 1..=17u32 {
let src = signal(1 << bits, f64::from(bits));
let mut want = src.clone();
for (i, &v) in src.iter().enumerate() {
want[rev(i, bits)] = v;
}
let mut got = src.clone();
bit_reverse(&mut got, bits);
assert!(want == got, "bit_reverse bits {bits}");
let mut got = src.clone();
bit_reverse_pairs(&mut got, bits);
assert!(want == got, "pairs bits {bits}");
for b in [3u32, 4, 5, 6] {
for parallel in [false, true] {
let mut got = src.clone();
bit_reverse_tiles(&mut got, bits, b, parallel);
assert!(want == got, "tiles bits {bits} b {b} parallel {parallel}");
}
}
}
}
#[test]
#[ignore]
fn time_fork_join() {
for jobs in [0usize, 1, 2, 3, 16, 64] {
let reps = 20000;
let mut best = f64::MAX;
for _ in 0..11 {
let t = std::time::Instant::now();
for _ in 0..reps {
slain_rt::for_range(crate::bench::slain(jobs), |k| {
crate::bench::slain(k);
});
}
best = best.min(t.elapsed().as_secs_f64());
}
println!("FORKJOIN jobs {jobs:3} {:.9} us per for_range", best / f64::from(reps) * 1e6);
}
}
#[test]
#[ignore]
fn time_tiled_reversal() {
for bits in 10u32..=16 {
let n = 1usize << bits;
let src = signal(n, 1.0);
let reps = ((1 << 24) / n).max(4);
let time = |f: &mut dyn FnMut(&mut [Complex])| {
let mut best = f64::MAX;
for _ in 0..11 {
let mut v = src.clone();
let t = std::time::Instant::now();
for _ in 0..reps {
f(crate::bench::slain(&mut v));
}
best = best.min(t.elapsed().as_secs_f64());
}
best / reps as f64 * 1e6
};
let pairs = time(&mut |v| bit_reverse_pairs(v, bits));
let mut line = format!("TILES n {n:6} pairs {pairs:.9} us");
for b in [3u32, 4, 5, 6] {
for parallel in [false, true] {
let t = time(&mut |v| bit_reverse_tiles(v, bits, b, parallel));
line += &format!(" b{b}{} {t:.9} us {:.9}x", if parallel { "par" } else { "" }, pairs / t);
}
}
println!("{line}");
}
}
#[test]
#[ignore]
fn time_bitrev_share() {
for n in [64usize, 256, 1024, 4096] {
let src = signal(n, 1.0);
let reps = (1 << 22) / n;
let time = |f: &mut dyn FnMut(&mut [Complex])| {
let mut best = f64::MAX;
for _ in 0..15 {
let mut v = src.clone();
let t = std::time::Instant::now();
for _ in 0..reps {
f(crate::bench::slain(&mut v));
}
best = best.min(t.elapsed().as_secs_f64());
}
best / reps as f64
};
let whole = time(&mut |v| fft_with(v, false, slain_rt::cpu::avx2()));
let inline = time(&mut |v| {
let bits = v.len().trailing_zeros();
for i in 0..v.len() {
let j = i.reverse_bits() >> (usize::BITS - bits);
if j > i {
v.swap(i, j);
}
}
});
let cached = time(&mut |v| bit_reverse(v, v.len().trailing_zeros()));
println!(
"BITREV n {n:5} whole {:.9} us inline {:.9} us cached {:.9} us cached share {:.9}% cached/inline {:.9}x",
whole * 1e6,
inline * 1e6,
cached * 1e6,
100.0 * cached / whole,
inline / cached
);
}
}
#[test]
#[ignore]
fn time_fft() {
for n in [64usize, 256, 1024, 4096] {
let src = signal(n, 1.0);
let mut rate = [0f64; 3];
for arm in 0..3 {
let reps = (1 << 22) / n;
let mut best = f64::MAX;
for _ in 0..15 {
let mut v = src.clone();
let t = std::time::Instant::now();
for r in 0..reps {
let v = crate::bench::slain(&mut v);
match arm {
0 => reference(v, r % 2 == 1),
1 => fft_with(v, r % 2 == 1, false),
_ => fft_with(v, r % 2 == 1, slain_rt::cpu::avx2()),
}
}
best = best.min(t.elapsed().as_secs_f64());
}
rate[arm] = reps as f64 / best;
}
println!(
"FFT n {n:5} reference {:.9} scalar+cache {:.9} avx2+cache {:.9} calls/s scalar+cache {:.9}x avx2+cache {:.9}x",
rate[0],
rate[1],
rate[2],
rate[1] / rate[0],
rate[2] / rate[0]
);
}
}
#[test]
fn round_trip_and_convolution() {
let orig: Vec<Complex> = (0..64).map(|i| Complex::new((i as f64 * 0.3).sin(), (i as f64 * 0.7).cos())).collect();
let mut v = orig.clone();
fft(&mut v, false);
let s = orig.iter().fold(Complex::default(), |a, b| Complex::new(a.re + b.re, a.im + b.im));
assert!((v[0].re - s.re).abs() < 1e-9 && (v[0].im - s.im).abs() < 1e-9);
fft(&mut v, true);
assert!(v.iter().zip(&orig).all(|(a, b)| (a.re - b.re).abs() < 1e-12 && (a.im - b.im).abs() < 1e-12));
let (a, b) = ([1.0, 2.0, 3.0], [0.5, -1.0]);
let c = convolve(&a, &b);
let want = [0.5, 0.0, -0.5, -3.0];
assert!(c.iter().zip(want).all(|(g, w)| (g - w).abs() < 1e-12));
}
}