1use std::hint::black_box;
  2
  3use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
  4use poulpy_core::{
  5    GGSWNoise, GLWEDecrypt, GLWEEncryptSk, GLWEExternalProduct, LWEEncryptSk, ScratchTakeCore,
  6    layouts::{
  7        Base2K, Degree, Dnum, Dsize, GGLWEToGGSWKeyLayout, GGSWLayout, GGSWPreparedFactory, GLWEAutomorphismKeyLayout,
  8        GLWELayout, GLWESecret, GLWESecretPrepared, GLWESecretPreparedFactory, GLWESwitchingKeyLayout, GLWEToLWEKeyLayout,
  9        LWESecret, Rank, TorusPrecision,
 10    },
 11};
 12
 13#[cfg(all(feature = "enable-avx", target_arch = "x86_64"))]
 14pub use poulpy_cpu_avx::FFT64Avx as BackendImpl;
 15
 16#[cfg(not(all(feature = "enable-avx", target_arch = "x86_64")))]
 17pub use poulpy_cpu_ref::FFT64Ref as BackendImpl;
 18
 19use poulpy_hal::{
 20    api::{ModuleN, ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxRotateInplace},
 21    layouts::{Backend, Module, Scratch, ScratchOwned},
 22    source::Source,
 23};
 24use poulpy_schemes::bin_fhe::{
 25    bdd_arithmetic::{
 26        Add, And, BDDKey, BDDKeyEncryptSk, BDDKeyLayout, BDDKeyPrepared, BDDKeyPreparedFactory, ExecuteBDDCircuit2WTo1W, FheUint,
 27        FheUintPrepare, FheUintPrepared, Or, Sll, Slt, Sltu, Sra, Srl, Sub, Xor,
 28    },
 29    blind_rotation::{
 30        BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory, BlindRotationKeyInfos, BlindRotationKeyLayout, CGGI,
 31    },
 32    circuit_bootstrapping::{
 33        CircuitBootstrappingKey, CircuitBootstrappingKeyEncryptSk, CircuitBootstrappingKeyLayout,
 34        CircuitBootstrappingKeyPrepared, CircuitBootstrappingKeyPreparedFactory, CirtuitBootstrappingExecute,
 35    },
 36};
 37
 38// Common setup data structure
 39struct BenchmarkSetup<BE: Backend, BRA: BlindRotationAlgo> {
 40    module: Module<BE>,
 41    scratch: ScratchOwned<BE>,
 42    a_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE>,
 43    b_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE>,
 44    bdd_key_prepared: BDDKeyPrepared<Vec<u8>, BRA, BE>,
 45    glwe_layout: GLWELayout,
 46}
 47
 48struct Params {
 49    name: String,
 50    block_size: usize,
 51    glwe_layout: GLWELayout,
 52    ggsw_layout: GGSWLayout,
 53    bdd_layout: BDDKeyLayout,
 54}
 55
 56fn setup_benchmark<BE: Backend, BRA: BlindRotationAlgo>(params: &Params) -> BenchmarkSetup<BE, BRA>
 57where
 58    Module<BE>: ModuleNew<BE>
 59        + ModuleN
 60        + GLWESecretPreparedFactory<BE>
 61        + GLWEExternalProduct<BE>
 62        + GLWEDecrypt<BE>
 63        + LWEEncryptSk<BE>
 64        + CircuitBootstrappingKeyEncryptSk<BRA, BE>
 65        + CircuitBootstrappingKeyPreparedFactory<BRA, BE>
 66        + CirtuitBootstrappingExecute<BRA, BE>
 67        + GGSWPreparedFactory<BE>
 68        + GGSWNoise<BE>
 69        + GLWEEncryptSk<BE>
 70        + VecZnxRotateInplace<BE>
 71        + BDDKeyEncryptSk<BRA, BE>
 72        + BDDKeyPreparedFactory<BRA, BE>
 73        + FheUintPrepare<BRA, BE>
 74        + ExecuteBDDCircuit2WTo1W<BE>,
 75    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
 76    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 77    Scratch<BE>: ScratchTakeCore<BE>,
 78{
 79    // Scratch space (16MB)
 80    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 24);
 81
 82    let n_glwe: poulpy_core::layouts::Degree = params.bdd_layout.cbt_layout.brk_layout.n_glwe();
 83    let n_lwe: poulpy_core::layouts::Degree = params.bdd_layout.cbt_layout.brk_layout.n_lwe();
 84    let rank: poulpy_core::layouts::Rank = params.bdd_layout.cbt_layout.brk_layout.rank;
 85
 86    let module: Module<BE> = Module::<BE>::new(n_glwe.as_u32() as u64);
 87
 88    let mut source_xs: Source = Source::new([1u8; 32]);
 89    let mut source_xa: Source = Source::new([1u8; 32]);
 90    let mut source_xe: Source = Source::new([1u8; 32]);
 91
 92    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
 93    sk_lwe.fill_binary_block(params.block_size, &mut source_xs);
 94
 95    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n_glwe, rank);
 96    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
 97
 98    // Circuit bootstrapping evaluation key
 99    let mut cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> =
