yum/gpu_fft

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

yumNaive DFT now respects workgroup sized144a072026-09-14

use criterion::{BatchSize, Criterion, black_box, criterion_group, criterion_main};
use gpu_fft::cpu::{
    fft_v1_hoist, fft_v2_double_buffer, fft_v3_iterative, fft_v4_radix_4, fft_v5_s0_opt,
    fft_v6_unsafe, fft_v7_radix_8, naive_dft, naive_fft, precompute_twiddles,
};
use num::complex::Complex;
use num::traits::Float;

// 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 get_data<const N: usize>() -> Vec<Complex<f32>> {
    let mut data: Vec<Complex<f32>> = Vec::new();
    data.resize(N, Complex::ZERO);
    // Randomize data.
    for i in 0..data.len() {
        data[i] = Complex::new(rand::random(), rand::random());
    }
    data
}

fn do_dft_benchmark_for_n<const N: usize>(c: &mut Criterion) {
    let data = get_data::<N>();

    let result_ref: Vec<rustfft::num_complex::Complex<f32>> = complex_to_t(&data);

    c.bench_function(format!("size {}, dft", N).as_str(), |b| {
        b.iter_batched_ref(
            || result_ref.clone(),
            |result_ref| {
                naive_dft::<f32>(result_ref);
                black_box(result_ref);
            },
            BatchSize::PerIteration,
        );
    });
}

fn do_naive_fft_benchmark_for_n<const N: usize>(c: &mut Criterion) {
    let data = get_data::<N>();

    let result_ref: Vec<rustfft::num_complex::Complex<f32>> = complex_to_t(&data);

    c.bench_function(format!("size {}, naive fft", N).as_str(), |b| {
        b.iter_batched_ref(
            || result_ref.clone(),
            |result_ref| {
                naive_fft::<f32>(result_ref);
                black_box(result_ref);
            },
            BatchSize::PerIteration,
        );
    });
}

fn do_fft_v1_benchmark_for_n<const N: usize>(c: &mut Criterion) {
    let data = get_data::<N>();

    let twiddles = precompute_twiddles(N);
    let result_ref: Vec<rustfft::num_complex::Complex<f32>> = complex_to_t(&data);

    c.bench_function(format!("size {}, fft v1", N).as_str(), |b| {
        b.iter_batched_ref(
            || result_ref.clone(),
            |result_ref| {
                fft_v1_hoist::<f32>(result_ref, &twiddles);
                black_box(result_ref);
            },
            BatchSize::PerIteration,
        );
    });
}

fn do_fft_v2_benchmark_for_n<const N: usize>(c: &mut Criterion) {
    let data = get_data::<N>();

    let twiddles = precompute_twiddles(N);
    let mut scratch = vec![Complex::new(0.0, 0.0); N];
    let result_ref: Vec<rustfft::num_complex::Complex<f32>> = complex_to_t(&data);

    c.bench_function(format!("size {}, fft v2", N).as_str(), |b| {
        b.iter_batched_ref(
            || result_ref.clone(),
            |result_ref| {
                fft_v2_double_buffer::<f32>(result_ref, &mut scratch, &twiddles);
                black_box(result_ref);
            },
            BatchSize::PerIteration,
        );
    });
}

fn do_fft_v3_benchmark_for_n<const N: usize>(c: &mut Criterion) {
    let data = get_data::<N>();

    let twiddles = precompute_twiddles(N);
    let mut scratch = vec![Complex::new(0.0, 0.0); N];
    let result_ref: Vec<rustfft::num_complex::Complex<f32>> = complex_to_t(&data);

    c.bench_function(format!("size {}, fft v3", N).as_str(), |b| {
        b.iter_batched_ref(
            || result_ref.clone(),
            |result_ref| {
                fft_v3_iterative::<f32>(result_ref, &mut scratch, &twiddles);
                black_box(result_ref);
            },
            BatchSize::PerIteration,
        );
    });
}

fn do_fft_v4_benchmark_for_n<const N: usize>(c: &mut Criterion) {
    let data = get_data::<N>();

    let twiddles = precompute_twiddles(N);
    let mut scratch = vec![Complex::new(0.0, 0.0); N];
    let result_ref: Vec<rustfft::num_complex::Complex<f32>> = complex_to_t(&data);

    c.bench_function(format!("size {}, fft v4", N).as_str(), |b| {
        b.iter_batched_ref(
            || result_ref.clone(),
            |result_ref| {
                fft_v4_radix_4::<f32>(result_ref, &mut scratch, &twiddles);
                black_box(result_ref);
            },
            BatchSize::PerIteration,
        );
    });
}

