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};
  9use poulpy_cpu_ref::FFT64Ref;
 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<FFT64Ref> = Module::<FFT64Ref>::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: GLWEAutomorphismKeyLayout = GLWEAutomorphismKeyLayout {
 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: GLWEAutomorphismKey<Vec<u8>> = GLWEAutomorphismKey::alloc_from_infos(&gglwe_atk_layout);
 66        let mut ct_in: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_in_layout);
 67        let mut ct_out: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_out_layout);
 68
 69        let mut scratch: ScratchOwned<FFT64Ref> = 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(&glwe_in_layout);
 85        sk_in.fill_ternary_prob(0.5, &mut source_xs);
 86
 87        let mut sk_in_dft: GLWESecretPrepared<Vec<u8>, FFT64Ref> = 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: GLWEAutomorphismKeyPrepared<Vec<u8>, _> =
108            GLWEAutomorphismKeyPrepared::alloc_from_infos(&module, &ksk);
109        ksk_prepared.prepare(&module, &ksk, scratch.borrow());
110
111        move || {
112            ct_out.automorphism(&module, &ct_in, &ksk_prepared, scratch.borrow());
113            black_box(());
114        }
115    }
116
117    let base2k: usize = 19;
118    let dsize = 1;
119
120    let params_set: Vec<Params> = vec![Params {
121        log_n: 15,
122        base2k: base2k.into(),
123        k_ct_in: (874 - dsize * base2k).into(),
124        k_ct_out: (874 - dsize * base2k).into(),
125        k_ksk: 874_u32.into(),
126        dsize: 1_u32.into(),
127        rank: 1_u32.into(),
128    }];
129
130    for params in params_set {
131        let id = BenchmarkId::new("KEYSWITCH_GLWE_FFT64", "");
132        let mut runner = runner(params);
133        group.sample_size(500);
134        group.measurement_time(Duration::from_secs(40));
135        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
136    }
137
138    group.finish();
139}
140
141fn bench_keyswitch_glwe_inplace_fft64(c: &mut Criterion) {
142    let mut group = c.benchmark_group("keyswitch_glwe_inplace_fft64");
143
144    struct Params {
145        log_n: usize,
146        base2k: Base2K,
147        k_ct: TorusPrecision,
148        k_ksk: TorusPrecision,
149        rank: Rank,
150    }
151
152    fn runner(p: Params) -> impl FnMut() {
153        let module: Module<FFT64Ref> = Module::<FFT64Ref>::new(1 << p.log_n);
154
155        let n: Degree = Degree(module.n() as u32);
156        let base2k: Base2K = p.base2k;
157        let k_ct: TorusPrecision = p.k_ct;
158        let k_ksk: TorusPrecision = p.k_ksk;
159        let rank: Rank = p.rank;
160        let dsize: Dsize = Dsize(1);
161
162        let dnum: Dnum = p.k_ct.div_ceil(p.base2k).into();
163
164        let gglwe_layout: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
165            n,
166            base2k,
167            k: k_ksk,
168            dnum,
169            dsize,
170            rank_in: rank,
171            rank_out: rank,
172        };
173
174        let glwe_layout: GLWELayout = GLWELayout {
175            n,
176            base2k,
177            k: k_ct,
178            rank,
179        };
180
181        let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&gglwe_layout);
182        let mut ct: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_layout);
183
184        let mut scratch: ScratchOwned<FFT64Ref> = ScratchOwned::alloc(
185            GLWESwitchingKey::encrypt_sk_tmp_bytes(&module, &gglwe_layout)
186                | GLWE::encrypt_sk_tmp_bytes(&module, &glwe_layout)
187                | GLWE::keyswitch_tmp_bytes(&module, &glwe_layout, &glwe_layout, &gglwe_layout),
188        );
189
190        let mut source_xs: Source = Source::new([0u8; 32]);
191        let mut source_xe: Source = Source::new([0u8; 32]);
192        let mut source_xa: Source = Source::new([0u8; 32]);
193
194        let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&glwe_layout);
195        sk_in.fill_ternary_prob(0.5, &mut source_xs);
196
197        let mut sk_in_dft: GLWESecretPrepared<Vec<u8>, FFT64Ref> = GLWESecretPrepared::alloc(&module, rank);
198        sk_in_dft.prepare(&module, &sk_in);
199
200        let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&glwe_layout);
201        sk_out.fill_ternary_prob(0.5, &mut source_xs);
202
203        ksk.encrypt_sk(
204            &module,
205            &sk_in,
206            &sk_out,
207            &mut source_xa,
208            &mut source_xe,
209            scratch.borrow(),
210        );
211
212        ct.encrypt_zero_sk(
213            &module,
214            &sk_in_dft,
215            &mut source_xa,
216            &mut source_xe,
217            scratch.borrow(),
218        );
219
220        let mut ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, _> = GLWESwitchingKeyPrepared::alloc_from_infos(&module, &ksk);
221        ksk_prepared.prepare(&module, &ksk, scratch.borrow());
222
223        move || {
224            ct.keyswitch_inplace(&module, &ksk_prepared, scratch.borrow());
225            black_box(());
226        }
227    }
228
229    let params_set: Vec<Params> = vec![Params {
230        log_n: 9,
231        base2k: 18_u32.into(),
232        k_ct: 27_u32.into(),
233        k_ksk: 27_u32.into(),
234        rank: 1_u32.into(),
235    }];
236
237    for params in params_set {
238        let id = BenchmarkId::new("KEYSWITCH_GLWE_INPLACE_FFT64", "");
239        let mut runner = runner(params);
240        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
241    }
242
243    group.finish();
244}
245
246criterion_group!(
247    benches,
248    bench_keyswitch_glwe_fft64,
249    bench_keyswitch_glwe_inplace_fft64
250);
251criterion_main!(benches);