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