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