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