poulpy/base2k/examples/vector_matrix_product.rs
base2k/examples/vector_matrix_product.rsRaw
1use base2k::{
2 BACKEND, Encoding, Infos, Module, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps, VecZnxOps, VmpPMat, VmpPMatOps,
3 alloc_aligned,
4};
5
6fn main() {
7 let log_n: i32 = 5;
8 let n: usize = 1 << log_n;
9
10 let module: Module = Module::new(n, BACKEND::FFT64);
11 let log_base2k: usize = 15;
12 let cols: usize = 5;
13 let log_k: usize = log_base2k * cols - 5;
14
15 let rows: usize = cols;
16 let cols: usize = cols + 1;
17
18 // Maximum size of the byte scratch needed
19 let tmp_bytes: usize = module.vmp_prepare_tmp_bytes(rows, cols) | module.vmp_apply_dft_tmp_bytes(cols, cols, rows, cols);
20
21 let mut buf: Vec<u8> = alloc_aligned(tmp_bytes);
22
23 let mut a_values: Vec<i64> = vec![i64::default(); n];
24 a_values[1] = (1 << log_base2k) + 1;
25
26 let mut a: VecZnx = module.new_vec_znx(1, rows);
27 a.encode_vec_i64(0, log_base2k, log_k, &a_values, 32);
28 a.normalize(log_base2k, &mut buf);
29
30 a.print(0, a.cols(), n);
31 println!();
32
33 let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
34
35 (0..a.cols()).for_each(|row_i| {
36 let mut tmp: VecZnx = module.new_vec_znx(1, cols);
37 tmp.at_mut(row_i)[1] = 1 as i64;
38 module.vmp_prepare_row(&mut vmp_pmat, tmp.raw(), row_i, &mut buf);
39 });
40
41 let mut c_dft: VecZnxDft = module.new_vec_znx_dft(cols);
42 module.vmp_apply_dft(&mut c_dft, &a, &vmp_pmat, &mut buf);
43
44 let mut c_big: VecZnxBig = c_dft.as_vec_znx_big();
45 module.vec_znx_idft_tmp_a(&mut c_big, &mut c_dft);
46
47 let mut res: VecZnx = module.new_vec_znx(1, rows);
48 module.vec_znx_big_normalize(log_base2k, &mut res, &c_big, &mut buf);
49
50 let mut values_res: Vec<i64> = vec![i64::default(); n];
51 res.decode_vec_i64(0, log_base2k, log_k, &mut values_res);
52
53 res.print(0, res.cols(), n);
54
55 module.free();
56
57 println!("{:?}", values_res)
58}