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