arnaucube
arnaucube/poulpyPublic Self-hosted repository
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, VecZnxVec, VmpPMat,
 3    VmpPMatOps, 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(cols);
27    a.encode_vec_i64(log_base2k, log_k, &a_values, 32);
28    a.normalize(log_base2k, &mut buf);
29
30    a.print(a.cols(), n);
31    println!();
32
33    let mut vecznx: Vec<VecZnx> = Vec::new();
34    (0..rows).for_each(|_| {
35        vecznx.push(module.new_vec_znx(cols));
36    });
37
38    (0..rows).for_each(|i| {
39        vecznx[i].raw_mut()[i * n + 1] = 1 as i64;
40    });
41
42    let slices: Vec<&[i64]> = vecznx.dblptr();
43
44    let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
45    module.vmp_prepare_dblptr(&mut vmp_pmat, &slices, &mut buf);
46
47    let mut c_dft: VecZnxDft = module.new_vec_znx_dft(cols);
48    module.vmp_apply_dft(&mut c_dft, &a, &vmp_pmat, &mut buf);
49
50    let mut c_big: VecZnxBig = c_dft.as_vec_znx_big();
51    module.vec_znx_idft_tmp_a(&mut c_big, &mut c_dft);
52
53    let mut res: VecZnx = module.new_vec_znx(cols);
54    module.vec_znx_big_normalize(log_base2k, &mut res, &c_big, &mut buf);
55
56    let mut values_res: Vec<i64> = vec![i64::default(); n];
57    res.decode_vec_i64(log_base2k, log_k, &mut values_res);
58
59    res.print(res.cols(), n);
60
61    module.free();
62
63    println!("{:?}", values_res)
64}