yum/gpu_fft

git clone https://git.yummers.dev/yum/gpu_fft

yumFormatting3c2c6032026-09-10

use num::complex::Complex;
use num::traits::{Float, FloatConst};

fn is_power_of_k(n: usize, k: usize) -> bool {
    match n {
        0 => false,
        1 => true,
        _ => n % k == 0 && is_power_of_k(n / k, k),
    }
}

fn usize_to_float<T: Float>(value: usize) -> T {
    num::cast(value).unwrap()
}

#[rustfmt::skip]
#[allow(dead_code)]
pub fn naive_dft<T: Float + FloatConst>(data: &mut [Complex<T>]) {
    let big_n = data.len();
    let mut result = vec![Complex::new(T::zero(), T::zero()); big_n];
    for k in 0..big_n {
        for n in 0..big_n {
            let k_t     = usize_to_float::<T>(k);
            let n_t     = usize_to_float::<T>(n);
            let big_n_t = usize_to_float::<T>(big_n);
            let phase = -T::TAU() * k_t * n_t / big_n_t;
            let factor = Complex::<T>::cis(phase);
            result[k] = result[k] + data[n] * factor;
        }
    }
    data.copy_from_slice(&result);
}

// Helper to naive_fft. Takes `data` along with 3 numbers that let us recreate an even-odd subset:
//  - `start_idx` tells us where the subset begins;
//  - `big_n` is the number of elements in the subset;
//  - `stride` is the distance between elements.
// We also use a double buffer, `scratch`, to avoid clobbering data while merging results.
#[rustfmt::skip]
fn _naive_fft<T: Float + FloatConst>(data: &mut [Complex<T>], start_idx: usize, big_n: usize, stride: usize, scratch: &mut [Complex<T>]) {
    if big_n == 1 {
        return;
    }
    // Compute DFT of even elements.
    _naive_fft(data, start_idx,        big_n/2, stride*2, scratch);
    // Odd elements.
    _naive_fft(data, start_idx+stride, big_n/2, stride*2, scratch);
    for k in 0..(big_n/2) {
        let p = data[start_idx + 2 * k       * stride];
        let q = data[start_idx + (2 * k + 1) * stride];
        let k_t     = usize_to_float::<T>(k);
        let big_n_t = usize_to_float::<T>(big_n);
        let phase = -T::TAU() * k_t / big_n_t;
        let factor = Complex::<T>::cis(phase);
        scratch[start_idx + k               * stride] = p + q * factor;
        scratch[start_idx + (k + big_n / 2) * stride] = p - q * factor;
    }
    data.copy_from_slice(scratch);
}

// Naive implementation of Cooley-Tukey FFT. Modifies `data`in place. Panics if data.len() is not a power of two.
#[allow(dead_code)]
pub fn naive_fft<T: Float + FloatConst>(data: &mut [Complex<T>]) {
    assert!(is_power_of_k(data.len(), 2));
    let mut scratch = Vec::from(data.as_ref());
    _naive_fft(data, 0, data.len(), 1, &mut scratch);
}

fn _fft_v1_hoist<T: Float + FloatConst>(
    data: &mut [Complex<T>],
    start_idx: usize,
    big_n: usize,
    stride: usize,
    scratch: &mut [Complex<T>],
    twiddles: &[Complex<T>],
) {
    if big_n == 1 {
        return;
    }
    // Compute DFT of even elements.
    _fft_v1_hoist(data, start_idx, big_n / 2, stride * 2, scratch, twiddles);
    // Odd elements.
    _fft_v1_hoist(
        data,
        start_idx + stride,
        big_n / 2,
        stride * 2,
        scratch,
        twiddles,
    );
    for k in 0..(big_n / 2) {
        let p = data[start_idx + 2 * k * stride];
        let q = data[start_idx + (2 * k + 1) * stride];
        let factor = twiddles[k * stride];
        scratch[start_idx + k * stride] = p + q * factor;
        scratch[start_idx + (k + big_n / 2) * stride] = p - q * factor;
    }
    data.copy_from_slice(scratch);
}

// Modification of fft_naive: hoist out and precompute twiddles.
pub fn fft_v1_hoist<T: Float + FloatConst>(data: &mut [Complex<T>], twiddles: &[Complex<T>]) {
    assert!(is_power_of_k(data.len(), 2));
    let mut scratch = Vec::from(data.as_ref());
    _fft_v1_hoist(data, 0, data.len(), 1, &mut scratch, &twiddles);
}

