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