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