#[inline(always)]
fn _fft_v2_double_buffer<T: Float + FloatConst>(
    src: &mut [Complex<T>],
    dst: &mut [Complex<T>],
    start_idx: usize,
    big_n: usize,
    stride: usize,
    twiddles: &[Complex<T>],
) {
    if big_n == 1 {
        return;
    }
    // Compute DFT of even elements.
    _fft_v2_double_buffer(dst, src, start_idx, big_n / 2, stride * 2, twiddles);
    // Odd elements.
    _fft_v2_double_buffer(
        dst,
        src,
        start_idx + stride,
        big_n / 2,
        stride * 2,
        twiddles,
    );
    for k in 0..(big_n / 2) {
        let p = src[start_idx + 2 * k * stride];
        let q = src[start_idx + (2 * k + 1) * stride];
        let factor = twiddles[k * stride];
        dst[start_idx + k * stride] = p + q * factor;
        dst[start_idx + (k + big_n / 2) * stride] = p - q * factor;
    }
}

pub fn fft_v2_double_buffer<T: Float + FloatConst>(
    src: &mut [Complex<T>],
    dst: &mut [Complex<T>],
    twiddles: &[Complex<T>],
) {
    assert!(is_power_of_k(src.len(), 2));
    dst.copy_from_slice(src);
    // Switching `src` and `dst` means that at the end, the result is in `src` - which is actually
    // what we want! We will be hiding `dst` and `twiddles` in a struct later on :)
    _fft_v2_double_buffer(dst, src, 0, src.len(), 1, twiddles);
}

// Evaluates the base-`k` logarithm of `n`.
// Precondition: `n` is a power of `k`.
const fn log_k_of<const K: usize>(mut n: usize) -> usize {
    let mut res = 0;
    while n > 1 {
        n /= K;
        res += 1;
    }
    res
}

pub fn fft_v3_iterative<T: Float + FloatConst>(
    src: &mut [Complex<T>],
    dst: &mut [Complex<T>],
    twiddles: &[Complex<T>],
) {
    assert!(is_power_of_k(src.len(), 2));
    let n_iter = log_k_of::<2>(src.len());

    dst.copy_from_slice(src);

    let (mut input, mut output) = if n_iter % 2 == 0 {
        (dst, src)
    } else {
        (src, dst)
    };
    let mut stride = input.len();
    let mut big_n = 1;
    for _ in 0..n_iter {
        stride /= 2;
        big_n *= 2;
        std::mem::swap(&mut input, &mut output);

        for start_idx in 0..stride {
            for k in 0..big_n / 2 {
                // Get odd and even elements.
                let p = input[start_idx + 2 * k * stride];
                let q = input[start_idx + (2 * k + 1) * stride];
                // Combine.
                let factor = twiddles[k * stride];
                output[start_idx + k * stride] = p + q * factor;
                output[start_idx + (k + big_n / 2) * stride] = p - q * factor;
            }
        }
    }
}

#[inline(always)]
fn mul_ni<T: Float + FloatConst>(x: Complex<T>) -> Complex<T> {
    Complex::new(x.im, -x.re)
}

fn fft_butterfly_radix_4<T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
    stride: usize,
    big_n: usize,
    twiddles: &[Complex<T>],
) {
    for start_idx in 0..stride {
        for k in 0..big_n / 4 {
            // Collect inputs.
            let i0 = input[start_idx + 4 * k * stride];
            let i1 = input[start_idx + (4 * k + 1) * stride];
            let i2 = input[start_idx + (4 * k + 2) * stride];
            let i3 = input[start_idx + (4 * k + 3) * stride];
            // Collect relevant twiddles.
            let ot1 = twiddles[1 * k * stride];
            let ot2 = twiddles[2 * k * stride];
            let ot3 = twiddles[3 * k * stride];

            let a = i0;
            let b = ot1 * i1;
            let c = ot2 * i2;
            let d = ot3 * i3;

            // To derive this, write the output assignments in terms of
            // a/b/c/d, then factor out!
            let ac_sum = a + c;
            let ac_diff = a - c;
            let bd_sum = b + d;
            let bd_diff_ni = mul_ni(b - d);

            output[start_idx + k * stride] = ac_sum + bd_sum;
            output[start_idx + (k + big_n / 4) * stride] = ac_diff + bd_diff_ni;
            output[start_idx + (k + big_n / 2) * stride] = ac_sum - bd_sum;
            output[start_idx + (k + 3 * big_n / 4) * stride] = ac_diff - bd_diff_ni;
        }
    }
}

