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