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