fn fft_butterfly_radix_4_s0<T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
) {
    let stride = input.len() / 4;
    let big_n = 4;

    for start_idx in 0..stride {
        for k in 0..big_n / 4 {
            // Collect inputs.
            let i0 = input[start_idx + 4 * k * stride];
            let i1 = input[start_idx + (4 * k + 1) * stride];
            let i2 = input[start_idx + (4 * k + 2) * stride];
            let i3 = input[start_idx + (4 * k + 3) * stride];

            let a = i0;
            let b = i1;
            let c = i2;
            let d = i3;

            // To derive this, write the output assignments in terms of
            // a/b/c/d, then factor out!
            let ac_sum = a + c;
            let ac_diff = a - c;
            let bd_sum = b + d;
            let bd_diff_ni = mul_ni(b - d);

            output[start_idx + k * stride] = ac_sum + bd_sum;
            output[start_idx + (k + big_n / 4) * stride] = ac_diff + bd_diff_ni;
            output[start_idx + (k + big_n / 2) * stride] = ac_sum - bd_sum;
            output[start_idx + (k + 3 * big_n / 4) * stride] = ac_diff - bd_diff_ni;
        }
    }
}

pub fn fft_v4_radix_4<T: Float + FloatConst>(
    src: &mut [Complex<T>],
    dst: &mut [Complex<T>],
    twiddles: &[Complex<T>],
) {
    assert!(is_power_of_k(src.len(), 4));
    let n_iter = log_k_of::<4>(src.len());

    dst.copy_from_slice(src);

    let (mut input, mut output) = if n_iter % 2 == 0 {
        (dst, src)
    } else {
        (src, dst)
    };
    let big_n = input.len();
    let mut stride = big_n;
    let mut big_n = 1;
    for _ in 0..n_iter {
        stride /= 4;
        big_n *= 4;
        std::mem::swap(&mut input, &mut output);

        fft_butterfly_radix_4(input, output, stride, big_n, twiddles);
    }
}

pub fn fft_v5_s0_opt<T: Float + FloatConst>(
    src: &mut [Complex<T>],
    dst: &mut [Complex<T>],
    twiddles: &[Complex<T>],
) {
    assert!(is_power_of_k(src.len(), 4));
    let n_iter = log_k_of::<4>(src.len());

    dst.copy_from_slice(src);

    let (mut input, mut output) = if n_iter % 2 == 0 {
        (dst, src)
    } else {
        (src, dst)
    };
    let big_n = input.len();
    let mut stride = big_n;
    let mut big_n = 1;
    for stage in 0..n_iter {
        stride /= 4;
        big_n *= 4;
        std::mem::swap(&mut input, &mut output);

        if stage == 0 {
            fft_butterfly_radix_4_s0(input, output);
        } else {
            fft_butterfly_radix_4(input, output, stride, big_n, twiddles);
        }
    }
}

fn fft_butterfly_radix_4_unsafe<T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
    stride: usize,
    big_n: usize,
    twiddles: &[Complex<T>],
) {
    let input_ptr = input.as_ptr();
    let output_ptr = output.as_mut_ptr();
    for start_idx in 0..stride {
        for k in 0..big_n / 4 {
            unsafe {
                // Collect inputs.
                let i0 = *input_ptr.add(start_idx + 4 * k * stride);
                let i1 = *input_ptr.add(start_idx + (4 * k + 1) * stride);
                let i2 = *input_ptr.add(start_idx + (4 * k + 2) * stride);
                let i3 = *input_ptr.add(start_idx + (4 * k + 3) * stride);
                // Collect relevant twiddles.
                let ot1 = twiddles.get_unchecked(1 * k * stride);
                let ot2 = twiddles.get_unchecked(2 * k * stride);
                let ot3 = twiddles.get_unchecked(3 * k * stride);

                let a = i0;
                let b = ot1 * i1;
                let c = ot2 * i2;
                let d = ot3 * i3;

                // To derive this, write the output assignments in terms of
                // a/b/c/d, then factor out!
                let ac_sum = a + c;
                let ac_diff = a - c;
                let bd_sum = b + d;
                let bd_diff_ni = mul_ni(b - d);

                *output_ptr.add(start_idx + k * stride) = ac_sum + bd_sum;
                *output_ptr.add(start_idx + (k + big_n / 4) * stride) = ac_diff + bd_diff_ni;
                *output_ptr.add(start_idx + (k + big_n / 2) * stride) = ac_sum - bd_sum;
                *output_ptr.add(start_idx + (k + 3 * big_n / 4) * stride) = ac_diff - bd_diff_ni;
            }
        }
    }
}

