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