poulpy/poulpy-schemes/examples/bdd_arithmetic.rs
poulpy-schemes/examples/bdd_arithmetic.rsRaw
1use std::collections::HashMap;
2
3use poulpy_core::{
4 GLWEDecrypt, GLWEEncryptSk, GLWEExternalProduct, LWEEncryptSk, ScratchTakeCore,
5 layouts::{
6 Base2K, Degree, Dnum, Dsize, GGLWEToGGSWKeyLayout, GGSWLayout, GGSWPreparedFactory, GLWEAutomorphismKeyLayout,
7 GLWELayout, GLWESecret, GLWESecretPrepared, GLWESecretPreparedFactory, GLWESwitchingKeyLayout, GLWEToLWEKeyLayout,
8 LWESecret, Rank, TorusPrecision,
9 },
10};
11use poulpy_hal::{
12 api::{ModuleN, ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxRotateInplace},
13 layouts::{Backend, Module, Scratch, ScratchOwned},
14 source::Source,
15};
16use poulpy_schemes::bin_fhe::{
17 bdd_arithmetic::{
18 Add, BDDKey, BDDKeyEncryptSk, BDDKeyLayout, BDDKeyPrepared, BDDKeyPreparedFactory, ExecuteBDDCircuit2WTo1W, FheUint,
19 FheUintPrepare, FheUintPrepared, GLWEBlindSelection, Xor,
20 },
21 blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory, BlindRotationKeyLayout, CGGI},
22 circuit_bootstrapping::CircuitBootstrappingKeyLayout,
23};
24use rand::Rng;
25
26#[cfg(all(feature = "enable-avx", target_arch = "x86_64"))]
27use poulpy_cpu_avx::FFT64Avx;
28#[cfg(not(all(feature = "enable-avx", target_arch = "x86_64")))]
29use poulpy_cpu_ref::FFT64Ref;
30
31// This example demonstrates and end-to-end example usage of the BDD arithmetic API
32// It includes all steps including:
33//
34// - Parameter Selection
35// - Key Generation
36// - Input Encryption
37//
38// - Key preparation
39// - Input Preparation
40// - Operation Execution
41//
42// - Result Decryption
43//
44
45fn example_bdd_arithmetic<BE: Backend, BRA: BlindRotationAlgo>()
46where
47 Module<BE>: ModuleNew<BE>
48 + ModuleN
49 + GLWESecretPreparedFactory<BE>
50 + GLWEExternalProduct<BE>
51 + GLWEDecrypt<BE>
52 + LWEEncryptSk<BE>
53 + GGSWPreparedFactory<BE>
54 + GLWEEncryptSk<BE>
55 + VecZnxRotateInplace<BE>
56 + BDDKeyEncryptSk<BRA, BE>
57 + BDDKeyPreparedFactory<BRA, BE>
58 + FheUintPrepare<BRA, BE>
59 + ExecuteBDDCircuit2WTo1W<BE>
60 + GLWEBlindSelection<u32, BE>,
61 BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>, // TODO find a way to remove this bound or move it to CBT KEY
62 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
63 Scratch<BE>: ScratchTakeCore<BE>,
64{
65 ////////// Parameter Selection
66 const N_GLWE: u32 = 1024;
67 const N_LWE: u32 = 567;
68 const BINARY_BLOCK_SIZE: u32 = 7;
69 const BASE2K: u32 = 17;
70 const RANK: u32 = 1;
71
72 // GLWE layout, used to generate GLWE Ciphertexts, keys, switching keys, etc
73 let glwe_layout = GLWELayout {
74 n: Degree(N_GLWE),
75 base2k: Base2K(BASE2K),
76 k: TorusPrecision(2 * BASE2K),
77 rank: Rank(RANK),
78 };
79
80 // Used to generate GGSW Ciphertexts
81 let ggsw_layout = GGSWLayout {
82 n: Degree(N_GLWE),
83 base2k: Base2K(BASE2K),
84 k: TorusPrecision(3 * BASE2K),
85 rank: Rank(RANK),
86 dnum: Dnum(3),
87 dsize: Dsize(1),
88 };
89
90 // Used to generate BDD Keys, for the arithmetic operations
91 let bdd_layout = BDDKeyLayout {
92 cbt_layout: CircuitBootstrappingKeyLayout {
93 brk_layout: BlindRotationKeyLayout {
94 n_glwe: Degree(N_GLWE),
95 n_lwe: Degree(N_LWE),
96 base2k: Base2K(BASE2K),
97 k: TorusPrecision(4 * BASE2K),
98 dnum: Dnum(4),
99 rank: Rank(RANK),
100 },
101 atk_layout: GLWEAutomorphismKeyLayout {
102 n: Degree(N_GLWE),
103 base2k: Base2K(BASE2K),
104 k: TorusPrecision(4 * BASE2K),
105 dnum: Dnum(4),
106 dsize: Dsize(1),
107 rank: Rank(RANK),
108 },
109 tsk_layout: GGLWEToGGSWKeyLayout {
110 n: Degree(N_GLWE),
111 base2k: Base2K(BASE2K),
112 k: TorusPrecision(4 * BASE2K),
113 dnum: Dnum(4),
114 dsize: Dsize(1),
115 rank: Rank(RANK),
116 },
117 },
118 ks_glwe_layout: Some(GLWESwitchingKeyLayout {
119 n: Degree(N_GLWE),
120 base2k: Base2K(BASE2K),
121 k: TorusPrecision(4 * BASE2K),
122 dnum: Dnum(4),
123 dsize: Dsize(1),
124 rank_in: Rank(RANK),
125 rank_out: Rank(1),
126 }),
127 ks_lwe_layout: GLWEToLWEKeyLayout {
128 n: Degree(N_GLWE),
129 base2k: Base2K(BASE2K),
130 k: TorusPrecision(4 * BASE2K),
131 rank_in: Rank(1),
132 dnum: Dnum(4),
133 },
134 };
135
136 let module = Module::<BE>::new(N_GLWE as u64);
137
138 // Secret key sampling source
139 let mut source_xs: Source = Source::new([1u8; 32]);
140
141 // Public randomness sampling source
142 let mut source_xa: Source = Source::new([1u8; 32]);
143
144 // Noise sampling source
145 let mut source_xe: Source = Source::new([1u8; 32]);
146
147 // Scratch space (4MB)
148 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
149
150 ////////// Key Generation and Preparation
151 // Generating the GLWE and LWE key
152 let mut sk_glwe = GLWESecret::alloc_from_infos(&glwe_layout);
153 sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
154
155 let mut sk_lwe = LWESecret::alloc(bdd_layout.cbt_layout.brk_layout.n_lwe);
156 sk_lwe.fill_binary_block(BINARY_BLOCK_SIZE as usize, &mut source_xs);
157
158 // Preparing the private keys
159 let mut sk_glwe_prepared = GLWESecretPrepared::alloc_from_infos(&module, &glwe_layout);
160 sk_glwe_prepared.prepare(&module, &sk_glwe);
161
162 // Creating the public BDD Key
163 // This key is required to prepare all Fhe Integers for operations,
164 // and for performing the operations themselves
165 let mut bdd_key: BDDKey<Vec<u8>, BRA> = BDDKey::alloc_from_infos(&bdd_layout);
166 bdd_key.encrypt_sk(
167 &module,
168 &sk_lwe,
169 &sk_glwe,
170 &mut source_xa,
171 &mut source_xe,
172 scratch.borrow(),
173 );
174
175 ////////// Input Encryption
176 // Encrypting the inputs
177 let input_a = 255_u32;
178 let input_b = 30_u32;
179
180 let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_layout);
181 a_enc.encrypt_sk(
182 &module,
183 input_a,
184 &sk_glwe_prepared,
185 &mut source_xa,
186 &mut source_xe,
187 scratch.borrow(),
188 );
189
190 let mut b_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_layout);
191 b_enc.encrypt_sk(
192 &module,
193 input_b,
194 &sk_glwe_prepared,
195 &mut source_xa,
196 &mut source_xe,
197 scratch.borrow(),
198 );
199
200 //////// Homomorphic computation starts here ////////
201
202 // Preparing the BDD Key
203 // The BDD key must be prepared once before any operation is performed
204 let mut bdd_key_prepared: BDDKeyPrepared<Vec<u8>, BRA, BE> = BDDKeyPrepared::alloc_from_infos(&module, &bdd_layout);
205 bdd_key_prepared.prepare(&module, &bdd_key, scratch.borrow());
206
207 // Input Preparation
208 // Before each operation, the inputs to that operation must be prepared
209 // Preparation extracts each bit of the integer into a seperate GLWE ciphertext and bootstraps it into a GGSW ciphertext
210 let mut a_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE> = FheUintPrepared::alloc_from_infos(&module, &ggsw_layout);
211 a_enc_prepared.prepare(&module, &a_enc, &bdd_key_prepared, scratch.borrow());
212
213 let mut b_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE> = FheUintPrepared::alloc_from_infos(&module, &ggsw_layout);
214 b_enc_prepared.prepare(&module, &b_enc, &bdd_key_prepared, scratch.borrow());
215
216 // Allocating the intermediate ciphertext c_enc
217 let mut c_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_layout);
218
219 // Performing the operation
220 c_enc.add(
221 &module,
222 &a_enc_prepared,
223 &b_enc_prepared,
224 &bdd_key_prepared,
225 scratch.borrow(),
226 );
227
228 // Preparing the intermediate result ciphertext, c_enc, for the next operation
229 let mut c_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE> = FheUintPrepared::alloc_from_infos(&module, &ggsw_layout);
230 c_enc_prepared.prepare(&module, &c_enc, &bdd_key_prepared, scratch.borrow());
231
232 // Creating the output ciphertext d_enc
233 let mut selected_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_layout);
234 selected_enc.xor(
235 &module,
236 &c_enc_prepared,
237 &a_enc_prepared,
238 &bdd_key_prepared,
239 scratch.borrow(),
240 );
241
242 //////// Homomorphic computation ends here ////////
243
244 // Decrypting the result
245 let d_dec = selected_enc.decrypt(&module, &sk_glwe_prepared, scratch.borrow());
246
247 // d = (a + b) ^ a
248 let d_correct = (input_a.wrapping_add(input_b)) ^ input_a;
249 println!("Result: {} == {}", d_dec, d_correct);
250
251 // List of available operations are:
252 // - add: addition
253 // - sub: subtraction
254 // - sll: left shift logical
255 // - sra: right shift arithmetic
256 // - srl: right shift logical
257 // - slt: less than
258 // - sltu: less than unsigned
259 // - and: bitwise and
260 // - or: bitwise or
261 // - xor: bitwise xor
262
263 ///////////////////////////// GLWE Blind Selection
264 // This example demonstrates the use of the GLWE Blind Selection operation
265 // It can choose between any number of encrypted fheuint inputs
266 // using an encrypted fheuint selector
267
268 let log_2_number_of_inputs: usize = 5;
269 let number_of_inputs: usize = 1 << log_2_number_of_inputs;
270 let inputs_a_vec: Vec<u32> = (0..number_of_inputs)
271 .map(|_| rand::rng().random_range(0..u32::MAX - 1))
272 .collect();
273 let input_selector: u32 = rand::rng().random_range(0..number_of_inputs as u32);
274
275 let mut inputs_a_enc_vec: Vec<FheUint<Vec<u8>, u32>> = Vec::new();
276 for input in &inputs_a_vec {
277 let mut next_input: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_layout);
278 next_input.encrypt_sk(
279 &module,
280 *input,
281 &sk_glwe_prepared,
282 &mut source_xa,
283 &mut source_xe,
284 scratch.borrow(),
285 );
286 inputs_a_enc_vec.push(next_input);
287 }
288
289 let mut inputs_a_enc_vec_map: HashMap<usize, &mut FheUint<Vec<u8>, u32>> = HashMap::new();
290 for (i, input) in inputs_a_enc_vec.iter_mut().enumerate() {
291 inputs_a_enc_vec_map.insert(i, input);
292 }
293
294 let mut input_selector_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_layout);
295 input_selector_enc.encrypt_sk(
296 &module,
297 input_selector,
298 &sk_glwe_prepared,
299 &mut source_xa,
300 &mut source_xe,
301 scratch.borrow(),
302 );
303 let mut input_selector_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE> =
304 FheUintPrepared::alloc_from_infos(&module, &ggsw_layout);
305 input_selector_enc_prepared.prepare(
306 &module,
307 &input_selector_enc,
308 &bdd_key_prepared,
309 scratch.borrow(),
310 );
311
312 module.glwe_blind_selection(
313 &mut selected_enc,
314 inputs_a_enc_vec_map,
315 &input_selector_enc_prepared,
316 0,
317 log_2_number_of_inputs,
318 scratch.borrow(),
319 );
320
321 let selected_dec = selected_enc.decrypt(&module, &sk_glwe_prepared, scratch.borrow());
322 let selected_correct = inputs_a_vec[input_selector as usize];
323 println!("Result: {} == {}", selected_dec, selected_correct);
324}
325
326fn main() {
327 #[cfg(all(feature = "enable-avx", target_arch = "x86_64"))]
328 example_bdd_arithmetic::<FFT64Avx, CGGI>();
329
330 #[cfg(not(all(feature = "enable-avx", target_arch = "x86_64")))]
331 example_bdd_arithmetic::<FFT64Ref, CGGI>();
332}