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