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(&module, &sk_lwe, &sk_glwe, &mut source_xa, &mut source_xe, scratch.borrow());
109
110        let mut res: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&params.ggsw_infos);
111        let mut cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, BE> =
112            CircuitBootstrappingKeyPrepared::alloc_from_infos(&module, &params.cbt_infos);
113        cbt_prepared.prepare(&module, &cbt_key, scratch.borrow());
114        move || {
115            cbt_prepared.execute_to_constant(
116                &module,
117                &mut res,
118                &ct_lwe,
119                params.k_pt,
120                params.extension_factor,
121                scratch.borrow(),
122            );
123            black_box(());
124        }
125    }
126
127    let params: Params = Params {
128        name: String::from("1-bit"),
129        extension_factor: 1,
130        k_pt: 1,
131        lwe_infos: LWELayout {
132            n: 574_u32.into(),
133            k: 13_u32.into(),
134            base2k: 13_u32.into(),
135        },
136        block_size: 7,
137        ggsw_infos: GGSWLayout {
138            n: 1024_u32.into(),
139            base2k: 13_u32.into(),
140            k: 26_u32.into(),
141            dnum: 2_u32.into(),
142            dsize: 1_u32.into(),
143            rank: 2_u32.into(),
144        },
145        cbt_infos: CircuitBootstrappingKeyLayout {
146            brk_layout: BlindRotationKeyLayout {
147                n_glwe: 1024_u32.into(),
148                n_lwe: 574_u32.into(),
149                base2k: 13_u32.into(),
150                k: 52_u32.into(),
151                dnum: 3_u32.into(),
152                rank: 2_u32.into(),
153            },
154            atk_layout: GLWEAutomorphismKeyLayout {
155                n: 1024_u32.into(),
156                base2k: 13_u32.into(),
157                k: 52_u32.into(),
158                dnum: 3_u32.into(),
159                dsize: Dsize(1),
160                rank: 2_u32.into(),
161            },
162            tsk_layout: GGLWEToGGSWKeyLayout {
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        },
171    };
172
173    let id: BenchmarkId = BenchmarkId::from_parameter(params.name.clone());
174    let mut runner = runner::<BE, BRA>(&params);
175    group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
176
177    group.finish();
178}
179
180fn bench_circuit_bootstrapping_fft64(c: &mut Criterion) {
181    #[cfg(all(feature = "enable-avx", target_arch = "x86_64"))]
182    let label = "fft64_avx";
183    #[cfg(not(all(feature = "enable-avx", target_arch = "x86_64")))]
184    let label = "fft64_ref";
185    benc_circuit_bootstrapping::<BackendImpl, CGGI>(c, label);
186}
187
188criterion_group!(benches, bench_circuit_bootstrapping_fft64);
189criterion_main!(benches);