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