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