1use std::hint::black_box;
  2
  3use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
  4use poulpy_core::{
  5    GLWEDecrypt, GLWEEncryptSk, 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        BDDKey, BDDKeyEncryptSk, BDDKeyLayout, BDDKeyPrepared, BDDKeyPreparedFactory, FheUint, FheUintPrepare, FheUintPrepared,
 27    },
 28    blind_rotation::{
 29        BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory, BlindRotationKeyInfos, BlindRotationKeyLayout, CGGI,
 30    },
 31    circuit_bootstrapping::{CircuitBootstrappingKeyEncryptSk, CircuitBootstrappingKeyLayout},
 32};
 33
 34pub fn benc_bdd_prepare<BE: Backend, BRA: BlindRotationAlgo>(c: &mut Criterion, label: &str)
 35where
 36    Module<BE>: ModuleNew<BE>
 37        + ModuleN
 38        + GLWESecretPreparedFactory<BE>
 39        + GLWEDecrypt<BE>
 40        + CircuitBootstrappingKeyEncryptSk<BRA, BE>
 41        + GGSWPreparedFactory<BE>
 42        + GLWEEncryptSk<BE>
 43        + BDDKeyEncryptSk<BRA, BE>
 44        + BDDKeyPreparedFactory<BRA, BE>
 45        + FheUintPrepare<BRA, BE>,
 46    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>, // TODO find a way to remove this bound or move it to CBT KEY
 47    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 48    Scratch<BE>: ScratchTakeCore<BE>,
 49{
 50    let group_name: String = format!("bdd_arithmetic::{label}");
 51
 52    let mut group = c.benchmark_group(group_name);
 53
 54    struct Params {
 55        name: String,
 56        block_size: usize,
 57        glwe_layout: GLWELayout,
 58        ggsw_layout: GGSWLayout,
 59        bdd_layout: BDDKeyLayout,
 60    }
 61
 62    fn runner<BE: Backend, BRA: BlindRotationAlgo>(params: &Params) -> impl FnMut()
 63    where
 64        Module<BE>: ModuleNew<BE>
 65            + ModuleN
 66            + GLWESecretPreparedFactory<BE>
 67            + GLWEDecrypt<BE>
 68            + CircuitBootstrappingKeyEncryptSk<BRA, BE>
 69            + GGSWPreparedFactory<BE>
 70            + GLWEEncryptSk<BE>
 71            + BDDKeyEncryptSk<BRA, BE>
 72            + BDDKeyPreparedFactory<BRA, BE>
 73            + FheUintPrepare<BRA, BE>,
 74        BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>, /* TODO find a way to remove this bound or move it to CBT KEY */
 75        ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 76        Scratch<BE>: ScratchTakeCore<BE>,
 77    {
 78        // Scratch space (4MB)
 79        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 25);
 80
 81        let n_glwe: poulpy_core::layouts::Degree = params.bdd_layout.cbt_layout.brk_layout.n_glwe();
 82        let n_lwe: poulpy_core::layouts::Degree = params.bdd_layout.cbt_layout.brk_layout.n_lwe();
 83        let rank: poulpy_core::layouts::Rank = params.bdd_layout.cbt_layout.brk_layout.rank;
 84
 85        let module: Module<BE> = Module::<BE>::new(n_glwe.as_u32() as u64);
 86
 87        let mut source_xs: Source = Source::new([1u8; 32]);
 88        let mut source_xa: Source = Source::new([1u8; 32]);
 89        let mut source_xe: Source = Source::new([1u8; 32]);
 90
 91        let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
 92        sk_lwe.fill_binary_block(params.block_size, &mut source_xs);
 93
 94        let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n_glwe, rank);
 95        sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
 96
 97        let mut sk_glwe_prepared = GLWESecretPrepared::alloc_from_infos(&module, &params.glwe_layout);
 98        sk_glwe_prepared.prepare(&module, &sk_glwe);
 99
