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(&module, &sk_lwe, &sk_glwe, &mut source_xa, &mut source_xe, scratch.borrow());
 95
 96    let mut cbt_key_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, BE> =
 97        CircuitBootstrappingKeyPrepared::alloc_from_infos(&module, &params.bdd_layout.cbt_layout);
 98    cbt_key_prepared.prepare(&module, &cbt_key, scratch.borrow());
 99
100    let mut sk_glwe_prepared = GLWESecretPrepared::alloc_from_infos(&module, &params.glwe_layout);
101    sk_glwe_prepared.prepare(&module, &sk_glwe);
102
103    let mut bdd_key: BDDKey<Vec<u8>, BRA> = BDDKey::alloc_from_infos(&params.bdd_layout);
104    bdd_key.encrypt_sk(&module, &sk_lwe, &sk_glwe, &mut source_xa, &mut source_xe, scratch.borrow());
105
106    let input_a = 255_u32;
107    let input_b = 30_u32;
108
109    let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&params.glwe_layout);
110    a_enc.encrypt_sk(
111        &module,
112        input_a,
113        &sk_glwe_prepared,
114        &mut source_xa,
115        &mut source_xe,
116        scratch.borrow(),
117    );
118
119    let mut b_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&params.glwe_layout);
120    b_enc.encrypt_sk(
121        &module,
122        input_b,
123        &sk_glwe_prepared,
124        &mut source_xa,
125        &mut source_xe,
126        scratch.borrow(),
127    );
128
129    //////// Homomorphic computation starts here ////////
130
131    // Preparing the BDD Key
132    // The BDD key must be prepared once before any operation is performed
133    let mut bdd_key_prepared: BDDKeyPrepared<Vec<u8>, BRA, BE> = BDDKeyPrepared::alloc_from_infos(&module, &params.bdd_layout);
134    bdd_key_prepared.prepare(&module, &bdd_key, scratch.borrow());
135
136    // Input Preparation
137    // Before each operation, the inputs to that operation must be prepared
138    // Preparation extracts each bit of the integer into a seperate GLWE ciphertext and bootstraps it into a GGSW ciphertext
139    let mut a_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE> = FheUintPrepared::alloc_from_infos(&module, &params.ggsw_layout);
140    a_enc_prepared.prepare(&module, &a_enc, &bdd_key_prepared, scratch.borrow());
141
142    let mut b_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE> = FheUintPrepared::alloc_from_infos(&module, &params.ggsw_layout);
143    b_enc_prepared.prepare(&module, &b_enc, &bdd_key_prepared, scratch.borrow());
144
145    BenchmarkSetup {
146        module,
147        scratch,
148        a_enc_prepared,
149        b_enc_prepared,
150        bdd_key_prepared,
151        glwe_layout: params.glwe_layout,
152    }
153}
154
155fn create_runner<BE: Backend, BRA: BlindRotationAlgo, F>(setup: BenchmarkSetup<BE, BRA>, operation: F) -> impl FnMut()
156where
157    Module<BE>: ExecuteBDDCircuit2WTo1W<BE>,
158    Scratch<BE>: ScratchTakeCore<BE>,
159    ScratchOwned<BE>: ScratchOwnedBorrow<BE>,
160    F: Fn(
161        &mut FheUint<Vec<u8>, u32>,
162        &Module<BE>,
163        &FheUintPrepared<Vec<u8>, u32, BE>,
164        &FheUintPrepared<Vec<u8>, u32, BE>,
165        &BDDKeyPrepared<Vec<u8>, BRA, BE>,
166        &mut Scratch<BE>,
167    ),
168{
169    let BenchmarkSetup {
170        module,
171        mut scratch,
172        a_enc_prepared,
173        b_enc_prepared,
174        bdd_key_prepared,
175        glwe_layout,
176    } = setup;
177
178    let mut c_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_layout);
179
180    move || {
181        operation(
182            &mut c_enc,
183            &module,
184            &a_enc_prepared,
185            &b_enc_prepared,
186            &bdd_key_prepared,
187            scratch.borrow(),
188        );
189        black_box(());
190    }
191}
192
193fn bench_operation<BE: Backend, BRA: BlindRotationAlgo, F>(
194    group: &mut criterion::BenchmarkGroup<'_, criterion::measurement::WallTime>,
195    params: &Params,
196    operation_name: &str,
197    operation: F,
198) where
199    Module<BE>: ModuleNew<BE>
200        + ModuleN
201        + GLWESecretPreparedFactory<BE>
202        + GLWEDecrypt<BE>
203        + LWEEncryptSk<BE>
204        + CircuitBootstrappingKeyEncryptSk<BRA, BE>
205        + GGSWPreparedFactory<BE>
206        + GLWEEncryptSk<BE>
207        + BDDKeyEncryptSk<BRA, BE>
208        + BDDKeyPreparedFactory<BRA, BE>
209        + FheUintPrepare<BRA, BE>
210        + ExecuteBDDCircuit2WTo1W<BE>,
211    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
212    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
213    Scratch<BE>: ScratchTakeCore<BE>,
214    F: Fn(
215            &mut FheUint<Vec<u8>, u32>,
216            &Module<BE>,
217            &FheUintPrepared<Vec<u8>, u32, BE>,
218            &FheUintPrepared<Vec<u8>, u32, BE>,
219            &BDDKeyPrepared<Vec<u8>, BRA, BE>,
220            &mut Scratch<BE>,
221        ) + 'static,
222{
223    let setup = setup_benchmark::<BE, BRA>(params);
224    let mut runner = create_runner(setup, operation);
225    let id = BenchmarkId::from_parameter(format!("{}_{}", params.name, operation_name));
226    group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
227}
228
229pub fn benc_bdd_arithmetic<BE: Backend, BRA: BlindRotationAlgo>(c: &mut Criterion, label: &str)
230where
231    Module<BE>: ModuleNew<BE>
232        + ModuleN
233        + GLWESecretPreparedFactory<BE>
234        + GLWEDecrypt<BE>
235        + LWEEncryptSk<BE>
236        + CircuitBootstrappingKeyEncryptSk<BRA, BE>
237        + GGSWPreparedFactory<BE>
238        + GLWEEncryptSk<BE>
239        + BDDKeyEncryptSk<BRA, BE>
240        + BDDKeyPreparedFactory<BRA, BE>
241        + FheUintPrepare<BRA, BE>
242        + ExecuteBDDCircuit2WTo1W<BE>,
243    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>, // TODO find a way to remove this bound or move it to CBT KEY
244    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
245    Scratch<BE>: ScratchTakeCore<BE>,
246{
247    let group_name: String = format!("bdd_arithmetic::{label}");
248
249    let mut group = c.benchmark_group(group_name);
250
251    const N_GLWE: u32 = 1024;
252    const N_LWE: u32 = 679;
253    const BINARY_BLOCK_SIZE: u32 = 7;
254    const BASE2K: u32 = 15;
255    const RANK: u32 = 2;
256
257    let params: Params = Params {
258        name: format!("n_glwe={N_GLWE}"),
259        block_size: BINARY_BLOCK_SIZE as usize,
260        glwe_layout: GLWELayout {
261            n: Degree(N_GLWE),
262            base2k: Base2K(BASE2K),
263            k: TorusPrecision(2 * BASE2K),
264            rank: Rank(RANK),
265        },
266        ggsw_layout: GGSWLayout {
267            n: Degree(N_GLWE),
268            base2k: Base2K(BASE2K),
269            k: TorusPrecision(3 * BASE2K),
270            rank: Rank(RANK),
271            dnum: Dnum(3),
272            dsize: Dsize(1),
273        },
274        bdd_layout: BDDKeyLayout {
275            cbt_layout: CircuitBootstrappingKeyLayout {
276                brk_layout: BlindRotationKeyLayout {
277                    n_glwe: Degree(N_GLWE),
278                    n_lwe: Degree(N_LWE),
279                    base2k: Base2K(BASE2K),
280                    k: TorusPrecision(4 * BASE2K),
281                    dnum: Dnum(4),
282                    rank: Rank(RANK),
283                },
284                atk_layout: GLWEAutomorphismKeyLayout {
285                    n: Degree(N_GLWE),
286                    base2k: Base2K(BASE2K),
287                    k: TorusPrecision(4 * BASE2K),
288                    dnum: Dnum(4),
289                    dsize: Dsize(1),
290                    rank: Rank(RANK),
291                },
292                tsk_layout: GGLWEToGGSWKeyLayout {
293                    n: Degree(N_GLWE),
294                    base2k: Base2K(BASE2K),
295                    k: TorusPrecision(4 * BASE2K),
296                    dnum: Dnum(4),
297                    dsize: Dsize(1),
298                    rank: Rank(RANK),
299                },
300            },
301            ks_glwe_layout: Some(GLWESwitchingKeyLayout {
302                n: Degree(N_GLWE),
303                base2k: Base2K(BASE2K),
304                k: TorusPrecision(4 * BASE2K),
305                dnum: Dnum(4),
306                dsize: Dsize(1),
307                rank_in: Rank(RANK),
308                rank_out: Rank(1),
309            }),
310            ks_lwe_layout: GLWEToLWEKeyLayout {
311                n: Degree(N_GLWE),
312                base2k: Base2K(BASE2K),
313                k: TorusPrecision(4 * BASE2K),
314                rank_in: Rank(1),
315                dnum: Dnum(4),
316            },
317        },
318    };
319
320    // Benchmark each operation
321    bench_operation::<BE, BRA, _>(&mut group, &params, "add", |c_enc, module, a, b, key, scratch| {
322        c_enc.add(module, a, b, key, scratch);
323    });
324
325    bench_operation::<BE, BRA, _>(&mut group, &params, "sub", |c_enc, module, a, b, key, scratch| {
326        c_enc.sub(module, a, b, key, scratch);
327    });
328
329    bench_operation::<BE, BRA, _>(&mut group, &params, "sll", |c_enc, module, a, b, key, scratch| {
330        c_enc.sll(module, a, b, key, scratch);
331    });
332
333    bench_operation::<BE, BRA, _>(&mut group, &params, "sra", |c_enc, module, a, b, key, scratch| {
334        c_enc.sra(module, a, b, key, scratch);
335    });
336
337    bench_operation::<BE, BRA, _>(&mut group, &params, "srl", |c_enc, module, a, b, key, scratch| {
338        c_enc.srl(module, a, b, key, scratch);
339    });
340
341    bench_operation::<BE, BRA, _>(&mut group, &params, "slt", |c_enc, module, a, b, key, scratch| {
342        c_enc.slt(module, a, b, key, scratch);
343    });
344
345    bench_operation::<BE, BRA, _>(&mut group, &params, "sltu", |c_enc, module, a, b, key, scratch| {
346        c_enc.sltu(module, a, b, key, scratch);
347    });
348
349    bench_operation::<BE, BRA, _>(&mut group, &params, "or", |c_enc, module, a, b, key, scratch| {
350        c_enc.or(module, a, b, key, scratch);
351    });
352
353    bench_operation::<BE, BRA, _>(&mut group, &params, "and", |c_enc, module, a, b, key, scratch| {
354        c_enc.and(module, a, b, key, scratch);
355    });
356
357    bench_operation::<BE, BRA, _>(&mut group, &params, "xor", |c_enc, module, a, b, key, scratch| {
358        c_enc.xor(module, a, b, key, scratch);
359    });
360
361    group.finish();
362}
363
364fn bench_bdd_arithmetic_fft64(c: &mut Criterion) {
365    #[cfg(all(feature = "enable-avx", target_arch = "x86_64"))]
366    let label = "fft64_avx";
367    #[cfg(not(all(feature = "enable-avx", target_arch = "x86_64")))]
368    let label = "fft64_ref";
369    benc_bdd_arithmetic::<BackendImpl, CGGI>(c, label);
370}
371
372criterion_group!(benches, bench_bdd_arithmetic_fft64);
373criterion_main!(benches);