poulpy/poulpy-schemes/benches/bdd_prepare.rs
poulpy-schemes/benches/bdd_prepare.rsRaw
1use std::hint::black_box;
2
3use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
4use poulpy_core::{
5 GLWEDecrypt, GLWEEncryptSk, ScratchTakeCore,
6 layouts::{
7 Base2K, Degree, Dnum, Dsize, GGLWEToGGSWKeyLayout, GGSWLayout, GGSWPreparedFactory, GLWEAutomorphismKeyLayout,
8 GLWELayout, GLWESecret, GLWESecretPrepared, GLWESecretPreparedFactory, GLWESwitchingKeyLayout, GLWEToLWEKeyLayout,
9 LWESecret, Rank, TorusPrecision,
10 },
11};
12
13#[cfg(all(feature = "enable-avx", target_arch = "x86_64"))]
14pub use poulpy_cpu_avx::FFT64Avx as BackendImpl;
15
16#[cfg(not(all(feature = "enable-avx", target_arch = "x86_64")))]
17pub use poulpy_cpu_ref::FFT64Ref as BackendImpl;
18
19use poulpy_hal::{
20 api::{ModuleN, ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow},
21 layouts::{Backend, Module, Scratch, ScratchOwned},
22 source::Source,
23};
24use poulpy_schemes::bin_fhe::{
25 bdd_arithmetic::{
26 BDDKey, BDDKeyEncryptSk, BDDKeyLayout, BDDKeyPrepared, BDDKeyPreparedFactory, FheUint, FheUintPrepare, FheUintPrepared,
27 },
28 blind_rotation::{
29 BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory, BlindRotationKeyInfos, BlindRotationKeyLayout, CGGI,
30 },
31 circuit_bootstrapping::{CircuitBootstrappingKeyEncryptSk, CircuitBootstrappingKeyLayout},
32};
33
34pub fn benc_bdd_prepare<BE: Backend, BRA: BlindRotationAlgo>(c: &mut Criterion, label: &str)
35where
36 Module<BE>: ModuleNew<BE>
37 + ModuleN
38 + GLWESecretPreparedFactory<BE>
39 + GLWEDecrypt<BE>
40 + CircuitBootstrappingKeyEncryptSk<BRA, BE>
41 + GGSWPreparedFactory<BE>
42 + GLWEEncryptSk<BE>
43 + BDDKeyEncryptSk<BRA, BE>
44 + BDDKeyPreparedFactory<BRA, BE>
45 + FheUintPrepare<BRA, 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!("bdd_arithmetic::{label}");
51
52 let mut group = c.benchmark_group(group_name);
53
54 struct Params {
55 name: String,
56 block_size: usize,
57 glwe_layout: GLWELayout,
58 ggsw_layout: GGSWLayout,
59 bdd_layout: BDDKeyLayout,
60 }
61
62 fn runner<BE: Backend, BRA: BlindRotationAlgo>(params: &Params) -> impl FnMut()
63 where
64 Module<BE>: ModuleNew<BE>
65 + ModuleN
66 + GLWESecretPreparedFactory<BE>
67 + GLWEDecrypt<BE>
68 + CircuitBootstrappingKeyEncryptSk<BRA, BE>
69 + GGSWPreparedFactory<BE>
70 + GLWEEncryptSk<BE>
71 + BDDKeyEncryptSk<BRA, BE>
72 + BDDKeyPreparedFactory<BRA, BE>
73 + FheUintPrepare<BRA, BE>,
74 BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>, /* TODO find a way to remove this bound or move it to CBT KEY */
75 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
76 Scratch<BE>: ScratchTakeCore<BE>,
77 {
78 // Scratch space (4MB)
79 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 25);
80
81 let n_glwe: poulpy_core::layouts::Degree = params.bdd_layout.cbt_layout.brk_layout.n_glwe();
82 let n_lwe: poulpy_core::layouts::Degree = params.bdd_layout.cbt_layout.brk_layout.n_lwe();
83 let rank: poulpy_core::layouts::Rank = params.bdd_layout.cbt_layout.brk_layout.rank;
84
85 let module: Module<BE> = Module::<BE>::new(n_glwe.as_u32() as u64);
86
87 let mut source_xs: Source = Source::new([1u8; 32]);
88 let mut source_xa: Source = Source::new([1u8; 32]);
89 let mut source_xe: Source = Source::new([1u8; 32]);
90
91 let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
92 sk_lwe.fill_binary_block(params.block_size, &mut source_xs);
93
94 let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n_glwe, rank);
95 sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
96
97 let mut sk_glwe_prepared = GLWESecretPrepared::alloc_from_infos(&module, ¶ms.glwe_layout);
98 sk_glwe_prepared.prepare(&module, &sk_glwe);
99
100 let mut bdd_key: BDDKey<Vec<u8>, BRA> = BDDKey::alloc_from_infos(¶ms.bdd_layout);
101 bdd_key.encrypt_sk(
102 &module,
103 &sk_lwe,
104 &sk_glwe,
105 &mut source_xa,
106 &mut source_xe,
107 scratch.borrow(),
108 );
109
110 let input_a = 255_u32;
111
112 let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(¶ms.glwe_layout);
113 a_enc.encrypt_sk(
114 &module,
115 input_a,
116 &sk_glwe_prepared,
117 &mut source_xa,
118 &mut source_xe,
119 scratch.borrow(),
120 );
121
122 let mut bdd_key_prepared: BDDKeyPrepared<Vec<u8>, BRA, BE> =
123 BDDKeyPrepared::alloc_from_infos(&module, ¶ms.bdd_layout);
124 bdd_key_prepared.prepare(&module, &bdd_key, scratch.borrow());
125
126 let mut a_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE> =
127 FheUintPrepared::alloc_from_infos(&module, ¶ms.ggsw_layout);
128
129 move || {
130 a_enc_prepared.prepare(&module, &a_enc, &bdd_key_prepared, scratch.borrow());
131 black_box(());
132 }
133 }
134
135 const N_GLWE: u32 = 1024;
136 const N_LWE: u32 = 679;
137 const BINARY_BLOCK_SIZE: u32 = 7;
138 const BASE2K: u32 = 15;
139 const RANK: u32 = 2;
140
141 let params: Params = Params {
142 name: format!("n_glwe={N_GLWE}"),
143 block_size: BINARY_BLOCK_SIZE as usize,
144 glwe_layout: GLWELayout {
145 n: Degree(N_GLWE),
146 base2k: Base2K(BASE2K),
147 k: TorusPrecision(2 * BASE2K),
148 rank: Rank(RANK),
149 },
150 ggsw_layout: GGSWLayout {
151 n: Degree(N_GLWE),
152 base2k: Base2K(BASE2K),
153 k: TorusPrecision(3 * BASE2K),
154 rank: Rank(RANK),
155 dnum: Dnum(3),
156 dsize: Dsize(1),
157 },
158 bdd_layout: BDDKeyLayout {
159 cbt_layout: CircuitBootstrappingKeyLayout {
160 brk_layout: BlindRotationKeyLayout {
161 n_glwe: Degree(N_GLWE),
162 n_lwe: Degree(N_LWE),
163 base2k: Base2K(BASE2K),
164 k: TorusPrecision(4 * BASE2K),
165 dnum: Dnum(4),
166 rank: Rank(RANK),
167 },
168 atk_layout: GLWEAutomorphismKeyLayout {
169 n: Degree(N_GLWE),
170 base2k: Base2K(BASE2K),
171 k: TorusPrecision(4 * BASE2K),
172 dnum: Dnum(4),
173 dsize: Dsize(1),
174 rank: Rank(RANK),
175 },
176 tsk_layout: GGLWEToGGSWKeyLayout {
177 n: Degree(N_GLWE),
178 base2k: Base2K(BASE2K),
179 k: TorusPrecision(4 * BASE2K),
180 dnum: Dnum(4),
181 dsize: Dsize(1),
182 rank: Rank(RANK),
183 },
184 },
185 ks_glwe_layout: Some(GLWESwitchingKeyLayout {
186 n: Degree(N_GLWE),
187 base2k: Base2K(BASE2K),
188 k: TorusPrecision(4 * BASE2K),
189 dnum: Dnum(4),
190 dsize: Dsize(1),
191 rank_in: Rank(RANK),
192 rank_out: Rank(1),
193 }),
194 ks_lwe_layout: GLWEToLWEKeyLayout {
195 n: Degree(N_GLWE),
196 base2k: Base2K(BASE2K),
197 k: TorusPrecision(4 * BASE2K),
198 rank_in: Rank(1),
199 dnum: Dnum(4),
200 },
201 },
202 };
203
204 let id: BenchmarkId = BenchmarkId::from_parameter(params.name.clone());
205 let mut runner = runner::<BE, BRA>(¶ms);
206 group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
207
208 group.finish();
209}
210
211fn bench_bdd_prepare_fft64(c: &mut Criterion) {
212 #[cfg(all(feature = "enable-avx", target_arch = "x86_64"))]
213 let label = "fft64_avx";
214 #[cfg(not(all(feature = "enable-avx", target_arch = "x86_64")))]
215 let label = "fft64_ref";
216 benc_bdd_prepare::<BackendImpl, CGGI>(c, label);
217}
218
219criterion_group!(benches, bench_bdd_prepare_fft64);
220criterion_main!(benches);