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