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