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