1use base2k::{FFT64, Module, ScalarZnxAlloc, ScratchOwned};
  2use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main};
  3use rlwe::{
  4    elem::Infos,
  5    encryption::EncryptSkScratchSpace,
  6    glwe::GLWECiphertext,
  7    keys::{SecretKey, SecretKeyFourier},
  8    keyswitch::{KeySwitch, KeySwitchInplace, KeySwitchInplaceScratchSpace, KeySwitchScratchSpace},
  9    keyswitch_key::GLWEKeySwitchKey,
 10};
 11use sampling::source::Source;
 12
 13fn bench_keyswitch_glwe_fft64(c: &mut Criterion) {
 14    let mut group = c.benchmark_group("keyswitch_glwe_fft64");
 15
 16    struct Params {
 17        log_n: usize,
 18        basek: usize,
 19        k_rlwe_in: usize,
 20        k_rlwe_out: usize,
 21        k_grlwe: 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_rlwe_in;
 29        let k_rlwe_out: usize = p.k_rlwe_out;
 30        let k_grlwe: usize = p.k_grlwe;
 31
 32        let rows: usize = (p.k_rlwe_in + p.basek - 1) / p.basek;
 33        let sigma: f64 = 3.2;
 34        let bound: f64 = sigma * 6.0;
 35
 36        let mut ct_grlwe: GLWEKeySwitchKey<Vec<u8>, FFT64> = GLWEKeySwitchKey::new(&module, basek, k_grlwe, rows);
 37        let mut ct_rlwe_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::new(&module, basek, k_rlwe_in);
 38        let mut ct_rlwe_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::new(&module, basek, k_rlwe_out);
 39        let pt_grlwe: base2k::ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
 40
 41        let mut scratch = ScratchOwned::new(
 42            GLWEKeySwitchKey::encrypt_sk_scratch_space(&module, ct_grlwe.size())
 43                | GLWECiphertext::encrypt_sk_scratch_space(&module, ct_rlwe_in.size())
 44                | GLWECiphertext::keyswitch_scratch_space(
 45                    &module,
 46                    ct_rlwe_out.size(),
 47                    ct_rlwe_in.size(),
 48                    ct_grlwe.size(),
 49                ),
 50        );
 51
 52        let mut source_xs = Source::new([0u8; 32]);
 53        let mut source_xe = Source::new([0u8; 32]);
 54        let mut source_xa = Source::new([0u8; 32]);
 55
 56        let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
 57        sk.fill_ternary_prob(0.5, &mut source_xs);
 58        let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module);
 59        sk_dft.dft(&module, &sk);
 60
 61        ct_grlwe.encrypt_sk(
 62            &module,
 63            &pt_grlwe,
 64            &sk_dft,
 65            &mut source_xa,
 66            &mut source_xe,
 67            sigma,
 68            bound,
 69            scratch.borrow(),
 70        );
 71
 72        ct_rlwe_in.encrypt_zero_sk(
 73            &module,
 74            &sk_dft,
 75            &mut source_xa,
 76            &mut source_xe,
 77            sigma,
 78            bound,
 79            scratch.borrow(),
 80        );
 81
 82        move || {
 83            ct_rlwe_out.keyswitch(
 84                black_box(&module),
 85                black_box(&ct_rlwe_in),
 86                black_box(&ct_grlwe),
 87                black_box(scratch.borrow()),
 88            );
 89        }
 90    }
 91
 92    let params_set: Vec<Params> = vec![Params {
 93        log_n: 16,
 94        basek: 50,
 95        k_rlwe_in: 1250,
 96        k_rlwe_out: 1250,
 97        k_grlwe: 1250 + 66,
 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_rlwe: usize,
116        k_grlwe: usize,
117    }
118
119    fn runner(p: Params) -> impl FnMut() {
120        let module: Module<FFT64> = Module::<FFT64>::new(1 << p.log_n);
121
122        let basek: usize = p.basek;
123        let k_rlwe: usize = p.k_rlwe;
124        let k_grlwe: usize = p.k_grlwe;
125
126        let rows: usize = (p.k_rlwe + p.basek - 1) / p.basek;
127        let sigma: f64 = 3.2;
128        let bound: f64 = sigma * 6.0;
129
130        let mut ct_grlwe: GLWEKeySwitchKey<Vec<u8>, FFT64> = GLWEKeySwitchKey::new(&module, basek, k_grlwe, rows);
131        let mut ct_rlwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::new(&module, basek, k_rlwe);
132        let pt_grlwe: base2k::ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
133
134        let mut scratch = ScratchOwned::new(
135            GLWEKeySwitchKey::encrypt_sk_scratch_space(&module, ct_grlwe.size())
136                | GLWECiphertext::encrypt_sk_scratch_space(&module, ct_rlwe.size())
137                | GLWECiphertext::keyswitch_inplace_scratch_space(&module, ct_rlwe.size(), ct_grlwe.size()),
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: SecretKey<Vec<u8>> = SecretKey::new(&module);
145        sk.fill_ternary_prob(0.5, &mut source_xs);
146        let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module);
147        sk_dft.dft(&module, &sk);
148
149        ct_grlwe.encrypt_sk(
150            &module,
151            &pt_grlwe,
152            &sk_dft,
153            &mut source_xa,
154            &mut source_xe,
155            sigma,
156            bound,
157            scratch.borrow(),
158        );
159
160        ct_rlwe.encrypt_zero_sk(
161            &module,
162            &sk_dft,
163            &mut source_xa,
164            &mut source_xe,
165            sigma,
166            bound,
167            scratch.borrow(),
168        );
169
170        move || {
171            ct_rlwe.keyswitch_inplace(
172                black_box(&module),
173                black_box(&ct_grlwe),
174                black_box(scratch.borrow()),
175            );
176        }
177    }
178
179    let params_set: Vec<Params> = vec![Params {
180        log_n: 9,
181        basek: 18,
182        k_rlwe: 27,
183        k_grlwe: 27,
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);