1use poulpy_core::layouts::{
  2    GGSWCiphertext, GLWECiphertext, GLWESecret, Infos,
  3    prepared::{GGSWCiphertextPrepared, GLWESecretPrepared, PrepareAlloc},
  4};
  5use std::hint::black_box;
  6
  7use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
  8
  9use poulpy_backend::cpu_spqlios::FFT64;
 10use poulpy_hal::{
 11    api::{ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow},
 12    layouts::{Module, ScalarZnx, ScratchOwned},
 13    source::Source,
 14};
 15
 16fn bench_external_product_glwe_fft64(c: &mut Criterion) {
 17    let mut group = c.benchmark_group("external_product_glwe_fft64");
 18
 19    struct Params {
 20        log_n: usize,
 21        basek: usize,
 22        k_ct_in: usize,
 23        k_ct_out: usize,
 24        k_ggsw: usize,
 25        rank: usize,
 26    }
 27
 28    fn runner(p: Params) -> impl FnMut() {
 29        let module: Module<FFT64> = Module::<FFT64>::new(1 << p.log_n);
 30
 31        let n: usize = module.n();
 32        let basek: usize = p.basek;
 33        let k_ct_in: usize = p.k_ct_in;
 34        let k_ct_out: usize = p.k_ct_out;
 35        let k_ggsw: usize = p.k_ggsw;
 36        let rank: usize = p.rank;
 37        let digits: usize = 1;
 38
 39        let rows: usize = 1; //(p.k_ct_in.div_ceil(p.basek);
 40
 41        let mut ct_ggsw: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_ggsw, rows, digits, rank);
 42        let mut ct_glwe_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_ct_in, rank);
 43        let mut ct_glwe_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_ct_out, rank);
 44        let pt_rgsw: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
 45
 46        let mut scratch: ScratchOwned<FFT64> = ScratchOwned::alloc(
 47            GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, ct_ggsw.k(), rank)
 48                | GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct_glwe_in.k())
 49                | GLWECiphertext::external_product_scratch_space(
 50                    &module,
 51                    basek,
 52                    ct_glwe_out.k(),
 53                    ct_glwe_in.k(),
 54                    ct_ggsw.k(),
 55                    digits,
 56                    rank,
 57                ),
 58        );
 59
 60        let mut source_xs = Source::new([0u8; 32]);
 61        let mut source_xe = Source::new([0u8; 32]);
 62        let mut source_xa = Source::new([0u8; 32]);
 63
 64        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
 65        sk.fill_ternary_prob(0.5, &mut source_xs);
 66        let sk_dft: GLWESecretPrepared<Vec<u8>, FFT64> = sk.prepare_alloc(&module, scratch.borrow());
 67
 68        ct_ggsw.encrypt_sk(
 69            &module,
 70            &pt_rgsw,
 71            &sk_dft,
 72            &mut source_xa,
 73            &mut source_xe,
 74            scratch.borrow(),
 75        );
 76
 77        ct_glwe_in.encrypt_zero_sk(
 78            &module,
 79            &sk_dft,
 80            &mut source_xa,
 81            &mut source_xe,
 82            scratch.borrow(),
 83        );
 84
 85        let ggsw_prepared: GGSWCiphertextPrepared<Vec<u8>, FFT64> = ct_ggsw.prepare_alloc(&module, scratch.borrow());
 86
 87        move || {
 88            ct_glwe_out.external_product(&module, &ct_glwe_in, &ggsw_prepared, scratch.borrow());
 89            black_box(());
 90        }
 91    }
 92
 93    let params_set: Vec<Params> = vec![Params {
 94        log_n: 11,
 95        basek: 22,
 96        k_ct_in: 44,
 97        k_ct_out: 44,
 98        k_ggsw: 54,
 99        rank: 1,
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_ct: usize,
118        k_ggsw: usize,
119        rank: usize,
120    }
121
122    fn runner(p: Params) -> impl FnMut() {
123        let module: Module<FFT64> = Module::<FFT64>::new(1 << p.log_n);
124
125        let n = module.n();
126        let basek: usize = p.basek;
127        let k_glwe: usize = p.k_ct;
128        let k_ggsw: usize = p.k_ggsw;
129        let rank: usize = p.rank;
130        let digits: usize = 1;
131
132        let rows: usize = p.k_ct.div_ceil(p.basek);
133
134        let mut ct_ggsw: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_ggsw, rows, digits, rank);
135        let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_glwe, rank);
136        let pt_rgsw: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
137
138        let mut scratch: ScratchOwned<FFT64> = ScratchOwned::alloc(
139            GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, ct_ggsw.k(), rank)
140                | GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct_glwe.k())
141                | GLWECiphertext::external_product_inplace_scratch_space(&module, basek, ct_glwe.k(), ct_ggsw.k(), digits, rank),
142        );
143
144        let mut source_xs = Source::new([0u8; 32]);
145        let mut source_xe = Source::new([0u8; 32]);
146        let mut source_xa = Source::new([0u8; 32]);
147
148        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
149        sk.fill_ternary_prob(0.5, &mut source_xs);
150        let sk_dft: GLWESecretPrepared<Vec<u8>, FFT64> = sk.prepare_alloc(&module, scratch.borrow());
151
152        ct_ggsw.encrypt_sk(
153            &module,
154            &pt_rgsw,
155            &sk_dft,
156            &mut source_xa,
157            &mut source_xe,
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            scratch.borrow(),
167        );
168
169        let ggsw_prepared: GGSWCiphertextPrepared<Vec<u8>, FFT64> = ct_ggsw.prepare_alloc(&module, scratch.borrow());
170
171        move || {
172            let scratch_borrow = scratch.borrow();
173            ct_glwe.external_product_inplace(&module, &ggsw_prepared, scratch_borrow);
174            black_box(());
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);