1use std::hint::black_box;
  2
  3use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
  4use poulpy_backend::{FFT64Avx, FFT64Ref};
  5use poulpy_core::{
  6    GGSWNoise, GLWEDecrypt, GLWEEncryptSk, GLWENoise, ScratchTakeCore,
  7    layouts::{GGSWLayout, GLWELayout, GLWESecretPreparedFactory, prepared::GLWESecretPrepared},
  8};
  9use poulpy_hal::{
 10    api::{ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow},
 11    layouts::{Backend, Module, Scratch, ScratchOwned},
 12    source::Source,
 13};
 14use rand::RngCore;
 15
 16use poulpy_schemes::tfhe::{
 17    bdd_arithmetic::{
 18        BDDKeyEncryptSk, BDDKeyPrepared, BDDKeyPreparedFactory, ExecuteBDDCircuit2WTo1W, FheUint, FheUintPrepare,
 19        FheUintPrepareDebug, FheUintPrepared, FheUintPreparedEncryptSk, FheUintPreparedFactory,
 20        tests::test_suite::TestContext,
 21    },
 22    blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory, CGGI},
 23};
 24
 25pub fn benc_bdd_prepare<BRA: BlindRotationAlgo, BE: Backend>(
 26    c: &mut Criterion,
 27    label: &str,
 28    test_context: &TestContext<BRA, BE>,
 29) where
 30    Module<BE>: ModuleNew<BE>
 31        + GLWESecretPreparedFactory<BE>
 32        + GLWEDecrypt<BE>
 33        + GLWENoise<BE>
 34        + FheUintPreparedFactory<u32, BE>
 35        + FheUintPreparedEncryptSk<u32, BE>
 36        + FheUintPrepareDebug<BRA, u32, BE>
 37        + BDDKeyEncryptSk<BRA, BE>
 38        + BDDKeyPreparedFactory<BRA, BE>
 39        + GGSWNoise<BE>
 40        + FheUintPrepare<BRA, BE>
 41        + ExecuteBDDCircuit2WTo1W<u32, BE>
 42        + GLWEEncryptSk<BE>,
 43    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
 44    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 45    Scratch<BE>: ScratchTakeCore<BE>,
 46{
 47    let group_name: String = format!("bdd_prepare::{label}");
 48
 49    let mut group = c.benchmark_group(group_name);
 50
 51    fn runner<BE: Backend, BRA: BlindRotationAlgo>(test_context: &TestContext<BRA, BE>) -> impl FnMut()
 52    where
 53        Module<BE>: ModuleNew<BE>
 54            + GLWESecretPreparedFactory<BE>
 55            + GLWEDecrypt<BE>
 56            + GLWENoise<BE>
 57            + FheUintPreparedFactory<u32, BE>
 58            + FheUintPreparedEncryptSk<u32, BE>
 59            + FheUintPrepareDebug<BRA, u32, BE>
 60            + BDDKeyEncryptSk<BRA, BE>
 61            + BDDKeyPreparedFactory<BRA, BE>
 62            + GGSWNoise<BE>
 63            + FheUintPrepare<BRA, BE>
 64            + ExecuteBDDCircuit2WTo1W<u32, BE>
 65            + GLWEEncryptSk<BE>,
 66        BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
 67        ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 68        Scratch<BE>: ScratchTakeCore<BE>,
 69    {
 70        let glwe_infos: GLWELayout = test_context.glwe_infos();
 71        let ggsw_infos: GGSWLayout = test_context.ggsw_infos();
 72
 73        let module: &Module<BE> = &test_context.module;
 74        let sk_glwe_prep: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
 75        let bdd_key_prepared: &BDDKeyPrepared<Vec<u8>, BRA, BE> = &test_context.bdd_key;
 76
 77        let mut source: Source = Source::new([6u8; 32]);
 78
 79        let mut source_xa: Source = Source::new([2u8; 32]);
 80        let mut source_xe: Source = Source::new([3u8; 32]);
 81
 82        let threads = 1;
 83
 84        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc((1 << 22) * threads);
 85
 86        // GLWE(value)
 87        let mut c_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_infos);
 88        let value: u32 = source.next_u32();
 89        c_enc.encrypt_sk(
 90            module,
 91            value,
 92            sk_glwe_prep,
 93            &mut source_xa,
 94            &mut source_xe,
 95            scratch.borrow(),
 96        );
 97
 98        // GGSW(0)
 99        let mut c_enc_prep: FheUintPrepared<Vec<u8>, u32, BE> =
100            FheUintPrepared::<Vec<u8>, u32, BE>::alloc_from_infos(module, &ggsw_infos);
101
102        // GGSW(value)
103        move || {
104            c_enc_prep.prepare_custom_multi_thread(threads, module, &c_enc, 0, 32, bdd_key_prepared, scratch.borrow());
105            black_box(());
106        }
107    }
108
109    let id: BenchmarkId = BenchmarkId::from_parameter(format!("n_glwe: {}", test_context.module.n()));
110    let mut runner = runner::<BE, BRA>(test_context);
111    group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
112
113    group.finish();
114}
115
116fn bench_bdd_prepare_cpu_ref_fft64(c: &mut Criterion) {
117    benc_bdd_prepare::<CGGI, FFT64Avx>(
118        c,
119        "bdd_prepare_fft64_ref",
120        &TestContext::<CGGI, FFT64Avx>::new(),
121    );
122}
123
124criterion_group!(benches, bench_bdd_prepare_cpu_ref_fft64,);
125
126criterion_main!(benches);