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