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