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