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