1use poulpy_core::layouts::{
  2    GGLWEAutomorphismKey, GGLWESwitchingKey, GLWECiphertext, GLWESecret, Infos,
  3    prepared::{GGLWEAutomorphismKeyPrepared, GGLWESwitchingKeyPrepared, GLWESecretPrepared, PrepareAlloc},
  4};
  5use std::{hint::black_box, time::Duration};
  6
  7use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
  8use poulpy_backend::cpu_spqlios::FFT64;
  9use poulpy_hal::{
 10    api::{ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow},
 11    layouts::{Module, ScratchOwned},
 12    source::Source,
 13};
 14
 15fn bench_keyswitch_glwe_fft64(c: &mut Criterion) {
 16    let mut group = c.benchmark_group("keyswitch_glwe_fft64");
 17
 18    struct Params {
 19        log_n: usize,
 20        basek: usize,
 21        k_ct_in: usize,
 22        k_ct_out: usize,
 23        k_ksk: usize,
 24        digits: usize,
 25        rank_in: usize,
 26        rank_out: usize,
 27    }
 28
 29    fn runner(p: Params) -> impl FnMut() {
 30        let module: Module<FFT64> = Module::<FFT64>::new(1 << p.log_n);
 31
 32        let n = module.n();
 33        let basek: usize = p.basek;
 34        let k_rlwe_in: usize = p.k_ct_in;
 35        let k_rlwe_out: usize = p.k_ct_out;
 36        let k_grlwe: usize = p.k_ksk;
 37        let rank_in: usize = p.rank_in;
 38        let rank_out: usize = p.rank_out;
 39        let digits: usize = p.digits;
 40
 41        let rows: usize = p.k_ct_in.div_ceil(p.basek * digits);
 42
 43        let mut ksk: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_grlwe, rows, digits, rank_out);
 44        let mut ct_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_rlwe_in, rank_in);
 45        let mut ct_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_rlwe_out, rank_out);
 46
 47        let mut scratch: ScratchOwned<FFT64> = ScratchOwned::alloc(
 48            GGLWESwitchingKey::encrypt_sk_scratch_space(&module, n, basek, ksk.k(), rank_in, rank_out)
 49                | GLWECiphertext::encrypt_sk_scratch_space(&module, n, basek, ct_in.k())
 50                | GLWECiphertext::keyswitch_scratch_space(
 51                    &module,
 52                    n,
 53                    basek,
 54                    ct_out.k(),
 55                    ct_in.k(),
 56                    ksk.k(),
 57                    digits,
 58                    rank_in,
 59                    rank_out,
 60                ),
 61        );
 62
 63        let mut source_xs = Source::new([0u8; 32]);
 64        let mut source_xe = Source::new([0u8; 32]);
 65        let mut source_xa = Source::new([0u8; 32]);
 66
 67        let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_in);
 68        sk_in.fill_ternary_prob(0.5, &mut source_xs);
 69        let sk_in_dft: GLWESecretPrepared<Vec<u8>, FFT64> = sk_in.prepare_alloc(&module, scratch.borrow());
 70
 71        let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
 72        sk_out.fill_ternary_prob(0.5, &mut source_xs);
 73
 74        ksk.encrypt_sk(
 75            &module,
 76            -1,
 77            &sk_in,
 78            &mut source_xa,
 79            &mut source_xe,
 80            scratch.borrow(),
 81        );
 82
 83        ct_in.encrypt_zero_sk(
 84            &module,
 85            &sk_in_dft,
 86            &mut source_xa,
 87            &mut source_xe,
 88            scratch.borrow(),
 89        );
 90
 91        let ksk_prepared: GGLWEAutomorphismKeyPrepared<Vec<u8>, _> = ksk.prepare_alloc(&module, scratch.borrow());
 92
 93        move || {
 94            ct_out.automorphism(&module, &ct_in, &ksk_prepared, scratch.borrow());
 95            black_box(());
 96        }
 97    }
 98
 99    let digits: usize = 1;
