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