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