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.layout_brk.n_glwe();
 89        let n_lwe: poulpy_core::layouts::Degree = params.cbt_infos.layout_brk.n_lwe();
 90        let rank: poulpy_core::layouts::Rank = params.cbt_infos.layout_brk.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        sk_lwe.fill_zero();
101
102        let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n_glwe, rank);
103        sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
104
105        let ct_lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(&params.lwe_infos);
106
107        // Circuit bootstrapping evaluation key
108        let mut cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::alloc_from_infos(&params.cbt_infos);
109        cbt_key.encrypt_sk(
110            &module,
111            &sk_lwe,
112            &sk_glwe,
113            &mut source_xa,
114            &mut source_xe,
115            scratch.borrow(),
116        );
117
118        let mut res: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&params.ggsw_infos);
119        let mut cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, BE> =
120            CircuitBootstrappingKeyPrepared::alloc_from_infos(&module, &params.cbt_infos);
121        cbt_prepared.prepare(&module, &cbt_key, scratch.borrow());
122        move || {
123            cbt_prepared.execute_to_constant(
124                &module,
125                &mut res,
126                &ct_lwe,
127                params.k_pt,
128                params.extension_factor,
129                scratch.borrow(),
130            );
131            black_box(());
132        }
133    }
134
135    let params: Params = Params {
136        name: String::from("1-bit"),
137        extension_factor: 1,
138        k_pt: 1,
139        lwe_infos: LWELayout {
140            n: 574_u32.into(),
141            k: 13_u32.into(),
142            base2k: 13_u32.into(),
143        },
144        block_size: 7,
145        ggsw_infos: GGSWLayout {
146            n: 1024_u32.into(),
147            base2k: 13_u32.into(),
148            k: 26_u32.into(),
149            dnum: 2_u32.into(),
150            dsize: 1_u32.into(),
151            rank: 2_u32.into(),
152        },
153        cbt_infos: CircuitBootstrappingKeyLayout {
154            layout_brk: BlindRotationKeyLayout {
155                n_glwe: 1024_u32.into(),
156                n_lwe: 574_u32.into(),
157                base2k: 13_u32.into(),
158                k: 52_u32.into(),
159                dnum: 3_u32.into(),
160                rank: 2_u32.into(),
161            },
162            layout_atk: GLWEAutomorphismKeyLayout {
163                n: 1024_u32.into(),
164                base2k: 13_u32.into(),
165                k: 52_u32.into(),
166                dnum: 3_u32.into(),
167                dsize: Dsize(1),
168                rank: 2_u32.into(),
169            },
170            layout_tsk: GGLWEToGGSWKeyLayout {
171                n: 1024_u32.into(),
172                base2k: 13_u32.into(),
173                k: 52_u32.into(),
174                dnum: 3_u32.into(),
175                dsize: Dsize(1),
176                rank: 2_u32.into(),
177            },
178        },
179    };
180
181    let id: BenchmarkId = BenchmarkId::from_parameter(params.name.clone());
182    let mut runner = runner::<BE, BRA>(&params);
183    group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
184
185    group.finish();
186}
187
188fn bench_circuit_bootstrapping_fft64(c: &mut Criterion) {
189    #[cfg(all(feature = "enable-avx", target_arch = "x86_64"))]
190    let label = "fft64_avx";
191    #[cfg(not(all(feature = "enable-avx", target_arch = "x86_64")))]
192    let label = "fft64_ref";
193    benc_circuit_bootstrapping::<BackendImpl, CGGI>(c, label);
194}
195
196criterion_group!(benches, bench_circuit_bootstrapping_fft64);
197criterion_main!(benches);