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