1use base2k::{
  2    Infos, BACKEND, Module, Sampling, SvpPPolOps, VecZnx, VecZnxDft, VecZnxDftOps, VecZnxOps,
  3    VmpPMat, alloc_aligned_u8,
  4};
  5use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
  6use rlwe::{
  7    ciphertext::{Ciphertext, new_gadget_ciphertext},
  8    elem::ElemCommon,
  9    encryptor::{encrypt_grlwe_sk, encrypt_grlwe_sk_tmp_bytes},
 10    gadget_product::{gadget_product_core, gadget_product_tmp_bytes},
 11    keys::SecretKey,
 12    parameters::{Parameters, ParametersLiteral},
 13};
 14use sampling::source::Source;
 15
 16fn bench_gadget_product_inplace(c: &mut Criterion) {
 17    fn runner<'a>(
 18        module: &'a Module,
 19        res_dft_0: &'a mut VecZnxDft,
 20        res_dft_1: &'a mut VecZnxDft,
 21        a: &'a VecZnx,
 22        a_cols: usize,
 23        b: &'a Ciphertext<VmpPMat>,
 24        b_cols: usize,
 25        tmp_bytes: &'a mut [u8],
 26    ) -> Box<dyn FnMut() + 'a> {
 27        Box::new(move || {
 28            gadget_product_core(
 29                module, res_dft_0, res_dft_1, a, a_cols, b, b_cols, tmp_bytes,
 30            );
 31        })
 32    }
 33
 34    let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> =
 35        c.benchmark_group("gadget_product_inplace");
 36
 37    for log_n in 10..11 {
 38        let params_lit: ParametersLiteral = ParametersLiteral {
 39            backend: BACKEND::FFT64,
 40            log_n: log_n,
 41            log_q: 32,
 42            log_p: 0,
 43            log_base2k: 16,
 44            log_scale: 20,
 45            xe: 3.2,
 46            xs: 128,
 47        };
 48
 49        let params: Parameters = Parameters::new(&params_lit);
 50
 51        let mut tmp_bytes: Vec<u8> = alloc_aligned_u8(
 52            params.encrypt_rlwe_sk_tmp_bytes(params.log_q())
 53                | gadget_product_tmp_bytes(
 54                    params.module(),
 55                    params.log_base2k(),
 56                    params.log_q(),
 57                    params.log_q(),
 58                    params.cols_q(),
 59                    params.log_qp(),
 60                )
 61                | encrypt_grlwe_sk_tmp_bytes(
 62                    params.module(),
 63                    params.log_base2k(),
 64                    params.cols_qp(),
 65                    params.log_qp(),
 66                ),
 67        );
 68
 69        let mut source: Source = Source::new([3; 32]);
 70
 71        let mut sk0: SecretKey = SecretKey::new(params.module());
 72        let mut sk1: SecretKey = SecretKey::new(params.module());
 73        sk0.fill_ternary_hw(params.xs(), &mut source);
 74        sk1.fill_ternary_hw(params.xs(), &mut source);
 75
 76        let mut source_xe: Source = Source::new([4; 32]);
 77        let mut source_xa: Source = Source::new([5; 32]);
 78
 79        let mut sk0_svp_ppol: base2k::SvpPPol = params.module().new_svp_ppol();
 80        params.module().svp_prepare(&mut sk0_svp_ppol, &sk0.0);
 81
 82        let mut sk1_svp_ppol: base2k::SvpPPol = params.module().new_svp_ppol();
 83        params.module().svp_prepare(&mut sk1_svp_ppol, &sk1.0);
 84
 85        let mut gadget_ct: Ciphertext<VmpPMat> = new_gadget_ciphertext(
 86            params.module(),
 87            params.log_base2k(),
 88            params.cols_q(),
 89            params.log_qp(),
 90        );
 91
 92        encrypt_grlwe_sk(
 93            params.module(),
 94            &mut gadget_ct,
 95            &sk0.0,
 96            &sk1_svp_ppol,
 97            &mut source_xa,
 98            &mut source_xe,
 99            params.xe(),
100            &mut tmp_bytes,
101        );
102
103        let mut ct: Ciphertext<VecZnx> = params.new_ciphertext(params.log_q());
104
105        params.encrypt_rlwe_sk(
106            &mut ct,
107            None,
108            &sk0_svp_ppol,
109            &mut source_xa,
110            &mut source_xe,
111            &mut tmp_bytes,
112        );
113
114        let mut res_dft_0: VecZnxDft = params.module().new_vec_znx_dft(gadget_ct.cols());
115        let mut res_dft_1: VecZnxDft = params.module().new_vec_znx_dft(gadget_ct.cols());
116
117        let mut a: VecZnx = params.module().new_vec_znx(params.cols_q());
118        params
119            .module()
120            .fill_uniform(params.log_base2k(), &mut a, params.cols_q(), &mut source_xa);
121
122        let a_cols: usize = a.cols();
123        let b_cols: usize = gadget_ct.cols();
124
125        let runners: [(String, Box<dyn FnMut()>); 1] = [(format!("gadget_product"), {
126            runner(
127                params.module(),
128                &mut res_dft_0,
129                &mut res_dft_1,
130                &mut a,
131                a_cols,
132                &gadget_ct,
133                b_cols,
134                &mut tmp_bytes,
135            )
136        })];
137
138        for (name, mut runner) in runners {
139            let id: BenchmarkId = BenchmarkId::new(name, format!("n={}", 1 << log_n));
140            b.bench_with_input(id, &(), |b: &mut criterion::Bencher<'_>, _| {
141                b.iter(&mut runner)
142            });
143        }
144    }
145}
146
147criterion_group!(benches, bench_gadget_product_inplace);
148criterion_main!(benches);