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