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