fn fft_butterfly_radix_4_s0_unsafe<T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
) {
    let stride = input.len() / 4;
    let big_n = 4;
    let input_ptr = input.as_ptr();
    let output_ptr = output.as_mut_ptr();
    for start_idx in 0..stride {
        for k in 0..big_n / 4 {
            unsafe {
                // Collect inputs.
                let i0 = *input_ptr.add(start_idx + 4 * k * stride);
                let i1 = *input_ptr.add(start_idx + (4 * k + 1) * stride);
                let i2 = *input_ptr.add(start_idx + (4 * k + 2) * stride);
                let i3 = *input_ptr.add(start_idx + (4 * k + 3) * stride);

                let a = i0;
                let b = i1;
                let c = i2;
                let d = i3;

                // To derive this, write the output assignments in terms of
                // a/b/c/d, then factor out!
                let ac_sum = a + c;
                let ac_diff = a - c;
                let bd_sum = b + d;
                let bd_diff_ni = mul_ni(b - d);

                *output_ptr.add(start_idx + k * stride) = ac_sum + bd_sum;
                *output_ptr.add(start_idx + (k + big_n / 4) * stride) = ac_diff + bd_diff_ni;
                *output_ptr.add(start_idx + (k + big_n / 2) * stride) = ac_sum - bd_sum;
                *output_ptr.add(start_idx + (k + 3 * big_n / 4) * stride) = ac_diff - bd_diff_ni;
            }
        }
    }
}

pub fn fft_v6_unsafe<T: Float + FloatConst>(
    src: &mut [Complex<T>],
    dst: &mut [Complex<T>],
    twiddles: &[Complex<T>],
) {
    assert!(is_power_of_k(src.len(), 4));
    assert_eq!(src.len(), dst.len());
    assert_eq!(twiddles.len(), src.len());
    let n_iter = log_k_of::<4>(src.len());

    dst.copy_from_slice(src);

    let (mut input, mut output) = if n_iter % 2 == 0 {
        (dst, src)
    } else {
        (src, dst)
    };
    let big_n = input.len();
    let mut stride = big_n;
    let mut big_n = 1;
    for stage in 0..n_iter {
        stride /= 4;
        big_n *= 4;
        std::mem::swap(&mut input, &mut output);

        if stage == 0 {
            fft_butterfly_radix_4_s0_unsafe(input, output);
        } else {
            fft_butterfly_radix_4_unsafe(input, output, stride, big_n, twiddles);
        }
    }
}

#[inline(always)]
fn rot_45<T: Float + FloatConst>(c: Complex<T>) -> Complex<T> {
    let s = T::FRAC_1_SQRT_2();
    // The standard 2D rotation matrix gives:
    //  [ cos(pi/4) -sin(pi/4)]   [ s -s ]
    //  [ sin(pi/4)  cos(pi/4)] = [ s  s ]
    Complex::<T>::new(c.re - c.im, c.re + c.im) * s
}

#[inline(always)]
fn rot_90<T: Float + FloatConst>(c: Complex<T>) -> Complex<T> {
    // The standard 2D rotation matrix gives:
    //  [ cos(pi/2) -sin(pi/2)]   [ 0 -1 ]
    //  [ sin(pi/2)  cos(pi/2)] = [ 1  0 ]
    Complex::<T>::new(-c.im, c.re)
}

#[inline(always)]
fn rot_180<T: Float + FloatConst>(c: Complex<T>) -> Complex<T> {
    // The standard 2D rotation matrix gives:
    //  [ cos(pi) -sin(pi)]   [ -1  0 ]
    //  [ sin(pi)  cos(pi)] = [  0 -1 ]
    Complex::<T>::new(-c.re, -c.im)
}

#[inline(always)]
fn rot_270<T: Float + FloatConst>(c: Complex<T>) -> Complex<T> {
    // The standard 2D rotation matrix gives:
    //  [ cos(3pi/2) -sin(3pi/2)]   [  0 1 ]
    //  [ sin(3pi/2)  cos(3pi/2)] = [ -1 0 ]
    Complex::<T>::new(c.im, -c.re)
}

