1use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
 2use spqlios::bindings::*;
 3use std::ffi::c_void;
 4
 5fn fft(c: &mut Criterion) {
 6    fn forward<'a>(
 7        m: u32,
 8        log_bound: u32,
 9        reim_fft_precomp: *mut spqlios::reim_fft_precomp,
10        a: &'a [i64],
11    ) -> Box<dyn FnMut() + 'a> {
12        unsafe {
13            let buf_a: *mut f64 = reim_fft_precomp_get_buffer(reim_fft_precomp, 0);
14            reim_from_znx64_simple(m as u32, log_bound as u32, buf_a as *mut c_void, a.as_ptr());
15            Box::new(move || reim_fft(reim_fft_precomp, buf_a))
16        }
17    }
18
19    fn backward<'a>(
20        m: u32,
21        log_bound: u32,
22        reim_ifft_precomp: *mut reim_ifft_precomp,
23        a: &'a [i64],
24    ) -> Box<dyn FnMut() + 'a> {
25        Box::new(move || unsafe {
26            let buf_a: *mut f64 = reim_ifft_precomp_get_buffer(reim_ifft_precomp, 0);
27            reim_from_znx64_simple(m as u32, log_bound as u32, buf_a as *mut c_void, a.as_ptr());
28            reim_ifft(reim_ifft_precomp, buf_a);
29        })
30    }
31
32    let q: u64 = 0x1fffffffffe00001u64;
33
34    let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> =
35        c.benchmark_group("fft");
36
37    for log_n in 10..17 {
38        let n: usize = 1 << log_n;
39        let m: usize = n >> 1;
40        let log_bound: u32 = 19;
41
42        let mut a: Vec<i64> = vec![i64::default(); n];
43        a.iter_mut().enumerate().for_each(|(i, x)| *x = i as i64);
44
45        unsafe {
46            let reim_fft_precomp: *mut reim_fft_precomp = new_reim_fft_precomp(m as u32, 1);
47            let reim_ifft_precomp: *mut reim_ifft_precomp = new_reim_ifft_precomp(m as u32, 1);
48
49            let runners: [(String, Box<dyn FnMut()>); 2] = [
50                (format!("forward"), {
51                    forward(m as u32, log_bound, reim_fft_precomp, &a)
52                }),
53                (format!("backward"), {
54                    backward(m as u32, log_bound, reim_ifft_precomp, &a)
55                }),
56            ];
57
58            for (name, mut runner) in runners {
59                let id: BenchmarkId = BenchmarkId::new(name, format!("n={}", 1 << log_n));
60                b.bench_with_input(id, &(), |b: &mut criterion::Bencher<'_>, _| {
61                    b.iter(&mut runner)
62                });
63            }
64        }
65    }
66}
67
68criterion_group!(benches, fft,);
69criterion_main!(benches);