arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/spqlios/examples/fft.rs
spqlios/examples/fft.rsRaw
 1use std::ffi::c_void;
 2use std::time::Instant;
 3
 4use spqlios::bindings::*;
 5
 6fn main() {
 7    let log_bound: usize = 19;
 8
 9    let n: usize = 2048;
10    let m: usize = n >> 1;
11
12    let mut a: Vec<i64> = vec![i64::default(); n];
13    let mut b: Vec<i64> = vec![i64::default(); n];
14    let mut c: Vec<i64> = vec![i64::default(); n];
15
16    a.iter_mut().enumerate().for_each(|(i, x)| *x = i as i64);
17    b[1] = 1;
18
19    println!("{:?}", b);
20
21    unsafe {
22        let reim_fft_precomp = new_reim_fft_precomp(m as u32, 2);
23        let reim_ifft_precomp = new_reim_ifft_precomp(m as u32, 1);
24
25        let buf_a = reim_fft_precomp_get_buffer(reim_fft_precomp, 0);
26        let buf_b = reim_fft_precomp_get_buffer(reim_fft_precomp, 1);
27        let buf_c = reim_ifft_precomp_get_buffer(reim_ifft_precomp, 0);
28
29        let now = Instant::now();
30        (0..1024).for_each(|_| {
31            reim_from_znx64_simple(m as u32, log_bound as u32, buf_a as *mut c_void, a.as_ptr());
32            reim_fft(reim_fft_precomp, buf_a);
33
34            reim_from_znx64_simple(m as u32, log_bound as u32, buf_b as *mut c_void, b.as_ptr());
35            reim_fft(reim_fft_precomp, buf_b);
36
37            reim_fftvec_mul_simple(
38                m as u32,
39                buf_c as *mut c_void,
40                buf_a as *mut c_void,
41                buf_b as *mut c_void,
42            );
43            reim_ifft(reim_ifft_precomp, buf_c);
44
45            reim_to_znx64_simple(
46                m as u32,
47                m as f64,
48                log_bound as u32,
49                c.as_mut_ptr(),
50                buf_c as *mut c_void,
51            )
52        });
53
54        println!("time: {}us", now.elapsed().as_micros());
55        println!("{:?}", &c[..16]);
56    }
57}