fn fft_butterfly_radix_8_unsafe<T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
    stride: usize,
    big_n: usize,
    twiddles: &[Complex<T>],
) {
    let input_ptr = input.as_ptr();
    let output_ptr = output.as_mut_ptr();
    for start_idx in 0..stride {
        for k in 0..big_n / 8 {
            unsafe {
                // Collect inputs.
                let i0 = *input_ptr.add(start_idx + 8 * k * stride);
                let i1 = *input_ptr.add(start_idx + (8 * k + 1) * stride);
                let i2 = *input_ptr.add(start_idx + (8 * k + 2) * stride);
                let i3 = *input_ptr.add(start_idx + (8 * k + 3) * stride);
                let i4 = *input_ptr.add(start_idx + (8 * k + 4) * stride);
                let i5 = *input_ptr.add(start_idx + (8 * k + 5) * stride);
                let i6 = *input_ptr.add(start_idx + (8 * k + 6) * stride);
                let i7 = *input_ptr.add(start_idx + (8 * k + 7) * stride);

                // Collect relevant twiddles.
                let ot1 = twiddles.get_unchecked(1 * k * stride);
                let ot2 = twiddles.get_unchecked(2 * k * stride);
                let ot3 = twiddles.get_unchecked(3 * k * stride);
                let ot4 = twiddles.get_unchecked(4 * k * stride);
                let ot5 = twiddles.get_unchecked(5 * k * stride);
                let ot6 = twiddles.get_unchecked(6 * k * stride);
                let ot7 = twiddles.get_unchecked(7 * k * stride);

                let a = i0;
                let b = ot1 * i1;
                let c = ot2 * i2;
                let d = ot3 * i3;
                let e = ot4 * i4;
                let f = ot5 * i5;
                let g = ot6 * i6;
                let h = ot7 * i7;

                let ae_sum = a + e;
                let ae_diff = a - e;
                let bf_sum = b + f;
                let bf_diff = b - f;
                let cg_sum = c + g;
                let cg_diff = c - g;
                let dh_sum = d + h;
                let dh_diff = d - h;

                let w00 = ae_sum + cg_sum;
                let w01 = ae_sum - cg_sum;
                let w10 = ae_diff + rot_270(cg_diff);
                let w11 = ae_diff - rot_270(cg_diff);
                let x00 = bf_sum + dh_sum;
                let x01 = rot_270(bf_sum) + rot_90(dh_sum);
                let x10 = rot_45(rot_270(bf_diff) + rot_180(dh_diff));
                let x11 = rot_45(rot_180(bf_diff) + rot_270(dh_diff));

                *output_ptr.add(start_idx + k * stride) = w00 + x00;
                *output_ptr.add(start_idx + (k + big_n / 8) * stride) = w10 + x10;
                *output_ptr.add(start_idx + (k + big_n / 4) * stride) = w01 + x01;
                *output_ptr.add(start_idx + (k + 3 * big_n / 8) * stride) = w11 + x11;
                *output_ptr.add(start_idx + (k + big_n / 2) * stride) = w00 - x00;
                *output_ptr.add(start_idx + (k + 5 * big_n / 8) * stride) = w10 - x10;
                *output_ptr.add(start_idx + (k + 3 * big_n / 4) * stride) = w01 - x01;
                *output_ptr.add(start_idx + (k + 7 * big_n / 8) * stride) = w11 - x11;
            }
        }
    }
}

