1use backend::{FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScratchOwned};
  2use core::{
  3    elem::Infos,
  4    ggsw_ciphertext::GGSWCiphertext,
  5    glwe_ciphertext::GLWECiphertext,
  6    keys::{SecretKey, SecretKeyFourier},
  7};
  8use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
  9use sampling::source::Source;
 10use std::hint::black_box;
 11
 12fn bench_external_product_glwe_fft64(c: &mut Criterion) {
 13    let mut group = c.benchmark_group("external_product_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_ggsw: usize,
 21        rank: 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_ct_in: usize = p.k_ct_in;
 29        let k_ct_out: usize = p.k_ct_out;
 30        let k_ggsw: usize = p.k_ggsw;
 31        let rank: usize = p.rank;
 32
 33        let rows: usize = (p.k_ct_in + p.basek - 1) / p.basek;
 34        let sigma: f64 = 3.2;
 35
 36        let mut ct_ggsw: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_ggsw, rows, rank);
 37        let mut ct_glwe_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_ct_in, rank);
 38        let mut ct_glwe_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_ct_out, rank);
 39        let pt_rgsw: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
 40
 41        let mut scratch = ScratchOwned::new(
 42            GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, ct_ggsw.k(), rank)
 43                | GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct_glwe_in.k())
 44                | GLWECiphertext::external_product_scratch_space(
 45                    &module,
 46                    basek,
 47                    ct_glwe_out.k(),
 48                    ct_glwe_in.k(),
 49                    ct_ggsw.k(),
 50                    rank,
 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::alloc(&module, rank);
 59        sk.fill_ternary_prob(0.5, &mut source_xs);
 60        let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::alloc(&module, rank);
 61        sk_dft.dft(&module, &sk);
 62
 63        ct_ggsw.encrypt_sk(
 64            &module,
 65            &pt_rgsw,
 66            &sk_dft,
 67            &mut source_xa,
 68            &mut source_xe,
 69            sigma,
 70            scratch.borrow(),
 71        );
 72
 73        ct_glwe_in.encrypt_zero_sk(
 74            &module,
 75            &sk_dft,
 76            &mut source_xa,
 77            &mut source_xe,
 78            sigma,
 79            scratch.borrow(),
 80        );
 81
 82        move || {
 83            ct_glwe_out.external_product(
 84                black_box(&module),
 85                black_box(&ct_glwe_in),
 86                black_box(&ct_ggsw),
 87                black_box(scratch.borrow()),
 88            );
 89        }
 90    }
 91
 92    let params_set: Vec<Params> = vec![Params {
 93        log_n: 10,
 94        basek: 7,
 95        k_ct_in: 27,
 96        k_ct_out: 27,
 97        k_ggsw: 27,
 98        rank: 1,
 99    }];
100
101    for params in params_set {
102        let id = BenchmarkId::new("EXTERNAL_PRODUCT_GLWE_FFT64", "");
103        let mut runner = runner(params);
104        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
105    }
106
107    group.finish();
108}
109
110fn bench_external_product_glwe_inplace_fft64(c: &mut Criterion) {
111    let mut group = c.benchmark_group("external_product_glwe_inplace_fft64");
112
113    struct Params {
114        log_n: usize,
115        basek: usize,
116        k_ct: usize,
117        k_ggsw: usize,
118        rank: 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_glwe: usize = p.k_ct;
126        let k_ggsw: usize = p.k_ggsw;
127        let rank: usize = p.rank;
128
129        let rows: usize = (p.k_ct + p.basek - 1) / p.basek;
130        let sigma: f64 = 3.2;
131
132        let mut ct_ggsw: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_ggsw, rows, rank);
133        let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_glwe, rank);
134        let pt_rgsw: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
135
136        let mut scratch = ScratchOwned::new(
137            GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, ct_ggsw.k(), rank)
138                | GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct_glwe.k())
139                | GLWECiphertext::external_product_inplace_scratch_space(&module, basek, ct_glwe.k(), ct_ggsw.k(), rank),
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::alloc(&module, rank);
147        sk.fill_ternary_prob(0.5, &mut source_xs);
148        let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::alloc(&module, rank);
149        sk_dft.dft(&module, &sk);
150
151        ct_ggsw.encrypt_sk(
152            &module,
153            &pt_rgsw,
154            &sk_dft,
155            &mut source_xa,
156            &mut source_xe,
157            sigma,
158            scratch.borrow(),
159        );
160
161        ct_glwe.encrypt_zero_sk(
162            &module,
163            &sk_dft,
164            &mut source_xa,
165            &mut source_xe,
166            sigma,
167            scratch.borrow(),
168        );
169
170        move || {
171            let scratch_borrow = scratch.borrow();
172            (0..687).for_each(|_| {
173                ct_glwe.external_product_inplace(
174                    black_box(&module),
175                    black_box(&ct_ggsw),
176                    black_box(scratch_borrow),
177                );
178            });
179        }
180    }
181
182    let params_set: Vec<Params> = vec![Params {
183        log_n: 12,
184        basek: 18,
185        k_ct: 54,
186        k_ggsw: 54,
187        rank: 1,
188    }];
189
190    for params in params_set {
191        let id = BenchmarkId::new("EXTERNAL_PRODUCT_GLWE_INPLACE_FFT64", "");
192        let mut runner = runner(params);
193        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
194    }
195
196    group.finish();
197}
198
199criterion_group!(
200    benches,
201    bench_external_product_glwe_fft64,
202    bench_external_product_glwe_inplace_fft64
203);
204criterion_main!(benches);