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