fn fft_butterfly_radix_8_s0_unsafe<T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
) {
    let stride = input.len() / 8;
    let big_n = 8;
    let input_ptr = input.as_ptr();
    let output_ptr = output.as_mut_ptr();
    for start_idx in 0..stride {
        for k in 0..big_n / 8 {
            unsafe {
                // Collect inputs.
                let i0 = *input_ptr.add(start_idx + 8 * k * stride);
                let i1 = *input_ptr.add(start_idx + (8 * k + 1) * stride);
                let i2 = *input_ptr.add(start_idx + (8 * k + 2) * stride);
                let i3 = *input_ptr.add(start_idx + (8 * k + 3) * stride);
                let i4 = *input_ptr.add(start_idx + (8 * k + 4) * stride);
                let i5 = *input_ptr.add(start_idx + (8 * k + 5) * stride);
                let i6 = *input_ptr.add(start_idx + (8 * k + 6) * stride);
                let i7 = *input_ptr.add(start_idx + (8 * k + 7) * stride);

                let a = i0;
                let b = i1;
                let c = i2;
                let d = i3;
                let e = i4;
                let f = i5;
                let g = i6;
                let h = i7;

                let ae_sum = a + e;
                let ae_diff = a - e;
                let bf_sum = b + f;
                let bf_diff = b - f;
                let cg_sum = c + g;
                let cg_diff = c - g;
                let dh_sum = d + h;
                let dh_diff = d - h;

                let w00 = ae_sum + cg_sum;
                let w01 = ae_sum - cg_sum;
                let w10 = ae_diff + rot_270(cg_diff);
                let w11 = ae_diff - rot_270(cg_diff);
                let x00 = bf_sum + dh_sum;
                let x01 = rot_270(bf_sum) + rot_90(dh_sum);
                let x10 = rot_45(rot_270(bf_diff) + rot_180(dh_diff));
                let x11 = rot_45(rot_180(bf_diff) + rot_270(dh_diff));

                *output_ptr.add(start_idx + k * stride) = w00 + x00;
                *output_ptr.add(start_idx + (k + big_n / 8) * stride) = w10 + x10;
                *output_ptr.add(start_idx + (k + big_n / 4) * stride) = w01 + x01;
                *output_ptr.add(start_idx + (k + 3 * big_n / 8) * stride) = w11 + x11;
                *output_ptr.add(start_idx + (k + big_n / 2) * stride) = w00 - x00;
                *output_ptr.add(start_idx + (k + 5 * big_n / 8) * stride) = w10 - x10;
                *output_ptr.add(start_idx + (k + 3 * big_n / 4) * stride) = w01 - x01;
                *output_ptr.add(start_idx + (k + 7 * big_n / 8) * stride) = w11 - x11;
            }
        }
    }
}

pub fn fft_v7_radix_8<T: Float + FloatConst>(
    src: &mut [Complex<T>],
    dst: &mut [Complex<T>],
    twiddles: &[Complex<T>],
) {
    assert!(is_power_of_k(src.len(), 8));
    assert_eq!(src.len(), dst.len());
    assert_eq!(twiddles.len(), src.len());
    let n_iter = log_k_of::<8>(src.len());

    dst.copy_from_slice(src);

    let (mut input, mut output) = if n_iter % 2 == 0 {
        (dst, src)
    } else {
        (src, dst)
    };
    let big_n = input.len();
    let mut stride = big_n;
    let mut big_n = 1;
    for stage in 0..n_iter {
        stride /= 8;
        big_n *= 8;
        std::mem::swap(&mut input, &mut output);

        if stage == 0 {
            fft_butterfly_radix_8_s0_unsafe(input, output);
        } else {
            fft_butterfly_radix_8_unsafe(input, output, stride, big_n, twiddles);
        }
    }
}

// Calculates the "twiddle factors" for an n-element FFT, aka all of the nth roots of unity.
pub fn precompute_twiddles<T: Float + FloatConst>(n: usize) -> Vec<Complex<T>> {
    let mut result = vec![Complex::<T>::new(T::zero(), T::zero()); n];

    let n_f64 = usize_to_float::<f64>(n);
    for i in 0..n {
        let tw_f64 = Complex::<f64>::cis(-f64::TAU() * usize_to_float::<f64>(i) / (n_f64));
        result[i] = Complex::new(T::from(tw_f64.re).unwrap(), T::from(tw_f64.im).unwrap());
    }

    result
}

// Precompute all twiddles factors for a given radix `r` and input size `n`.
pub fn precompute_all_twiddles<T: Float + FloatConst>(r: usize, n: usize) -> Vec<Vec<Complex<T>>> {
    assert!(is_power_of_k(n, r));

    let mut result = Vec::new();
    let mut n_cur = n;
    while n_cur > 1 {
        result.push(precompute_twiddles(n_cur));
        n_cur /= r;
    }
    result
}

#[cfg(test)]
mod tests {
    use rand::RngExt;
    use rand::SeedableRng;
    use rand::rngs::StdRng;

    use super::*;
    use num::complex::Complex;
    use std::time::Instant;

    // Cast one float type to another, truncating if needed.
    fn t0_to_t1<T0: Float, T1: Float>(val: T0) -> T1 {
        return num::cast(val).unwrap();
    }

