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