1use core::{GGSWCiphertext, GGSWCiphertextExec, GLWECiphertext, GLWESecret, GLWESecretExec, Infos};
  2use std::hint::black_box;
  3
  4use backend::{
  5    hal::{
  6        api::{ModuleNew, ScalarZnxAlloc, ScratchOwnedAlloc, ScratchOwnedBorrow},
  7        layouts::{Module, ScalarZnx, ScratchOwned},
  8    },
  9    implementation::cpu_spqlios::FFT64,
 10};
 11use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
 12use sampling::source::Source;
 13
 14fn bench_external_product_glwe_fft64(c: &mut Criterion) {
 15    let mut group = c.benchmark_group("external_product_glwe_fft64");
 16
 17    struct Params {
 18        log_n: usize,
 19        basek: usize,
 20        k_ct_in: usize,
 21        k_ct_out: usize,
 22        k_ggsw: usize,
 23        rank: 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_ct_in: usize = p.k_ct_in;
 31        let k_ct_out: usize = p.k_ct_out;
 32        let k_ggsw: usize = p.k_ggsw;
 33        let rank: usize = p.rank;
 34        let digits: usize = 1;
 35
 36        let rows: usize = 1; //(p.k_ct_in.div_ceil(p.basek);
 37        let sigma: f64 = 3.2;
 38
 39        let mut ct_ggsw: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(&module, basek, k_ggsw, rows, digits, rank);
 40        let mut ct_glwe_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_ct_in, rank);
 41        let mut ct_glwe_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_ct_out, rank);
 42        let pt_rgsw: ScalarZnx<Vec<u8>> = module.scalar_znx_alloc(1);
 43
 44        let mut scratch: ScratchOwned<FFT64> = ScratchOwned::alloc(
 45            GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, ct_ggsw.k(), rank)
 46                | GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct_glwe_in.k())
 47                | GLWECiphertext::external_product_scratch_space(
 48                    &module,
 49                    basek,
 50                    ct_glwe_out.k(),
 51                    ct_glwe_in.k(),
 52                    ct_ggsw.k(),
 53                    digits,
 54                    rank,
 55                ),
 56        );
 57
 58        let mut source_xs = Source::new([0u8; 32]);
 59        let mut source_xe = Source::new([0u8; 32]);
 60        let mut source_xa = Source::new([0u8; 32]);
 61
 62        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
 63        sk.fill_ternary_prob(0.5, &mut source_xs);
 64        let sk_dft: GLWESecretExec<Vec<u8>, FFT64> = GLWESecretExec::from(&module, &sk);
 65
 66        ct_ggsw.encrypt_sk(
 67            &module,
 68            &pt_rgsw,
 69            &sk_dft,
 70            &mut source_xa,
 71            &mut source_xe,
 72            sigma,
 73            scratch.borrow(),
 74        );
 75
 76        ct_glwe_in.encrypt_zero_sk(
 77            &module,
 78            &sk_dft,
 79            &mut source_xa,
 80            &mut source_xe,
 81            sigma,
 82            scratch.borrow(),
 83        );
 84
 85        let ggsw_exec: GGSWCiphertextExec<Vec<u8>, FFT64> = GGSWCiphertextExec::from(&module, &ct_ggsw, scratch.borrow());
 86
 87        move || {
 88            black_box(ct_glwe_out.external_product(&module, &ct_glwe_in, &ggsw_exec, scratch.borrow()));
 89        }
 90    }
 91
 92    let params_set: Vec<Params> = vec![Params {
 93        log_n: 11,
 94        basek: 22,
 95        k_ct_in: 44,
 96        k_ct_out: 44,
 97        k_ggsw: 54,
 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        let digits: usize = 1;
129
130        let rows: usize = p.k_ct.div_ceil(p.basek);
131        let sigma: f64 = 3.2;
132
133        let mut ct_ggsw: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(&module, basek, k_ggsw, rows, digits, rank);
134        let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_glwe, rank);
135        let pt_rgsw: ScalarZnx<Vec<u8>> = module.scalar_znx_alloc(1);
136
137        let mut scratch: ScratchOwned<FFT64> = ScratchOwned::alloc(
138            GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, ct_ggsw.k(), rank)
139                | GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct_glwe.k())
140                | GLWECiphertext::external_product_inplace_scratch_space(&module, basek, ct_glwe.k(), ct_ggsw.k(), digits, rank),
141        );
142
143        let mut source_xs = Source::new([0u8; 32]);
144        let mut source_xe = Source::new([0u8; 32]);
145        let mut source_xa = Source::new([0u8; 32]);
146
147        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
148        sk.fill_ternary_prob(0.5, &mut source_xs);
149        let sk_dft: GLWESecretExec<Vec<u8>, FFT64> = GLWESecretExec::from(&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        let ggsw_exec: GGSWCiphertextExec<Vec<u8>, FFT64> = GGSWCiphertextExec::from(&module, &ct_ggsw, scratch.borrow());
171
172        move || {
173            let scratch_borrow = scratch.borrow();
174            black_box(ct_glwe.external_product_inplace(&module, &ggsw_exec, scratch_borrow));
175        }
176    }
177
178    let params_set: Vec<Params> = vec![Params {
179        log_n: 12,
180        basek: 18,
181        k_ct: 54,
182        k_ggsw: 54,
183        rank: 1,
184    }];
185
186    for params in params_set {
187        let id = BenchmarkId::new("EXTERNAL_PRODUCT_GLWE_INPLACE_FFT64", "");
188        let mut runner = runner(params);
189        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
190    }
191
192    group.finish();
193}
194
195criterion_group!(
196    benches,
197    bench_external_product_glwe_fft64,
198    bench_external_product_glwe_inplace_fft64
199);
200criterion_main!(benches);