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