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        Dsize, GGLWEToGGSWKeyLayout, GGSW, GGSWLayout, GGSWPreparedFactory, GLWEAutomorphismKeyLayout, GLWESecret,
  8        GLWESecretPreparedFactory, LWE, LWELayout, LWESecret,
  9    },
 10};
 11
 12#[cfg(all(feature = "enable-avx", target_arch = "x86_64"))]
 13pub use poulpy_cpu_avx::FFT64Avx as BackendImpl;
 14
 15#[cfg(not(all(feature = "enable-avx", target_arch = "x86_64")))]
 16pub use poulpy_cpu_ref::FFT64Ref as BackendImpl;
 17
 18use poulpy_hal::{
 19    api::{ModuleN, ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxRotateInplace},
 20    layouts::{Backend, Module, Scratch, ScratchOwned},
 21    source::Source,
 22};
 23use poulpy_schemes::bin_fhe::{
 24    blind_rotation::{
 25        BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory, BlindRotationKeyInfos, BlindRotationKeyLayout, CGGI,
 26    },
 27    circuit_bootstrapping::{
 28        CircuitBootstrappingKey, CircuitBootstrappingKeyEncryptSk, CircuitBootstrappingKeyLayout,
 29        CircuitBootstrappingKeyPrepared, CircuitBootstrappingKeyPreparedFactory, CirtuitBootstrappingExecute,
 30    },
 31};
 32
 33pub fn benc_circuit_bootstrapping<BE: Backend, BRA: BlindRotationAlgo>(c: &mut Criterion, label: &str)
 34where
 35    Module<BE>: ModuleNew<BE>
 36        + ModuleN
 37        + GLWESecretPreparedFactory<BE>
 38        + GLWEExternalProduct<BE>
 39        + GLWEDecrypt<BE>
 40        + LWEEncryptSk<BE>
 41        + CircuitBootstrappingKeyEncryptSk<BRA, BE>
 42        + CircuitBootstrappingKeyPreparedFactory<BRA, BE>
 43        + CirtuitBootstrappingExecute<BRA, BE>
 44        + GGSWPreparedFactory<BE>
 45        + GGSWNoise<BE>
 46        + GLWEEncryptSk<BE>
 47        + VecZnxRotateInplace<BE>,
 48    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>, // TODO find a way to remove this bound or move it to CBT KEY
 49    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 50    Scratch<BE>: ScratchTakeCore<BE>,
 51{
 52    let group_name: String = format!("circuit_bootstrapping::{label}");
 53
 54    let mut group = c.benchmark_group(group_name);
 55
 56    struct Params {
 57        name: String,
 58        extension_factor: usize,
 59        k_pt: usize,
 60        block_size: usize,
 61        lwe_infos: LWELayout,
 62        ggsw_infos: GGSWLayout,
 63        cbt_infos: CircuitBootstrappingKeyLayout,
 64    }
 65
 66    fn runner<BE: Backend, BRA: BlindRotationAlgo>(params: &Params) -> impl FnMut()
 67    where
 68        Module<BE>: ModuleNew<BE>
 69            + ModuleN
 70            + GLWESecretPreparedFactory<BE>
 71            + GLWEExternalProduct<BE>
 72            + GLWEDecrypt<BE>
 73            + LWEEncryptSk<BE>
 74            + CircuitBootstrappingKeyEncryptSk<BRA, BE>
 75            + CircuitBootstrappingKeyPreparedFactory<BRA, BE>
 76            + CirtuitBootstrappingExecute<BRA, BE>
 77            + GGSWPreparedFactory<BE>
 78            + GGSWNoise<BE>
 79            + GLWEEncryptSk<BE>
 80            + VecZnxRotateInplace<BE>,
 81        BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>, /* TODO find a way to remove this bound or move it to CBT KEY */
 82        ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 83        Scratch<BE>: ScratchTakeCore<BE>,
 84    {
 85        // Scratch space (4MB)
 86        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
 87
 88        let n_glwe: poulpy_core::layouts::Degree = params.cbt_infos.brk_layout.n_glwe();
 89        let n_lwe: poulpy_core::layouts::Degree = params.cbt_infos.brk_layout.n_lwe();
 90        let rank: poulpy_core::layouts::Rank = params.cbt_infos.brk_layout.rank;
 91
 92        let module: Module<BE> = Module::<BE>::new(n_glwe.as_u32() as u64);
 93
 94        let mut source_xs: Source = Source::new([1u8; 32]);
 95        let mut source_xa: Source = Source::new([1u8; 32]);
 96        let mut source_xe: Source = Source::new([1u8; 32]);
 97
 98        let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
 99        sk_lwe.fill_binary_block(params.block_size, &mut source_xs);
100
101        let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n_glwe, rank);
102        sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
103
104        let ct_lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(&params.lwe_infos);
105
106        // Circuit bootstrapping evaluation key
107        let mut cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::alloc_from_infos(&params.cbt_infos);
108        cbt_key.encrypt_sk(
109            &module,
110            &sk_lwe,
111            &sk_glwe,
112            &mut source_xa,
113            &mut source_xe,
114            scratch.borrow(),
115        );
116
117        let mut res: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&params.ggsw_infos);
118        let mut cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, BE> =
119            CircuitBootstrappingKeyPrepared::alloc_from_infos(&module, &params.cbt_infos);
120        cbt_prepared.prepare(&module, &cbt_key, scratch.borrow());
121        move || {
122            cbt_prepared.execute_to_constant(
123                &module,
124                &mut res,
125                &ct_lwe,
126                params.k_pt,
127                params.extension_factor,
128                scratch.borrow(),
129            );
130            black_box(());
131        }
132    }
133
134    let params: Params = Params {
135        name: String::from("1-bit"),
136        extension_factor: 1,
137        k_pt: 1,
138        lwe_infos: LWELayout {
139            n: 574_u32.into(),
140            k: 13_u32.into(),
141            base2k: 13_u32.into(),
142        },
143        block_size: 7,
144        ggsw_infos: GGSWLayout {
145            n: 1024_u32.into(),
146            base2k: 13_u32.into(),
147            k: 26_u32.into(),
148            dnum: 2_u32.into(),
149            dsize: 1_u32.into(),
150            rank: 2_u32.into(),
151        },
152        cbt_infos: CircuitBootstrappingKeyLayout {
153            brk_layout: BlindRotationKeyLayout {
154                n_glwe: 1024_u32.into(),
155                n_lwe: 574_u32.into(),
156                base2k: 13_u32.into(),
157                k: 52_u32.into(),
158                dnum: 3_u32.into(),
159                rank: 2_u32.into(),
160            },
161            atk_layout: GLWEAutomorphismKeyLayout {
162                n: 1024_u32.into(),
163                base2k: 13_u32.into(),
164                k: 52_u32.into(),
165                dnum: 3_u32.into(),
166                dsize: Dsize(1),
167                rank: 2_u32.into(),
168            },
169            tsk_layout: GGLWEToGGSWKeyLayout {
170                n: 1024_u32.into(),
171                base2k: 13_u32.into(),
172                k: 52_u32.into(),
173                dnum: 3_u32.into(),
174                dsize: Dsize(1),
175                rank: 2_u32.into(),
176            },
177        },
178    };
179
180    let id: BenchmarkId = BenchmarkId::from_parameter(params.name.clone());
181    let mut runner = runner::<BE, BRA>(&params);
182    group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
183
184    group.finish();
185}
186
187fn bench_circuit_bootstrapping_fft64(c: &mut Criterion) {
188    #[cfg(all(feature = "enable-avx", target_arch = "x86_64"))]
189    let label = "fft64_avx";
190    #[cfg(not(all(feature = "enable-avx", target_arch = "x86_64")))]
191    let label = "fft64_ref";
192    benc_circuit_bootstrapping::<BackendImpl, CGGI>(c, label);
193}
194
195criterion_group!(benches, bench_circuit_bootstrapping_fft64);
196criterion_main!(benches);