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