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