    // Casts one Complex<T> to another.
    fn complex_to_t<T0: Float, T1: Float>(data: &[Complex<T0>]) -> Vec<Complex<T1>> {
        data.iter()
            .map(|c| Complex::new(t0_to_t1(c.re), t0_to_t1(c.im)))
            .collect()
    }

    fn evaluate_results(
        result_ref_64: &Vec<Complex<f64>>,
        result_cur_64: &Vec<Complex<f64>>,
        duration_dft: std::time::Duration,
        duration_cur: std::time::Duration,
        algo_name: &str,
    ) {
        assert_eq!(result_ref_64.len(), result_cur_64.len());

        println!("  Algorithm:  {}", algo_name);
        println!("    Duration:  {:?}", duration_cur);
        println!(
            "    Speedup:   {:?}",
            duration_dft.as_nanos() as f32 / duration_cur.as_nanos() as f32
        );

        let mut max_err: f64 = 0.0;
        let mut sum_err: f64 = 0.0;
        // TODO: median, 90p, 99p

        let result_64: Vec<Complex<f64>> = complex_to_t(&result_cur_64);

        for i in 0..result_ref_64.len() {
            max_err = max_err.max((result_ref_64[i] - result_64[i]).norm());
            sum_err += (result_ref_64[i] - result_64[i]).norm();
        }
        println!("    Max err:   {}", max_err);
        println!("    Avg err:   {}", sum_err / (result_ref_64.len() as f64));
    }

