poulpy/base2k/examples/fft.rs
base2k/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}