1use poulpy_core::layouts::{
  2    Base2K, Degree, Dnum, Dsize, GLWE, GLWEAutomorphismKey, GLWEAutomorphismKeyLayout, GLWELayout, GLWESecret, GLWESwitchingKey,
  3    GLWESwitchingKeyLayout, GLWESwitchingKeyPrepared, Rank, TorusPrecision,
  4    prepared::{GLWEAutomorphismKeyPrepared, GLWESecretPrepared},
  5};
  6use std::{hint::black_box, time::Duration};
  7
  8use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
  9
 10#[cfg(all(feature = "enable-avx", target_arch = "x86_64"))]
 11pub use poulpy_cpu_avx::FFT64Avx as BackendImpl;
 12
 13#[cfg(not(all(feature = "enable-avx", target_arch = "x86_64")))]
 14pub use poulpy_cpu_ref::FFT64Ref as BackendImpl;
 15
 16use poulpy_hal::{
 17    api::{ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow},
 18    layouts::{Module, ScratchOwned},
 19    source::Source,
 20};
 21
 22fn bench_keyswitch_glwe_fft64(c: &mut Criterion) {
 23    let mut group = c.benchmark_group("keyswitch_glwe_fft64");
 24
 25    struct Params {
 26        log_n: usize,
 27        base2k: Base2K,
 28        k_ct_in: TorusPrecision,
 29        k_ct_out: TorusPrecision,
 30        k_ksk: TorusPrecision,
 31        dsize: Dsize,
 32        rank: Rank,
 33    }
 34
 35    fn runner(p: Params) -> impl FnMut() {
 36        let module: Module<BackendImpl> = Module::<BackendImpl>::new(1 << p.log_n);
 37
 38        let n: Degree = Degree(module.n() as u32);
 39        let base2k: Base2K = p.base2k;
 40        let k_glwe_in: TorusPrecision = p.k_ct_in;
 41        let k_glwe_out: TorusPrecision = p.k_ct_out;
 42        let k_gglwe: TorusPrecision = p.k_ksk;
 43        let rank: Rank = p.rank;
 44        let dsize: Dsize = p.dsize;
 45
 46        let dnum: Dnum = p.k_ct_in.div_ceil(p.base2k.0 * dsize.0).into();
 47
 48        let gglwe_atk_layout: GLWEAutomorphismKeyLayout = GLWEAutomorphismKeyLayout {
 49            n,
 50            base2k,
 51            k: k_gglwe,
 52            dnum,
 53            rank,
 54            dsize,
 55        };
 56
 57        let glwe_in_layout: GLWELayout = GLWELayout {
 58            n,
 59            base2k,
 60            k: k_glwe_in,
 61            rank,
 62        };
 63
 64        let glwe_out_layout: GLWELayout = GLWELayout {
 65            n,
 66            base2k,
 67            k: k_glwe_out,
 68            rank,
 69        };
 70
 71        let mut ksk: GLWEAutomorphismKey<Vec<u8>> = GLWEAutomorphismKey::alloc_from_infos(&gglwe_atk_layout);
 72        let mut ct_in: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_in_layout);
 73        let mut ct_out: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_out_layout);
 74
 75        let mut scratch: ScratchOwned<BackendImpl> = ScratchOwned::alloc(
 76            GLWESwitchingKey::encrypt_sk_tmp_bytes(&module, &gglwe_atk_layout)
 77                | GLWE::encrypt_sk_tmp_bytes(&module, &glwe_in_layout)
 78                | GLWE::keyswitch_tmp_bytes(
 79                    &module,
 80                    &glwe_out_layout,
 81                    &glwe_in_layout,
 82                    &gglwe_atk_layout,
 83                ),
 84        );
 85
 86        let mut source_xs: Source = Source::new([0u8; 32]);
 87        let mut source_xe: Source = Source::new([0u8; 32]);
 88        let mut source_xa: Source = Source::new([0u8; 32]);
 89
 90        let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&glwe_in_layout);
 91        sk_in.fill_ternary_prob(0.5, &mut source_xs);
 92
 93        let mut sk_in_dft: GLWESecretPrepared<Vec<u8>, BackendImpl> = GLWESecretPrepared::alloc(&module, rank);
 94        sk_in_dft.prepare(&module, &sk_in);
 95
 96        ksk.encrypt_sk(
 97            &module,
 98            -1,
 99            &sk_in,
100            &mut source_xa,
101            &mut source_xe,
102            scratch.borrow(),
103        );
104
105        ct_in.encrypt_zero_sk(
106            &module,
107            &sk_in_dft,
108            &mut source_xa,
109            &mut source_xe,
110            scratch.borrow(),
111        );
112
113        let mut ksk_prepared: GLWEAutomorphismKeyPrepared<Vec<u8>, _> =
114            GLWEAutomorphismKeyPrepared::alloc_from_infos(&module, &ksk);
115        ksk_prepared.prepare(&module, &ksk, scratch.borrow());
116
117        move || {
118            ct_out.automorphism(&module, &ct_in, &ksk_prepared, scratch.borrow());
119            black_box(());
120        }
121    }
122
123    let base2k: usize = 19;
124    let dsize = 1;
125
126    let params_set: Vec<Params> = vec![Params {
127        log_n: 15,
128        base2k: base2k.into(),
129        k_ct_in: (874 - dsize * base2k).into(),
130        k_ct_out: (874 - dsize * base2k).into(),
131        k_ksk: 874_u32.into(),
132        dsize: 1_u32.into(),
133        rank: 1_u32.into(),
134    }];
135
136    for params in params_set {
137        let id = BenchmarkId::new("KEYSWITCH_GLWE_FFT64", "");
138        let mut runner = runner(params);
139        group.sample_size(500);
140        group.measurement_time(Duration::from_secs(40));
141        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
142    }
143
144    group.finish();
145}
146
147fn bench_keyswitch_glwe_inplace_fft64(c: &mut Criterion) {
148    let mut group = c.benchmark_group("keyswitch_glwe_inplace_fft64");
149
150    struct Params {
151        log_n: usize,
152        base2k: Base2K,
153        k_ct: TorusPrecision,
154        k_ksk: TorusPrecision,
155        rank: Rank,
156    }
157
158    fn runner(p: Params) -> impl FnMut() {
159        let module: Module<BackendImpl> = Module::<BackendImpl>::new(1 << p.log_n);
160
161        let n: Degree = Degree(module.n() as u32);
162        let base2k: Base2K = p.base2k;
163        let k_ct: TorusPrecision = p.k_ct;
164        let k_ksk: TorusPrecision = p.k_ksk;
165        let rank: Rank = p.rank;
166        let dsize: Dsize = Dsize(1);
167
168        let dnum: Dnum = p.k_ct.div_ceil(p.base2k).into();
169
170        let gglwe_layout: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
171            n,
172            base2k,
173            k: k_ksk,
174            dnum,
175            dsize,
176            rank_in: rank,
177            rank_out: rank,
178        };
179
180        let glwe_layout: GLWELayout = GLWELayout {
181            n,
182            base2k,
183            k: k_ct,
184            rank,
185        };
186
187        let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&gglwe_layout);
188        let mut ct: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_layout);
189
190        let mut scratch: ScratchOwned<BackendImpl> = ScratchOwned::alloc(
191            GLWESwitchingKey::encrypt_sk_tmp_bytes(&module, &gglwe_layout)
192                | GLWE::encrypt_sk_tmp_bytes(&module, &glwe_layout)
193                | GLWE::keyswitch_tmp_bytes(&module, &glwe_layout, &glwe_layout, &gglwe_layout),
194        );
195
196        let mut source_xs: Source = Source::new([0u8; 32]);
197        let mut source_xe: Source = Source::new([0u8; 32]);
198        let mut source_xa: Source = Source::new([0u8; 32]);
199
200        let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&glwe_layout);
201        sk_in.fill_ternary_prob(0.5, &mut source_xs);
202
203        let mut sk_in_dft: GLWESecretPrepared<Vec<u8>, BackendImpl> = GLWESecretPrepared::alloc(&module, rank);
204        sk_in_dft.prepare(&module, &sk_in);
205
206        let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&glwe_layout);
207        sk_out.fill_ternary_prob(0.5, &mut source_xs);
208
209        ksk.encrypt_sk(
210            &module,
211            &sk_in,
212            &sk_out,
213            &mut source_xa,
214            &mut source_xe,
215            scratch.borrow(),
216        );
217
218        ct.encrypt_zero_sk(
219            &module,
220            &sk_in_dft,
221            &mut source_xa,
222            &mut source_xe,
223            scratch.borrow(),
224        );
225
226        let mut ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, _> = GLWESwitchingKeyPrepared::alloc_from_infos(&module, &ksk);
227        ksk_prepared.prepare(&module, &ksk, scratch.borrow());
228
229        move || {
230            ct.keyswitch_inplace(&module, &ksk_prepared, scratch.borrow());
231            black_box(());
232        }
233    }
234
235    let params_set: Vec<Params> = vec![Params {
236        log_n: 9,
237        base2k: 18_u32.into(),
238        k_ct: 27_u32.into(),
239        k_ksk: 27_u32.into(),
240        rank: 1_u32.into(),
241    }];
242
243    for params in params_set {
244        let id = BenchmarkId::new("KEYSWITCH_GLWE_INPLACE_FFT64", "");
245        let mut runner = runner(params);
246        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
247    }
248
249    group.finish();
250}
251
252criterion_group!(
253    benches,
254    bench_keyswitch_glwe_fft64,
255    bench_keyswitch_glwe_inplace_fft64
256);
257criterion_main!(benches);