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