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}