100        CircuitBootstrappingKey::alloc_from_infos(&params.bdd_layout.cbt_layout);
101    cbt_key.encrypt_sk(
102        &module,
103        &sk_lwe,
104        &sk_glwe,
105        &mut source_xa,
106        &mut source_xe,
107        scratch.borrow(),
108    );
109
110    let mut cbt_key_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, BE> =
111        CircuitBootstrappingKeyPrepared::alloc_from_infos(&module, &params.bdd_layout.cbt_layout);
112    cbt_key_prepared.prepare(&module, &cbt_key, scratch.borrow());
113
114    let mut sk_glwe_prepared = GLWESecretPrepared::alloc_from_infos(&module, &params.glwe_layout);
115    sk_glwe_prepared.prepare(&module, &sk_glwe);
116
117    let mut bdd_key: BDDKey<Vec<u8>, BRA> = BDDKey::alloc_from_infos(&params.bdd_layout);
118    bdd_key.encrypt_sk(
119        &module,
120        &sk_lwe,
121        &sk_glwe,
122        &mut source_xa,
123        &mut source_xe,
124        scratch.borrow(),
125    );
126
127    let input_a = 255_u32;
128    let input_b = 30_u32;
129
130    let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&params.glwe_layout);
131    a_enc.encrypt_sk(
132        &module,
133        input_a,
134        &sk_glwe_prepared,
135        &mut source_xa,
136        &mut source_xe,
137        scratch.borrow(),
138    );
139
140    let mut b_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&params.glwe_layout);
141    b_enc.encrypt_sk(
142        &module,
143        input_b,
144        &sk_glwe_prepared,
145        &mut source_xa,
146        &mut source_xe,
147        scratch.borrow(),
148    );
149
150    //////// Homomorphic computation starts here ////////
151
152    // Preparing the BDD Key
153    // The BDD key must be prepared once before any operation is performed
154    let mut bdd_key_prepared: BDDKeyPrepared<Vec<u8>, BRA, BE> = BDDKeyPrepared::alloc_from_infos(&module, &params.bdd_layout);
155    bdd_key_prepared.prepare(&module, &bdd_key, scratch.borrow());
156
157    // Input Preparation
158    // Before each operation, the inputs to that operation must be prepared
159    // Preparation extracts each bit of the integer into a seperate GLWE ciphertext and bootstraps it into a GGSW ciphertext
160    let mut a_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE> = FheUintPrepared::alloc_from_infos(&module, &params.ggsw_layout);
161    a_enc_prepared.prepare(&module, &a_enc, &bdd_key_prepared, scratch.borrow());
162
163    let mut b_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE> = FheUintPrepared::alloc_from_infos(&module, &params.ggsw_layout);
164    b_enc_prepared.prepare(&module, &b_enc, &bdd_key_prepared, scratch.borrow());
165
166    BenchmarkSetup {
167        module,
168        scratch,
169        a_enc_prepared,
170        b_enc_prepared,
171        bdd_key_prepared,
172        glwe_layout: params.glwe_layout.clone(),
173    }
174}
175
176fn create_runner<BE: Backend, BRA: BlindRotationAlgo, F>(setup: BenchmarkSetup<BE, BRA>, operation: F) -> impl FnMut()
177where
178    Module<BE>: ExecuteBDDCircuit2WTo1W<BE>,
179    Scratch<BE>: ScratchTakeCore<BE>,
180    ScratchOwned<BE>: ScratchOwnedBorrow<BE>,
181    F: Fn(
182        &mut FheUint<Vec<u8>, u32>,
183        &Module<BE>,
184        &FheUintPrepared<Vec<u8>, u32, BE>,
185        &FheUintPrepared<Vec<u8>, u32, BE>,
186        &BDDKeyPrepared<Vec<u8>, BRA, BE>,
187        &mut Scratch<BE>,
188    ),
189{
190    let BenchmarkSetup {
191        module,
192        mut scratch,
193        a_enc_prepared,
194        b_enc_prepared,
195        bdd_key_prepared,
196        glwe_layout,
197    } = setup;
198
199    let mut c_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_layout);
200
201    move || {
202        operation(
203            &mut c_enc,
204            &module,
205            &a_enc_prepared,
206            &b_enc_prepared,
207            &bdd_key_prepared,
208            scratch.borrow(),
209        );
210        black_box(());
211    }
212}
213
214fn bench_operation<BE: Backend, BRA: BlindRotationAlgo, F>(
215    group: &mut criterion::BenchmarkGroup<'_, criterion::measurement::WallTime>,
216    params: &Params,
217    operation_name: &str,
218    operation: F,
219) where
220    Module<BE>: ModuleNew<BE>
221        + ModuleN
222        + GLWESecretPreparedFactory<BE>
223        + GLWEExternalProduct<BE>
224        + GLWEDecrypt<BE>
225        + LWEEncryptSk<BE>
226        + CircuitBootstrappingKeyEncryptSk<BRA, BE>
227        + CircuitBootstrappingKeyPreparedFactory<BRA, BE>
228        + CirtuitBootstrappingExecute<BRA, BE>
229        + GGSWPreparedFactory<BE>
230        + GGSWNoise<BE>
231        + GLWEEncryptSk<BE>
232        + VecZnxRotateInplace<BE>
233        + BDDKeyEncryptSk<BRA, BE>
234        + BDDKeyPreparedFactory<BRA, BE>
235        + FheUintPrepare<BRA, BE>
236        + ExecuteBDDCircuit2WTo1W<BE>,
237    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
238    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
239    Scratch<BE>: ScratchTakeCore<BE>,
240    F: Fn(
241            &mut FheUint<Vec<u8>, u32>,
242            &Module<BE>,
243            &FheUintPrepared<Vec<u8>, u32, BE>,
244            &FheUintPrepared<Vec<u8>, u32, BE>,
245            &BDDKeyPrepared<Vec<u8>, BRA, BE>,
246            &mut Scratch<BE>,
247        ) + 'static,
248{
249    let setup = setup_benchmark::<BE, BRA>(params);
250    let mut runner = create_runner(setup, operation);
251    let id = BenchmarkId::from_parameter(format!("{}_{}", params.name, operation_name));
252    group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
253}
254
255pub fn benc_bdd_arithmetic<BE: Backend, BRA: BlindRotationAlgo>(c: &mut Criterion, label: &str)
256where
257    Module<BE>: ModuleNew<BE>
258        + ModuleN
259        + GLWESecretPreparedFactory<BE>
260        + GLWEExternalProduct<BE>
261        + GLWEDecrypt<BE>
262        + LWEEncryptSk<BE>
263        + CircuitBootstrappingKeyEncryptSk<BRA, BE>
264        + CircuitBootstrappingKeyPreparedFactory<BRA, BE>
265        + CirtuitBootstrappingExecute<BRA, BE>
266        + GGSWPreparedFactory<BE>
267        + GGSWNoise<BE>
268        + GLWEEncryptSk<BE>
269        + VecZnxRotateInplace<BE>
270        + BDDKeyEncryptSk<BRA, BE>
271        + BDDKeyPreparedFactory<BRA, BE>
272        + FheUintPrepare<BRA, BE>
273        + ExecuteBDDCircuit2WTo1W<BE>,
274    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>, // TODO find a way to remove this bound or move it to CBT KEY
275    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
276    Scratch<BE>: ScratchTakeCore<BE>,
277{
278    let group_name: String = format!("bdd_arithmetic::{label}");
279
280    let mut group = c.benchmark_group(group_name);
281
282    const N_GLWE: u32 = 1024;
283    const N_LWE: u32 = 679;
284    const BINARY_BLOCK_SIZE: u32 = 7;
285    const BASE2K: u32 = 15;
286    const RANK: u32 = 2;
287
288    let params: Params = Params {
289        name: String::from(format!("n_glwe={N_GLWE}")),
290        block_size: BINARY_BLOCK_SIZE as usize,
291        glwe_layout: GLWELayout {
292            n: Degree(N_GLWE),
293            base2k: Base2K(BASE2K),
294            k: TorusPrecision(2 * BASE2K),
295            rank: Rank(RANK),
296        },
297        ggsw_layout: GGSWLayout {
298            n: Degree(N_GLWE),
299            base2k: Base2K(BASE2K),
300            k: TorusPrecision(3 * BASE2K),
301            rank: Rank(RANK),
302            dnum: Dnum(3),
303            dsize: Dsize(1),
304        },
305        bdd_layout: BDDKeyLayout {
306            cbt_layout: CircuitBootstrappingKeyLayout {
307                brk_layout: BlindRotationKeyLayout {
308                    n_glwe: Degree(N_GLWE),
309                    n_lwe: Degree(N_LWE),
310                    base2k: Base2K(BASE2K),
311                    k: TorusPrecision(4 * BASE2K),
312                    dnum: Dnum(4),
313                    rank: Rank(RANK),
314                },
315                atk_layout: GLWEAutomorphismKeyLayout {
316                    n: Degree(N_GLWE),
317                    base2k: Base2K(BASE2K),
318                    k: TorusPrecision(4 * BASE2K),
319                    dnum: Dnum(4),
320                    dsize: Dsize(1),
321                    rank: Rank(RANK),
322                },
323                tsk_layout: GGLWEToGGSWKeyLayout {
324                    n: Degree(N_GLWE),
325                    base2k: Base2K(BASE2K),
326                    k: TorusPrecision(4 * BASE2K),
327                    dnum: Dnum(4),
328                    dsize: Dsize(1),
329                    rank: Rank(RANK),
330                },
331            },
332            ks_glwe_layout: Some(GLWESwitchingKeyLayout {
333                n: Degree(N_GLWE),
334                base2k: Base2K(BASE2K),
335                k: TorusPrecision(4 * BASE2K),
336                dnum: Dnum(4),
337                dsize: Dsize(1),
338                rank_in: Rank(RANK),
339                rank_out: Rank(1),
340            }),
341            ks_lwe_layout: GLWEToLWEKeyLayout {
342                n: Degree(N_GLWE),
343                base2k: Base2K(BASE2K),
344                k: TorusPrecision(4 * BASE2K),
345                rank_in: Rank(1),
346                dnum: Dnum(4),
347            },
348        },
349    };
350
351    // Benchmark each operation
352    bench_operation::<BE, BRA, _>(
353        &mut group,
354        &params,
355        "add",
356        |c_enc, module, a, b, key, scratch| {
357            c_enc.add(module, a, b, key, scratch);
358        },
359    );
360
361    bench_operation::<BE, BRA, _>(
362        &mut group,
363        &params,
364        "sub",
365        |c_enc, module, a, b, key, scratch| {
366            c_enc.sub(module, a, b, key, scratch);
367        },
368    );
369
370    bench_operation::<BE, BRA, _>(
371        &mut group,
372        &params,
373        "sll",
374        |c_enc, module, a, b, key, scratch| {
375            c_enc.sll(module, a, b, key, scratch);
376        },
377    );
378
379    bench_operation::<BE, BRA, _>(
380        &mut group,
381        &params,
382        "sra",
383        |c_enc, module, a, b, key, scratch| {
384            c_enc.sra(module, a, b, key, scratch);
385        },
386    );
387
388    bench_operation::<BE, BRA, _>(
389        &mut group,
390        &params,
391        "srl",
392        |c_enc, module, a, b, key, scratch| {
393            c_enc.srl(module, a, b, key, scratch);
394        },
395    );
396
397    bench_operation::<BE, BRA, _>(
398        &mut group,
399        &params,
400        "slt",
401        |c_enc, module, a, b, key, scratch| {
402            c_enc.slt(module, a, b, key, scratch);
403        },
404    );
405
406    bench_operation::<BE, BRA, _>(
407        &mut group,
408        &params,
409        "sltu",
410        |c_enc, module, a, b, key, scratch| {
411            c_enc.sltu(module, a, b, key, scratch);
412        },
413    );
414
415    bench_operation::<BE, BRA, _>(
416        &mut group,
417        &params,
418        "or",
419        |c_enc, module, a, b, key, scratch| {
420            c_enc.or(module, a, b, key, scratch);
421        },
422    );
423
424    bench_operation::<BE, BRA, _>(
425        &mut group,
426        &params,
427        "and",
428        |c_enc, module, a, b, key, scratch| {
429            c_enc.and(module, a, b, key, scratch);
430        },
431    );
432
433    bench_operation::<BE, BRA, _>(
434        &mut group,
435        &params,
436        "xor",
437        |c_enc, module, a, b, key, scratch| {
438            c_enc.xor(module, a, b, key, scratch);
439        },
440    );
441
442    group.finish();
443}
444
445fn bench_bdd_arithmetic_fft64(c: &mut Criterion) {
446    #[cfg(all(feature = "enable-avx", target_arch = "x86_64"))]
447    let label = "fft64_avx";
448    #[cfg(not(all(feature = "enable-avx", target_arch = "x86_64")))]
449    let label = "fft64_ref";
450    benc_bdd_arithmetic::<BackendImpl, CGGI>(c, label);
451}
452
453criterion_group!(benches, bench_bdd_arithmetic_fft64);
454criterion_main!(benches);