arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/poulpy-schemes/benches/circuit_bootstrapping.rs
poulpy-schemes/benches/circuit_bootstrapping.rsRaw
  1use std::hint::black_box;
  2
  3use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
  4use poulpy_backend::{FFT64Avx, FFT64Ref, FFT64Spqlios};
  5use poulpy_core::layouts::{
  6    Dsize, GGLWEAutomorphismKeyLayout, GGLWETensorKeyLayout, GGSWCiphertext, GGSWCiphertextLayout, GLWESecret, LWECiphertext,
  7    LWECiphertextLayout, LWESecret, prepared::PrepareAlloc,
  8};
  9use poulpy_hal::{
 10    api::{
 11        ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDft, SvpApplyDftToDftInplace, SvpPPolAlloc,
 12        SvpPPolAllocBytes, SvpPrepare, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphism,
 13        VecZnxAutomorphismInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAlloc, VecZnxBigAllocBytes,
 14        VecZnxBigAutomorphismInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallNegateInplace, VecZnxCopy,
 15        VecZnxDftAddInplace, VecZnxDftAlloc, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxDftCopy, VecZnxFillUniform,
 16        VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNegateInplace, VecZnxNormalize, VecZnxNormalizeInplace,
 17        VecZnxNormalizeTmpBytes, VecZnxRotate, VecZnxRotateInplace, VecZnxRotateInplaceTmpBytes, VecZnxRshInplace, VecZnxSub,
 18        VecZnxSubInplace, VecZnxSwitchRing, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes, VmpPMatAlloc,
 19        VmpPrepare, ZnAddNormal, ZnFillUniform, ZnNormalizeInplace,
 20    },
 21    layouts::{Backend, Module, ScratchOwned},
 22    oep::{
 23        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeMatZnxImpl, TakeScalarZnxImpl, TakeSliceImpl,
 24        TakeSvpPPolImpl, TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxDftSliceImpl, TakeVecZnxImpl, TakeVecZnxSliceImpl,
 25    },
 26    source::Source,
 27};
 28use poulpy_schemes::tfhe::{
 29    blind_rotation::{
 30        BlincRotationExecute, BlindRotationAlgo, BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyEncryptSk,
 31        BlindRotationKeyInfos, BlindRotationKeyLayout, BlindRotationKeyPrepared, CGGI,
 32    },
 33    circuit_bootstrapping::{
 34        CircuitBootstrappingKey, CircuitBootstrappingKeyEncryptSk, CircuitBootstrappingKeyLayout,
 35        CircuitBootstrappingKeyPrepared, CirtuitBootstrappingExecute,
 36    },
 37};
 38
 39pub fn benc_circuit_bootstrapping<B: Backend, BRA: BlindRotationAlgo>(c: &mut Criterion, label: &str)
 40where
 41    Module<B>: ModuleNew<B>
 42        + VecZnxFillUniform
 43        + VecZnxAddNormal
 44        + VecZnxNormalizeInplace<B>
 45        + VecZnxDftAllocBytes
 46        + VecZnxBigNormalize<B>
 47        + VecZnxDftApply<B>
 48        + SvpApplyDftToDftInplace<B>
 49        + VecZnxIdftApplyConsume<B>
 50        + VecZnxNormalizeTmpBytes
 51        + VecZnxSubInplace
 52        + VecZnxAddInplace
 53        + VecZnxNormalize<B>
 54        + VecZnxSub
 55        + VecZnxAddScalarInplace
 56        + VecZnxAutomorphism
 57        + VecZnxSwitchRing
 58        + VecZnxBigAllocBytes
 59        + VecZnxIdftApplyTmpA<B>
 60        + SvpApplyDftToDft<B>
 61        + VecZnxBigAddInplace<B>
 62        + VecZnxBigAddSmallInplace<B>
 63        + VecZnxBigAlloc<B>
 64        + VecZnxDftAlloc<B>
 65        + VecZnxBigNormalizeTmpBytes
 66        + VmpPMatAlloc<B>
 67        + VmpPrepare<B>
 68        + SvpPrepare<B>
 69        + SvpPPolAlloc<B>
 70        + VmpApplyDftToDftTmpBytes
 71        + VmpApplyDftToDft<B>
 72        + VmpApplyDftToDftAdd<B>
 73        + SvpPPolAllocBytes
 74        + VecZnxRotateInplace<B>
 75        + VecZnxBigAutomorphismInplace<B>
 76        + VecZnxRshInplace<B>
 77        + VecZnxDftCopy<B>
 78        + VecZnxNegateInplace
 79        + VecZnxCopy
 80        + VecZnxAutomorphismInplace<B>
 81        + VecZnxBigSubSmallNegateInplace<B>
 82        + VecZnxRotateInplaceTmpBytes
 83        + VecZnxBigAllocBytes
 84        + VecZnxDftAddInplace<B>
 85        + VecZnxRotate
 86        + ZnFillUniform
 87        + ZnAddNormal
 88        + ZnNormalizeInplace<B>,
 89    ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
 90    B: Backend
 91        + ScratchOwnedAllocImpl<B>
 92        + ScratchOwnedBorrowImpl<B>
 93        + TakeVecZnxDftImpl<B>
 94        + ScratchAvailableImpl<B>
 95        + TakeVecZnxImpl<B>
 96        + TakeScalarZnxImpl<B>
 97        + TakeSvpPPolImpl<B>
 98        + TakeVecZnxBigImpl<B>
 99        + TakeVecZnxDftSliceImpl<B>