fn do_fft_v5_benchmark_for_n<const N: usize>(c: &mut Criterion) {
    let data = get_data::<N>();

    let twiddles = precompute_twiddles(N);
    let mut scratch = vec![Complex::new(0.0, 0.0); N];
    let result_ref: Vec<rustfft::num_complex::Complex<f32>> = complex_to_t(&data);

    c.bench_function(format!("size {}, fft v5", N).as_str(), |b| {
        b.iter_batched_ref(
            || result_ref.clone(),
            |result_ref| {
                fft_v5_s0_opt::<f32>(result_ref, &mut scratch, &twiddles);
                black_box(result_ref);
            },
            BatchSize::PerIteration,
        );
    });
}

fn do_fft_v6_benchmark_for_n<const N: usize>(c: &mut Criterion) {
    let data = get_data::<N>();

    let twiddles = precompute_twiddles(N);
    let mut scratch = vec![Complex::new(0.0, 0.0); N];
    let result_ref: Vec<rustfft::num_complex::Complex<f32>> = complex_to_t(&data);

    c.bench_function(format!("size {}, fft v6", N).as_str(), |b| {
        b.iter_batched_ref(
            || result_ref.clone(),
            |result_ref| {
                fft_v6_unsafe::<f32>(result_ref, &mut scratch, &twiddles);
                black_box(result_ref);
            },
            BatchSize::PerIteration,
        );
    });
}

fn do_fft_v7_benchmark_for_n<const N: usize>(c: &mut Criterion) {
    let data = get_data::<N>();

    let twiddles = precompute_twiddles(N);
    let mut scratch = vec![Complex::new(0.0, 0.0); N];
    let result_ref: Vec<rustfft::num_complex::Complex<f32>> = complex_to_t(&data);

    c.bench_function(format!("size {}, fft v7", N).as_str(), |b| {
        b.iter_batched_ref(
            || result_ref.clone(),
            |result_ref| {
                fft_v7_radix_8::<f32>(result_ref, &mut scratch, &twiddles);
                black_box(result_ref);
            },
            BatchSize::PerIteration,
        );
    });
}

fn do_ref_benchmark_for_n<const N: usize>(c: &mut Criterion) {
    let data = get_data::<N>();

    let result_ref: Vec<rustfft::num_complex::Complex<f32>> = complex_to_t(&data);
    let mut ref_planner = rustfft::FftPlanner::<f32>::new();
    let mut scratch: Vec<rustfft::num_complex::Complex<f32>> =
        vec![Complex::<f32>::new(0.0, 0.0); N];
    let ref_fft = ref_planner.plan_fft_forward(result_ref.len());

    c.bench_function(format!("size {}, reference", N).as_str(), |b| {
        b.iter_batched_ref(
            || result_ref.clone(),
            |result_ref| {
                ref_fft.process_with_scratch(result_ref, &mut scratch);
                black_box(result_ref);
            },
            BatchSize::PerIteration,
        );
    });
}

pub fn criterion_benchmark(c: &mut Criterion) {
    do_dft_benchmark_for_n::<4096>(c);

    do_fft_v6_benchmark_for_n::<256>(c);
    do_ref_benchmark_for_n::<256>(c);

    do_fft_v7_benchmark_for_n::<512>(c);
    do_ref_benchmark_for_n::<512>(c);

    do_fft_v6_benchmark_for_n::<1024>(c);
    do_ref_benchmark_for_n::<1024>(c);

    do_dft_benchmark_for_n::<4096>(c);
    do_naive_fft_benchmark_for_n::<4096>(c);
    do_fft_v1_benchmark_for_n::<4096>(c);
    do_fft_v2_benchmark_for_n::<4096>(c);
    do_fft_v3_benchmark_for_n::<4096>(c);
    do_fft_v4_benchmark_for_n::<4096>(c);
    do_fft_v5_benchmark_for_n::<4096>(c);
    do_fft_v6_benchmark_for_n::<4096>(c);
    do_fft_v7_benchmark_for_n::<4096>(c);
    do_ref_benchmark_for_n::<4096>(c);
}

criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);