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            ct_glwe_out.external_product(
 79                black_box(&module),
 80                black_box(&ct_glwe_in),
 81                black_box(&ct_ggsw),
 82                black_box(scratch.borrow()),
 83            );
 84        }
 85    }
 86
 87    let params_set: Vec<Params> = vec![Params {
 88        log_n: 10,
 89        basek: 7,
 90        k_ct_in: 27,
 91        k_ct_out: 27,
 92        k_ggsw: 27,
 93        rank: 1,
 94    }];
 95
 96    for params in params_set {
 97        let id = BenchmarkId::new("EXTERNAL_PRODUCT_GLWE_FFT64", "");
 98        let mut runner = runner(params);
 99        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
100    }
101
102    group.finish();
103}
104
105fn bench_external_product_glwe_inplace_fft64(c: &mut Criterion) {
106    let mut group = c.benchmark_group("external_product_glwe_inplace_fft64");
107
108    struct Params {
109        log_n: usize,
110        basek: usize,
111        k_ct: usize,
112        k_ggsw: usize,
113        rank: usize,
114    }
115
116    fn runner(p: Params) -> impl FnMut() {
117        let module: Module<FFT64> = Module::<FFT64>::new(1 << p.log_n);
118
119        let basek: usize = p.basek;
120        let k_glwe: usize = p.k_ct;
121        let k_ggsw: usize = p.k_ggsw;
122        let rank: usize = p.rank;
123        let digits: usize = 1;
124
125        let rows: usize = (p.k_ct + p.basek - 1) / p.basek;
126        let sigma: f64 = 3.2;
127
128        let mut ct_ggsw: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_ggsw, rows, digits, rank);
129        let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_glwe, rank);
130        let pt_rgsw: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
131
132        let mut scratch = ScratchOwned::new(
133            GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, ct_ggsw.k(), rank)
134                | GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct_glwe.k())
135                | GLWECiphertext::external_product_inplace_scratch_space(&module, basek, ct_glwe.k(), ct_ggsw.k(), digits, rank),
136        );
137
138        let mut source_xs = Source::new([0u8; 32]);
139        let mut source_xe = Source::new([0u8; 32]);
140        let mut source_xa = Source::new([0u8; 32]);
141
142        let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
143        sk.fill_ternary_prob(&module, 0.5, &mut source_xs);
144
145        ct_ggsw.encrypt_sk(
146            &module,
147            &pt_rgsw,
148            &sk,
149            &mut source_xa,
150            &mut source_xe,
151            sigma,
152            scratch.borrow(),
153        );
154
155        ct_glwe.encrypt_zero_sk(
156            &module,
157            &sk,
158            &mut source_xa,
159            &mut source_xe,
160            sigma,
161            scratch.borrow(),
162        );
163
164        move || {
165            let scratch_borrow = scratch.borrow();
166            (0..687).for_each(|_| {
167                ct_glwe.external_product_inplace(
168                    black_box(&module),
169                    black_box(&ct_ggsw),
170                    black_box(scratch_borrow),
171                );
172            });
173        }
174    }
175
176    let params_set: Vec<Params> = vec![Params {
177        log_n: 12,
178        basek: 18,
179        k_ct: 54,
180        k_ggsw: 54,
181        rank: 1,
182    }];
183
184    for params in params_set {
185        let id = BenchmarkId::new("EXTERNAL_PRODUCT_GLWE_INPLACE_FFT64", "");
186        let mut runner = runner(params);
187        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
188    }
189
190    group.finish();
191}
192
193criterion_group!(
194    benches,
195    bench_external_product_glwe_fft64,
196    bench_external_product_glwe_inplace_fft64
197);
198criterion_main!(benches);