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