    #[test]
    fn naive_dft_matches_dft() {
        println!("Testing naive DFT against naive FFT");
        {
            let mut data: Vec<Complex<f64>> = Vec::new();
            data.resize((2 as usize).pow(12), Complex::ZERO);
            // Randomize data.
            let mut rng = StdRng::seed_from_u64(42);
            for i in 0..data.len() {
                data[i] = Complex::new(rng.random(), rng.random());
            }

            let mut result_ref: Vec<rustfft::num_complex::Complex<f64>> = complex_to_t(&data);
            let mut ref_planner = rustfft::FftPlanner::<f64>::new();
            let ref_fft = ref_planner.plan_fft_forward(result_ref.len());
            ref_fft.process(&mut result_ref);
            let result_ref_64: Vec<Complex<f64>> = complex_to_t(&result_ref);

            println!("Input size {}", data.len());

            let mut result_dft: Vec<Complex<f32>> = complex_to_t(&data);
            let mut begin_instant = Instant::now();
            naive_dft(&mut result_dft);
            let duration_dft = begin_instant.elapsed();
            let result_dft_64: Vec<Complex<f64>> = complex_to_t(&result_dft);
            evaluate_results(
                &result_ref_64,
                &result_dft_64,
                duration_dft,
                duration_dft,
                "Naive DFT",
            );

            {
                let mut result_fft: Vec<Complex<f32>> = complex_to_t(&data);
                begin_instant = Instant::now();
                naive_fft(&mut result_fft);
                let duration_fft = begin_instant.elapsed();
                let result_fft_64: Vec<Complex<f64>> = complex_to_t(&result_fft);
                evaluate_results(
                    &result_ref_64,
                    &result_fft_64,
                    duration_dft,
                    duration_fft,
                    "Naive FFT",
                );
            }

            {
                let mut result_fft: Vec<Complex<f32>> = complex_to_t(&data);
                let twiddles = precompute_twiddles(data.len());
                begin_instant = Instant::now();
                fft_v1_hoist(&mut result_fft, &twiddles);
                let duration_fft = begin_instant.elapsed();
                let result_fft_64: Vec<Complex<f64>> = complex_to_t(&result_fft);
                evaluate_results(
                    &result_ref_64,
                    &result_fft_64,
                    duration_dft,
                    duration_fft,
                    "FFT v1 (hoist twiddles)",
                );
            }

            {
                let mut result_fft: Vec<Complex<f32>> = complex_to_t(&data);
                let twiddles = precompute_twiddles(data.len());
                let mut scratch: Vec<Complex<f32>> = vec![Complex::new(0.0, 0.0); data.len()];
                begin_instant = Instant::now();
                fft_v2_double_buffer(&mut result_fft, &mut scratch, &twiddles);
                let duration_fft = begin_instant.elapsed();
                let result_fft_64: Vec<Complex<f64>> = complex_to_t(&result_fft);
                evaluate_results(
                    &result_ref_64,
                    &result_fft_64,
                    duration_dft,
                    duration_fft,
                    "FFT v2 (double buffer)",
                );
            }

            {
                let mut result_fft: Vec<Complex<f32>> = complex_to_t(&data);
                let twiddles = precompute_twiddles(data.len());
                let mut scratch: Vec<Complex<f32>> = vec![Complex::new(0.0, 0.0); data.len()];
                begin_instant = Instant::now();
                fft_v3_iterative(&mut result_fft, &mut scratch, &twiddles);
                let duration_fft = begin_instant.elapsed();
                let result_fft_64: Vec<Complex<f64>> = complex_to_t(&result_fft);
                evaluate_results(
                    &result_ref_64,
                    &result_fft_64,
                    duration_dft,
                    duration_fft,
                    "FFT v3 (iterative)",
                );
            }

            {
                let mut result_fft: Vec<Complex<f32>> = complex_to_t(&data);
                let twiddles = precompute_twiddles(data.len());
                let mut scratch: Vec<Complex<f32>> = vec![Complex::new(0.0, 0.0); data.len()];
                begin_instant = Instant::now();
                fft_v4_radix_4(&mut result_fft, &mut scratch, &twiddles);
                let duration_fft = begin_instant.elapsed();
                let result_fft_64: Vec<Complex<f64>> = complex_to_t(&result_fft);
                evaluate_results(
                    &result_ref_64,
                    &result_fft_64,
                    duration_dft,
                    duration_fft,
                    "FFT v4 (radix 4)",
                );
            }

            {
                let mut result_fft: Vec<Complex<f32>> = complex_to_t(&data);
                let twiddles = precompute_twiddles(data.len());
                let mut scratch: Vec<Complex<f32>> = vec![Complex::new(0.0, 0.0); data.len()];
                begin_instant = Instant::now();
                fft_v5_s0_opt(&mut result_fft, &mut scratch, &twiddles);
                let duration_fft = begin_instant.elapsed();
                let result_fft_64: Vec<Complex<f64>> = complex_to_t(&result_fft);
                evaluate_results(
                    &result_ref_64,
                    &result_fft_64,
                    duration_dft,
                    duration_fft,
                    "FFT v5 (stage 0 opt)",
                );
            }

            {
                let mut result_fft: Vec<Complex<f32>> = complex_to_t(&data);
                let twiddles = precompute_twiddles(data.len());
                let mut scratch: Vec<Complex<f32>> = vec![Complex::new(0.0, 0.0); data.len()];
                begin_instant = Instant::now();
                fft_v6_unsafe(&mut result_fft, &mut scratch, &twiddles);
                let duration_fft = begin_instant.elapsed();
                let result_fft_64: Vec<Complex<f64>> = complex_to_t(&result_fft);
                evaluate_results(
                    &result_ref_64,
                    &result_fft_64,
                    duration_dft,
                    duration_fft,
                    "FFT v6 (unsafe)",
                );
            }

            {
                let mut result_fft: Vec<Complex<f32>> = complex_to_t(&data);
                let twiddles = precompute_twiddles(data.len());
                let mut scratch: Vec<Complex<f32>> = vec![Complex::new(0.0, 0.0); data.len()];
                begin_instant = Instant::now();
                fft_v7_radix_8(&mut result_fft, &mut scratch, &twiddles);
                let duration_fft = begin_instant.elapsed();
                let result_fft_64: Vec<Complex<f64>> = complex_to_t(&result_fft);
                evaluate_results(
                    &result_ref_64,
                    &result_fft_64,
                    duration_dft,
                    duration_fft,
                    "FFT v7 (radix-8)",
                );
            }

            let mut result_rfft: Vec<rustfft::num_complex::Complex<f32>> = complex_to_t(&data);
            let mut rfft_planner = rustfft::FftPlanner::<f32>::new();
            let rfft = rfft_planner.plan_fft_forward(result_ref.len());
            begin_instant = Instant::now();
            rfft.process(&mut result_rfft);
            let duration_rfft = begin_instant.elapsed();
            let result_rfft_64: Vec<Complex<f64>> = complex_to_t(&result_rfft);
            evaluate_results(
                &result_ref_64,
                &result_rfft_64,
                duration_dft,
                duration_rfft,
                "RFFT",
            );
        }
    }
}