1use backend::{FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScratchOwned};
  2use core::{GGSWCiphertext, GLWECiphertext, GLWESecret, Infos};
  3use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
  4use sampling::source::Source;
  5use std::hint::black_box;
  6
  7fn bench_external_product_glwe_fft64(c: &mut Criterion) {
  8    let mut group = c.benchmark_group("external_product_glwe_fft64");
  9
 10    struct Params {
 11        log_n: usize,
 12        basek: usize,
 13        k_ct_in: usize,
 14        k_ct_out: usize,
 15        k_ggsw: usize,
 16        rank: usize,
 17    }
 18
 19    fn runner(p: Params) -> impl FnMut() {
 20        let module: Module<FFT64> = Module::<FFT64>::new(1 << p.log_n);
 21
 22        let basek: usize = p.basek;
 23        let k_ct_in: usize = p.k_ct_in;
 24        let k_ct_out: usize = p.k_ct_out;
 25        let k_ggsw: usize = p.k_ggsw;
 26        let rank: usize = p.rank;
 27        let digits: usize = 1;
 28
 29        let rows: usize = (p.k_ct_in + p.basek - 1) / p.basek;
 30        let sigma: f64 = 3.2;
 31
 32        let mut ct_ggsw: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_ggsw, rows, digits, rank);
 33        let mut ct_glwe_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_ct_in, rank);
 34        let mut ct_glwe_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_ct_out, rank);
 35        let pt_rgsw: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
 36
 37        let mut scratch = ScratchOwned::new(
 38            GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, ct_ggsw.k(), rank)
 39                | GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct_glwe_in.k())
 40                | GLWECiphertext::external_product_scratch_space(
 41                    &module,
 42                    basek,
 43                    ct_glwe_out.k(),
 44                    ct_glwe_in.k(),
 45                    ct_ggsw.k(),
 46                    digits,
 47                    rank,
 48                ),
 49        );
 50
 51        let mut source_xs = Source::new([0u8; 32]);
 52        let mut source_xe = Source::new([0u8; 32]);
 53        let mut source_xa = Source::new([0u8; 32]);
 54
 55        let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
 56        sk.fill_ternary_prob(&module, 0.5, &mut source_xs);
 57
 58        ct_ggsw.encrypt_sk(
 59            &module,
 60            &pt_rgsw,
 61            &sk,
 62            &mut source_xa,
 63            &mut source_xe,
 64            sigma,
 65            scratch.borrow(),
 66        );
 67
 68        ct_glwe_in.encrypt_zero_sk(
 69            &module,
 70            &sk,
 71            &mut source_xa,
 72            &mut source_xe,
 73            sigma,
 74            scratch.borrow(),
 75        );
 76
 77        move || {
 78            black_box(ct_glwe_out.external_product(&module, &ct_glwe_in, &ct_ggsw, scratch.borrow()));
 79        }
 80    }
 81
 82    let params_set: Vec<Params> = vec![Params {
 83        log_n: 10,
 84        basek: 7,
 85        k_ct_in: 27,
 86        k_ct_out: 27,
 87        k_ggsw: 27,
 88        rank: 1,
 89    }];
 90
 91    for params in params_set {
 92        let id = BenchmarkId::new("EXTERNAL_PRODUCT_GLWE_FFT64", "");
 93        let mut runner = runner(params);
 94        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
 95    }
 96
 97    group.finish();
 98}
 99
100fn bench_external_product_glwe_inplace_fft64(c: &mut Criterion) {
101    let mut group = c.benchmark_group("external_product_glwe_inplace_fft64");
102
103    struct Params {
104        log_n: usize,
105        basek: usize,
106        k_ct: usize,
107        k_ggsw: usize,
108        rank: usize,
109    }
110
111    fn runner(p: Params) -> impl FnMut() {
112        let module: Module<FFT64> = Module::<FFT64>::new(1 << p.log_n);
113
114        let basek: usize = p.basek;
115        let k_glwe: usize = p.k_ct;
116        let k_ggsw: usize = p.k_ggsw;
117        let rank: usize = p.rank;
118        let digits: usize = 1;
119
120        let rows: usize = (p.k_ct + p.basek - 1) / p.basek;
121        let sigma: f64 = 3.2;
122
123        let mut ct_ggsw: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_ggsw, rows, digits, rank);
124        let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_glwe, rank);
125        let pt_rgsw: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
126
127        let mut scratch = ScratchOwned::new(
128            GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, ct_ggsw.k(), rank)
129                | GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct_glwe.k())
130                | GLWECiphertext::external_product_inplace_scratch_space(&module, basek, ct_glwe.k(), ct_ggsw.k(), digits, rank),
131        );
132
133        let mut source_xs = Source::new([0u8; 32]);
134        let mut source_xe = Source::new([0u8; 32]);
135        let mut source_xa = Source::new([0u8; 32]);
136
137        let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
138        sk.fill_ternary_prob(&module, 0.5, &mut source_xs);
139
140        ct_ggsw.encrypt_sk(
141            &module,
142            &pt_rgsw,
143            &sk,
144            &mut source_xa,
145            &mut source_xe,
146            sigma,
147            scratch.borrow(),
148        );
149
150        ct_glwe.encrypt_zero_sk(
151            &module,
152            &sk,
153            &mut source_xa,
154            &mut source_xe,
155            sigma,
156            scratch.borrow(),
157        );
158
159        move || {
160            let scratch_borrow = scratch.borrow();
161            black_box(ct_glwe.external_product_inplace(&module, &ct_ggsw, scratch_borrow));
162        }
163    }
164
165    let params_set: Vec<Params> = vec![Params {
166        log_n: 12,
167        basek: 18,
168        k_ct: 54,
169        k_ggsw: 54,
170        rank: 1,
171    }];
172
173    for params in params_set {
174        let id = BenchmarkId::new("EXTERNAL_PRODUCT_GLWE_INPLACE_FFT64", "");
175        let mut runner = runner(params);
176        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
177    }
178
179    group.finish();
180}
181
182criterion_group!(
183    benches,
184    bench_external_product_glwe_fft64,
185    bench_external_product_glwe_inplace_fft64
186);
187criterion_main!(benches);