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