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