100        + TakeMatZnxImpl<B>
101        + TakeVecZnxSliceImpl<B>
102        + TakeSliceImpl<B>,
103    BlindRotationKey<Vec<u8>, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
104    BlindRotationKeyPrepared<Vec<u8>, BRA, B>: BlincRotationExecute<B>,
105    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<B>,
106{
107    let group_name: String = format!("circuit_bootstrapping::{label}");
108
109    let mut group = c.benchmark_group(group_name);
110
111    struct Params {
112        name: String,
113        extension_factor: usize,
114        k_pt: usize,
115        block_size: usize,
116        lwe_infos: LWECiphertextLayout,
117        ggsw_infos: GGSWCiphertextLayout,
118        cbt_infos: CircuitBootstrappingKeyLayout,
119    }
120
121    fn runner<B: Backend, BRA: BlindRotationAlgo>(params: &Params) -> impl FnMut()
122    where
123        Module<B>: ModuleNew<B>
124            + VecZnxFillUniform
125            + VecZnxAddNormal
126            + VecZnxNormalizeInplace<B>
127            + VecZnxDftAllocBytes
128            + VecZnxBigNormalize<B>
129            + VecZnxDftApply<B>
130            + SvpApplyDftToDftInplace<B>
131            + VecZnxIdftApplyConsume<B>
132            + VecZnxNormalizeTmpBytes
133            + VecZnxSubInplace
134            + VecZnxAddInplace
135            + VecZnxNormalize<B>
136            + VecZnxSub
137            + VecZnxAddScalarInplace
138            + VecZnxAutomorphism
139            + VecZnxSwitchRing
140            + VecZnxBigAllocBytes
141            + VecZnxIdftApplyTmpA<B>
142            + SvpApplyDftToDft<B>
143            + VecZnxBigAddInplace<B>
144            + VecZnxBigAddSmallInplace<B>
145            + VecZnxBigAlloc<B>
146            + VecZnxDftAlloc<B>
147            + VecZnxBigNormalizeTmpBytes
148            + VmpPMatAlloc<B>
149            + VmpPrepare<B>
150            + SvpPrepare<B>
151            + SvpPPolAlloc<B>
152            + VmpApplyDftToDftTmpBytes
153            + VmpApplyDftToDft<B>
154            + VmpApplyDftToDftAdd<B>
155            + SvpPPolAllocBytes
156            + VecZnxRotateInplace<B>
157            + VecZnxBigAutomorphismInplace<B>
158            + VecZnxRshInplace<B>
159            + VecZnxDftCopy<B>
160            + VecZnxNegateInplace
161            + VecZnxCopy
162            + VecZnxAutomorphismInplace<B>
163            + VecZnxBigSubSmallNegateInplace<B>
164            + VecZnxRotateInplaceTmpBytes
165            + VecZnxBigAllocBytes
166            + VecZnxDftAddInplace<B>
167            + VecZnxRotate
168            + ZnFillUniform
169            + ZnAddNormal
170            + ZnNormalizeInplace<B>,
171        B: Backend
172            + ScratchOwnedAllocImpl<B>
173            + ScratchOwnedBorrowImpl<B>
174            + TakeVecZnxDftImpl<B>
175            + ScratchAvailableImpl<B>
176            + TakeVecZnxImpl<B>
177            + TakeScalarZnxImpl<B>
178            + TakeSvpPPolImpl<B>
179            + TakeVecZnxBigImpl<B>
180            + TakeVecZnxDftSliceImpl<B>
181            + TakeMatZnxImpl<B>
182            + TakeVecZnxSliceImpl<B>
183            + TakeSliceImpl<B>,
184        BlindRotationKey<Vec<u8>, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
185        BlindRotationKeyPrepared<Vec<u8>, BRA, B>: BlincRotationExecute<B>,
186        BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<B>,
187    {
188        // Scratch space (4MB)
189        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(1 << 22);
190
191        let n_glwe: poulpy_core::layouts::Degree = params.cbt_infos.layout_brk.n_glwe();
192        let n_lwe: poulpy_core::layouts::Degree = params.cbt_infos.layout_brk.n_lwe();
193        let rank: poulpy_core::layouts::Rank = params.cbt_infos.layout_brk.rank;
194
195        let module: Module<B> = Module::<B>::new(n_glwe.as_u32() as u64);
196
197        let mut source_xs: Source = Source::new([1u8; 32]);
198        let mut source_xa: Source = Source::new([1u8; 32]);
199        let mut source_xe: Source = Source::new([1u8; 32]);
200
201        let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
202        sk_lwe.fill_binary_block(params.block_size, &mut source_xs);
203        sk_lwe.fill_zero();
204
205        let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc_with(n_glwe, rank);
206        sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
207
208        let ct_lwe: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(&params.lwe_infos);
209
210        // Circuit bootstrapping evaluation key
211        let cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::encrypt_sk(
212            &module,
213            &sk_lwe,
214            &sk_glwe,
215            &params.cbt_infos,
216            &mut source_xa,
217            &mut source_xe,
218            scratch.borrow(),
219        );
220
221        let mut res: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(&params.ggsw_infos);
222        let cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B> = cbt_key.prepare_alloc(&module, scratch.borrow());
223
224        move || {
225            cbt_prepared.execute_to_constant(
226                &module,
227                &mut res,
228                &ct_lwe,
229                params.k_pt,
230                params.extension_factor,
231                scratch.borrow(),
232            );
233            black_box(());
234        }
235    }
236
237    for params in [Params {
238        name: String::from("1-bit"),
239        extension_factor: 1,
240        k_pt: 1,
241        lwe_infos: LWECiphertextLayout {
242            n: 574_u32.into(),
243            k: 13_u32.into(),
244            base2k: 13_u32.into(),
245        },
246        block_size: 7,
247        ggsw_infos: GGSWCiphertextLayout {
248            n: 1024_u32.into(),
249            base2k: 13_u32.into(),
250            k: 26_u32.into(),
251            dnum: 2_u32.into(),
252            dsize: 1_u32.into(),
253            rank: 2_u32.into(),
254        },
255        cbt_infos: CircuitBootstrappingKeyLayout {
256            layout_brk: BlindRotationKeyLayout {
257                n_glwe: 1024_u32.into(),
258                n_lwe: 574_u32.into(),
259                base2k: 13_u32.into(),
260                k: 52_u32.into(),
261                dnum: 3_u32.into(),
262                rank: 2_u32.into(),
263            },
264            layout_atk: GGLWEAutomorphismKeyLayout {
265                n: 1024_u32.into(),
266                base2k: 13_u32.into(),
267                k: 52_u32.into(),
268                dnum: 3_u32.into(),
269                dsize: Dsize(1),
270                rank: 2_u32.into(),
271            },
272            layout_tsk: GGLWETensorKeyLayout {
273                n: 1024_u32.into(),
274                base2k: 13_u32.into(),
275                k: 52_u32.into(),
276                dnum: 3_u32.into(),
277                dsize: Dsize(1),
278                rank: 2_u32.into(),
279            },
280        },
281    }] {
282        let id: BenchmarkId = BenchmarkId::from_parameter(params.name.clone());
283        let mut runner = runner::<B, BRA>(&params);
284        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
285    }
286
287    group.finish();
288}
289
290fn bench_circuit_bootstrapping_cpu_ref_fft64(c: &mut Criterion) {
291    benc_circuit_bootstrapping::<FFT64Ref, CGGI>(c, "fft64_ref");
292}
293
294fn bench_circuit_bootstrapping_cpu_avx_fft64(c: &mut Criterion) {
295    benc_circuit_bootstrapping::<FFT64Avx, CGGI>(c, "fft64_avx");
296}
297
298fn bench_circuit_bootstrapping_cpu_spqlios_fft64(c: &mut Criterion) {
299    benc_circuit_bootstrapping::<FFT64Spqlios, CGGI>(c, "fft64_spqlios");
300}
301
302criterion_group!(
303    benches,
304    bench_circuit_bootstrapping_cpu_ref_fft64,
305    bench_circuit_bootstrapping_cpu_avx_fft64,
306    bench_circuit_bootstrapping_cpu_spqlios_fft64,
307);
308
309criterion_main!(benches);