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