1use base2k::{FFT64, Module, SvpPPolOps, VecZnx, VmpPMat, alloc_aligned_u8};
  2use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
  3use rlwe::{
  4    ciphertext::{Ciphertext, new_gadget_ciphertext},
  5    elem::Elem,
  6    encryptor::{encrypt_grlwe_sk_thread_safe, encrypt_grlwe_sk_tmp_bytes},
  7    evaluator::{gadget_product_inplace, gadget_product_tmp_bytes},
  8    key_generator::gen_switching_key_thread_safe_tmp_bytes,
  9    keys::SecretKey,
 10    parameters::{Parameters, ParametersLiteral},
 11};
 12use sampling::source::Source;
 13
 14fn bench_gadget_product_inplace(c: &mut Criterion) {
 15    fn runner<'a>(
 16        module: &'a Module,
 17        elem: &'a mut Elem<VecZnx>,
 18        gadget_ct: &'a Ciphertext<VmpPMat>,
 19        tmp_bytes: &'a mut [u8],
 20    ) -> Box<dyn FnMut() + 'a> {
 21        Box::new(move || gadget_product_inplace::<true, _>(module, elem, gadget_ct, tmp_bytes))
 22    }
 23
 24    let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> =
 25        c.benchmark_group("gadget_product_inplace");
 26
 27    for log_n in 10..11 {
 28        let params_lit: ParametersLiteral = ParametersLiteral {
 29            log_n: log_n,
 30            log_q: 32,
 31            log_p: 0,
 32            log_base2k: 16,
 33            log_scale: 20,
 34            xe: 3.2,
 35            xs: 128,
 36        };
 37
 38        let params: Parameters = Parameters::new::<FFT64>(&params_lit);
 39
 40        let mut tmp_bytes: Vec<u8> = alloc_aligned_u8(
 41            params.decrypt_rlwe_thread_safe_tmp_byte(params.log_q())
 42                | params.encrypt_rlwe_sk_tmp_bytes(params.log_q())
 43                | gen_switching_key_thread_safe_tmp_bytes(
 44                    params.module(),
 45                    params.log_base2k(),
 46                    params.limbs_q(),
 47                    params.log_q(),
 48                )
 49                | gadget_product_tmp_bytes(
 50                    params.module(),
 51                    params.log_base2k(),
 52                    params.log_q(),
 53                    params.log_q(),
 54                    params.limbs_q(),
 55                    params.log_qp(),
 56                )
 57                | encrypt_grlwe_sk_tmp_bytes(
 58                    params.module(),
 59                    params.log_base2k(),
 60                    params.limbs_qp(),
 61                    params.log_qp(),
 62                ),
 63            64,
 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.limbs_q(),
 86            params.log_qp(),
 87        );
 88
 89        encrypt_grlwe_sk_thread_safe(
 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_thread_safe(
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 runners: [(String, Box<dyn FnMut()>); 1] = [(format!("gadget_product"), {
112            runner(params.module(), &mut ct.0, &gadget_ct, &mut tmp_bytes)
113        })];
114
115        for (name, mut runner) in runners {
116            let id: BenchmarkId = BenchmarkId::new(name, format!("n={}", 1 << log_n));
117            b.bench_with_input(id, &(), |b: &mut criterion::Bencher<'_>, _| {
118                b.iter(&mut runner)
119            });
120        }
121    }
122}
123
124criterion_group!(benches, bench_gadget_product_inplace);
125criterion_main!(benches);