100        let mut bdd_key: BDDKey<Vec<u8>, BRA> = BDDKey::alloc_from_infos(&params.bdd_layout);
101        bdd_key.encrypt_sk(&module, &sk_lwe, &sk_glwe, &mut source_xa, &mut source_xe, scratch.borrow());
102
103        let input_a = 255_u32;
104
105        let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&params.glwe_layout);
106        a_enc.encrypt_sk(
107            &module,
108            input_a,
109            &sk_glwe_prepared,
110            &mut source_xa,
111            &mut source_xe,
112            scratch.borrow(),
113        );
114
115        let mut bdd_key_prepared: BDDKeyPrepared<Vec<u8>, BRA, BE> =
116            BDDKeyPrepared::alloc_from_infos(&module, &params.bdd_layout);
117        bdd_key_prepared.prepare(&module, &bdd_key, scratch.borrow());
118
119        let mut a_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE> =
120            FheUintPrepared::alloc_from_infos(&module, &params.ggsw_layout);
121
122        move || {
123            a_enc_prepared.prepare(&module, &a_enc, &bdd_key_prepared, scratch.borrow());
124            black_box(());
125        }
126    }
127
128    const N_GLWE: u32 = 1024;
129    const N_LWE: u32 = 679;
130    const BINARY_BLOCK_SIZE: u32 = 7;
131    const BASE2K: u32 = 15;
132    const RANK: u32 = 2;
133
134    let params: Params = Params {
135        name: format!("n_glwe={N_GLWE}"),
136        block_size: BINARY_BLOCK_SIZE as usize,
137        glwe_layout: GLWELayout {
138            n: Degree(N_GLWE),
139            base2k: Base2K(BASE2K),
140            k: TorusPrecision(2 * BASE2K),
141            rank: Rank(RANK),
142        },
143        ggsw_layout: GGSWLayout {
144            n: Degree(N_GLWE),
145            base2k: Base2K(BASE2K),
146            k: TorusPrecision(3 * BASE2K),
147            rank: Rank(RANK),
148            dnum: Dnum(3),
149            dsize: Dsize(1),
150        },
151        bdd_layout: BDDKeyLayout {
152            cbt_layout: CircuitBootstrappingKeyLayout {
153                brk_layout: BlindRotationKeyLayout {
154                    n_glwe: Degree(N_GLWE),
155                    n_lwe: Degree(N_LWE),
156                    base2k: Base2K(BASE2K),
157                    k: TorusPrecision(4 * BASE2K),
158                    dnum: Dnum(4),
159                    rank: Rank(RANK),
160                },
161                atk_layout: GLWEAutomorphismKeyLayout {
162                    n: Degree(N_GLWE),
163                    base2k: Base2K(BASE2K),
164                    k: TorusPrecision(4 * BASE2K),
165                    dnum: Dnum(4),
166                    dsize: Dsize(1),
167                    rank: Rank(RANK),
168                },
169                tsk_layout: GGLWEToGGSWKeyLayout {
170                    n: Degree(N_GLWE),
171                    base2k: Base2K(BASE2K),
172                    k: TorusPrecision(4 * BASE2K),
173                    dnum: Dnum(4),
174                    dsize: Dsize(1),
175                    rank: Rank(RANK),
176                },
177            },
178            ks_glwe_layout: Some(GLWESwitchingKeyLayout {
179                n: Degree(N_GLWE),
180                base2k: Base2K(BASE2K),
181                k: TorusPrecision(4 * BASE2K),
182                dnum: Dnum(4),
183                dsize: Dsize(1),
184                rank_in: Rank(RANK),
185                rank_out: Rank(1),
186            }),
187            ks_lwe_layout: GLWEToLWEKeyLayout {
188                n: Degree(N_GLWE),
189                base2k: Base2K(BASE2K),
190                k: TorusPrecision(4 * BASE2K),
191                rank_in: Rank(1),
192                dnum: Dnum(4),
193            },
194        },
195    };
196
197    let id: BenchmarkId = BenchmarkId::from_parameter(params.name.clone());
198    let mut runner = runner::<BE, BRA>(&params);
199    group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
200
201    group.finish();
202}
203
204fn bench_bdd_prepare_fft64(c: &mut Criterion) {
205    #[cfg(all(feature = "enable-avx", target_arch = "x86_64"))]
206    let label = "fft64_avx";
207    #[cfg(not(all(feature = "enable-avx", target_arch = "x86_64")))]
208    let label = "fft64_ref";
209    benc_bdd_prepare::<BackendImpl, CGGI>(c, label);
210}
211
212criterion_group!(benches, bench_bdd_prepare_fft64);
213criterion_main!(benches);