100    let basek: usize = 19;
101
102    let params_set: Vec<Params> = vec![Params {
103        log_n: 15,
104        basek,
105        k_ct_in: 874 - digits * basek,
106        k_ct_out: 874 - digits * basek,
107        k_ksk: 874,
108        digits,
109        rank_in: 1,
110        rank_out: 1,
111    }];
112
113    for params in params_set {
114        let id = BenchmarkId::new("KEYSWITCH_GLWE_FFT64", "");
115        let mut runner = runner(params);
116        group.sample_size(500);
117        group.measurement_time(Duration::from_secs(40));
118        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
119    }
120
121    group.finish();
122}
123
124fn bench_keyswitch_glwe_inplace_fft64(c: &mut Criterion) {
125    let mut group = c.benchmark_group("keyswitch_glwe_inplace_fft64");
126
127    struct Params {
128        log_n: usize,
129        basek: usize,
130        k_ct: usize,
131        k_ksk: usize,
132        rank: usize,
133    }
134
135    fn runner(p: Params) -> impl FnMut() {
136        let module: Module<FFT64> = Module::<FFT64>::new(1 << p.log_n);
137
138        let n = module.n();
139        let basek: usize = p.basek;
140        let k_ct: usize = p.k_ct;
141        let k_ksk: usize = p.k_ksk;
142        let rank: usize = p.rank;
143        let digits: usize = 1;
144
145        let rows: usize = p.k_ct.div_ceil(p.basek);
146
147        let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank, rank);
148        let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_ct, rank);
149
150        let mut scratch: ScratchOwned<FFT64> = ScratchOwned::alloc(
151            GGLWESwitchingKey::encrypt_sk_scratch_space(&module, n, basek, ksk.k(), rank, rank)
152                | GLWECiphertext::encrypt_sk_scratch_space(&module, n, basek, ct.k())
153                | GLWECiphertext::keyswitch_inplace_scratch_space(&module, n, basek, ct.k(), ksk.k(), digits, rank),
154        );
155
156        let mut source_xs: Source = Source::new([0u8; 32]);
157        let mut source_xe: Source = Source::new([0u8; 32]);
158        let mut source_xa: Source = Source::new([0u8; 32]);
159
160        let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
161        sk_in.fill_ternary_prob(0.5, &mut source_xs);
162        let sk_in_dft: GLWESecretPrepared<Vec<u8>, FFT64> = sk_in.prepare_alloc(&module, scratch.borrow());
163
164        let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
165        sk_out.fill_ternary_prob(0.5, &mut source_xs);
166
167        ksk.encrypt_sk(
168            &module,
169            &sk_in,
170            &sk_out,
171            &mut source_xa,
172            &mut source_xe,
173            scratch.borrow(),
174        );
175
176        ct.encrypt_zero_sk(
177            &module,
178            &sk_in_dft,
179            &mut source_xa,
180            &mut source_xe,
181            scratch.borrow(),
182        );
183
184        let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, FFT64> = ksk.prepare_alloc(&module, scratch.borrow());
185
186        move || {
187            ct.keyswitch_inplace(&module, &ksk_prepared, scratch.borrow());
188            black_box(());
189        }
190    }
191
192    let params_set: Vec<Params> = vec![Params {
193        log_n: 9,
194        basek: 18,
195        k_ct: 27,
196        k_ksk: 27,
197        rank: 1,
198    }];
199
200    for params in params_set {
201        let id = BenchmarkId::new("KEYSWITCH_GLWE_INPLACE_FFT64", "");
202        let mut runner = runner(params);
203        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
204    }
205
206    group.finish();
207}
208
209criterion_group!(
210    benches,
211    bench_keyswitch_glwe_fft64,
212    bench_keyswitch_glwe_inplace_fft64
213);
214criterion_main!(benches);