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    external_product::{
  7        ExternalProduct, ExternalProductInplace, ExternalProductInplaceScratchSpace, ExternalProductScratchSpace,
  8    },
  9    ggsw::GGSWCiphertext,
 10    glwe::GLWECiphertext,
 11    keys::{SecretKey, SecretKeyFourier},
 12};
 13use sampling::source::Source;
 14
 15fn bench_external_product_glwe_fft64(c: &mut Criterion) {
 16    let mut group = c.benchmark_group("external_product_glwe_fft64");
 17
 18    struct Params {
 19        log_n: usize,
 20        basek: usize,
 21        k_rlwe_in: usize,
 22        k_rlwe_out: usize,
 23        k_rgsw: usize,
 24    }
 25
 26    fn runner(p: Params) -> impl FnMut() {
 27        let module: Module<FFT64> = Module::<FFT64>::new(1 << p.log_n);
 28
 29        let basek: usize = p.basek;
 30        let k_rlwe_in: usize = p.k_rlwe_in;
 31        let k_rlwe_out: usize = p.k_rlwe_out;
 32        let k_rgsw: usize = p.k_rgsw;
 33
 34        let rows: usize = (p.k_rlwe_in + p.basek - 1) / p.basek;
 35        let sigma: f64 = 3.2;
 36        let bound: f64 = sigma * 6.0;
 37
 38        let mut ct_rgsw: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k_rgsw, rows);
 39        let mut ct_rlwe_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::new(&module, basek, k_rlwe_in);
 40        let mut ct_rlwe_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::new(&module, basek, k_rlwe_out);
 41        let pt_rgsw: base2k::ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
 42
 43        let mut scratch = ScratchOwned::new(
 44            GGSWCiphertext::encrypt_sk_scratch_space(&module, ct_rgsw.size())
 45                | GLWECiphertext::encrypt_sk_scratch_space(&module, ct_rlwe_in.size())
 46                | GLWECiphertext::external_product_scratch_space(
 47                    &module,
 48                    ct_rlwe_out.size(),
 49                    ct_rlwe_in.size(),
 50                    ct_rgsw.size(),
 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: SecretKey<Vec<u8>> = SecretKey::new(&module);
 59        sk.fill_ternary_prob(0.5, &mut source_xs);
 60        let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module);
 61        sk_dft.dft(&module, &sk);
 62
 63        ct_rgsw.encrypt_sk(
 64            &module,
 65            &pt_rgsw,
 66            &sk_dft,
 67            &mut source_xa,
 68            &mut source_xe,
 69            sigma,
 70            bound,
 71            scratch.borrow(),
 72        );
 73
 74        ct_rlwe_in.encrypt_zero_sk(
 75            &module,
 76            &sk_dft,
 77            &mut source_xa,
 78            &mut source_xe,
 79            sigma,
 80            bound,
 81            scratch.borrow(),
 82        );
 83
 84        move || {
 85            ct_rlwe_out.external_product(
 86                black_box(&module),
 87                black_box(&ct_rlwe_in),
 88                black_box(&ct_rgsw),
 89                black_box(scratch.borrow()),
 90            );
 91        }
 92    }
 93
 94    let params_set: Vec<Params> = vec![Params {
 95        log_n: 10,
 96        basek: 7,
 97        k_rlwe_in: 27,
 98        k_rlwe_out: 27,
 99        k_rgsw: 27,
100    }];
101
102    for params in params_set {
103        let id = BenchmarkId::new("EXTERNAL_PRODUCT_GLWE_FFT64", "");
104        let mut runner = runner(params);
105        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
106    }
107
108    group.finish();
109}
110
111fn bench_external_product_glwe_inplace_fft64(c: &mut Criterion) {
112    let mut group = c.benchmark_group("external_product_glwe_inplace_fft64");
113
114    struct Params {
115        log_n: usize,
116        basek: usize,
117        k_rlwe: usize,
118        k_rgsw: usize,
119    }
120
121    fn runner(p: Params) -> impl FnMut() {
122        let module: Module<FFT64> = Module::<FFT64>::new(1 << p.log_n);
123
124        let basek: usize = p.basek;
125        let k_rlwe: usize = p.k_rlwe;
126        let k_rgsw: usize = p.k_rgsw;
127
128        let rows: usize = (p.k_rlwe + p.basek - 1) / p.basek;
129        let sigma: f64 = 3.2;
130        let bound: f64 = sigma * 6.0;
131
132        let mut ct_rgsw: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k_rgsw, rows);
133        let mut ct_rlwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::new(&module, basek, k_rlwe);
134        let pt_rgsw: base2k::ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
135
136        let mut scratch = ScratchOwned::new(
137            GGSWCiphertext::encrypt_sk_scratch_space(&module, ct_rgsw.size())
138                | GLWECiphertext::encrypt_sk_scratch_space(&module, ct_rlwe.size())
139                | GLWECiphertext::external_product_inplace_scratch_space(&module, ct_rlwe.size(), ct_rgsw.size()),
140        );
141
142        let mut source_xs = Source::new([0u8; 32]);
143        let mut source_xe = Source::new([0u8; 32]);
144        let mut source_xa = Source::new([0u8; 32]);
145
146        let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
147        sk.fill_ternary_prob(0.5, &mut source_xs);
148        let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module);
149        sk_dft.dft(&module, &sk);
150
151        ct_rgsw.encrypt_sk(
152            &module,
153            &pt_rgsw,
154            &sk_dft,
155            &mut source_xa,
156            &mut source_xe,
157            sigma,
158            bound,
159            scratch.borrow(),
160        );
161
162        ct_rlwe.encrypt_zero_sk(
163            &module,
164            &sk_dft,
165            &mut source_xa,
166            &mut source_xe,
167            sigma,
168            bound,
169            scratch.borrow(),
170        );
171
172        move || {
173            let scratch_borrow = scratch.borrow();
174            (0..1374).for_each(|i| {
175                ct_rlwe.external_product_inplace(
176                    black_box(&module),
177                    black_box(&ct_rgsw),
178                    black_box(scratch_borrow),
179                );
180            });
181        }
182    }
183
184    let params_set: Vec<Params> = vec![Params {
185        log_n: 9,
186        basek: 18,
187        k_rlwe: 27,
188        k_rgsw: 27,
189    }];
190
191    for params in params_set {
192        let id = BenchmarkId::new("EXTERNAL_PRODUCT_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_external_product_glwe_fft64,
203    bench_external_product_glwe_inplace_fft64
204);
205criterion_main!(benches);