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