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(&module, &sk_lwe, &sk_glwe, &mut source_xa, &mut source_xe, scratch.borrow());
166
167 ////////// Input Encryption
168 // Encrypting the inputs
169 let input_a = 255_u32;
170 let input_b = 30_u32;
171
172 let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_layout);
173 a_enc.encrypt_sk(
174 &module,
175 input_a,
176 &sk_glwe_prepared,
177 &mut source_xa,
178 &mut source_xe,
179 scratch.borrow(),
180 );
181
182 let mut b_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_layout);
183 b_enc.encrypt_sk(
184 &module,
185 input_b,
186 &sk_glwe_prepared,
187 &mut source_xa,
188 &mut source_xe,
189 scratch.borrow(),
190 );
191
192 //////// Homomorphic computation starts here ////////
193
194 // Preparing the BDD Key
195 // The BDD key must be prepared once before any operation is performed
196 let mut bdd_key_prepared: BDDKeyPrepared<Vec<u8>, BRA, BE> = BDDKeyPrepared::alloc_from_infos(&module, &bdd_layout);
197 bdd_key_prepared.prepare(&module, &bdd_key, scratch.borrow());
198
199 // Input Preparation
200 // Before each operation, the inputs to that operation must be prepared
201 // Preparation extracts each bit of the integer into a seperate GLWE ciphertext and bootstraps it into a GGSW ciphertext
202 let mut a_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE> = FheUintPrepared::alloc_from_infos(&module, &ggsw_layout);
203 a_enc_prepared.prepare(&module, &a_enc, &bdd_key_prepared, scratch.borrow());
204
205 let mut b_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE> = FheUintPrepared::alloc_from_infos(&module, &ggsw_layout);
206 b_enc_prepared.prepare(&module, &b_enc, &bdd_key_prepared, scratch.borrow());
207
208 // Allocating the intermediate ciphertext c_enc
209 let mut c_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_layout);
210
211 // Performing the operation
212 c_enc.add(&module, &a_enc_prepared, &b_enc_prepared, &bdd_key_prepared, scratch.borrow());
213
214 // Preparing the intermediate result ciphertext, c_enc, for the next operation
215 let mut c_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE> = FheUintPrepared::alloc_from_infos(&module, &ggsw_layout);
216 c_enc_prepared.prepare(&module, &c_enc, &bdd_key_prepared, scratch.borrow());
217
218 // Creating the output ciphertext d_enc
219 let mut selected_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_layout);
220 selected_enc.xor(&module, &c_enc_prepared, &a_enc_prepared, &bdd_key_prepared, scratch.borrow());
221
222 //////// Homomorphic computation ends here ////////
223
224 // Decrypting the result
225 let d_dec = selected_enc.decrypt(&module, &sk_glwe_prepared, scratch.borrow());
226
227 // d = (a + b) ^ a
228 let d_correct = (input_a.wrapping_add(input_b)) ^ input_a;
229 println!("Result: {} == {}", d_dec, d_correct);
230
231 // List of available operations are:
232 // - add: addition
233 // - sub: subtraction
234 // - sll: left shift logical
235 // - sra: right shift arithmetic
236 // - srl: right shift logical
237 // - slt: less than
238 // - sltu: less than unsigned
239 // - and: bitwise and
240 // - or: bitwise or
241 // - xor: bitwise xor
242
243 ///////////////////////////// GLWE Blind Selection
244 // This example demonstrates the use of the GLWE Blind Selection operation
245 // It can choose between any number of encrypted fheuint inputs
246 // using an encrypted fheuint selector
247
248 let log_2_number_of_inputs: usize = 5;
249 let number_of_inputs: usize = 1 << log_2_number_of_inputs;
250 let inputs_a_vec: Vec<u32> = (0..number_of_inputs)
251 .map(|_| rand::rng().random_range(0..u32::MAX - 1))
252 .collect();
253 let input_selector: u32 = rand::rng().random_range(0..number_of_inputs as u32);
254
255 let mut inputs_a_enc_vec: Vec<FheUint<Vec<u8>, u32>> = Vec::new();
256 for input in &inputs_a_vec {
257 let mut next_input: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_layout);
258 next_input.encrypt_sk(
259 &module,
260 *input,
261 &sk_glwe_prepared,
262 &mut source_xa,
263 &mut source_xe,
264 scratch.borrow(),
265 );
266 inputs_a_enc_vec.push(next_input);
267 }
268
269 let mut inputs_a_enc_vec_map: HashMap<usize, &mut FheUint<Vec<u8>, u32>> = HashMap::new();
270 for (i, input) in inputs_a_enc_vec.iter_mut().enumerate() {
271 inputs_a_enc_vec_map.insert(i, input);
272 }
273
274 let mut input_selector_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_layout);
275 input_selector_enc.encrypt_sk(
276 &module,
277 input_selector,
278 &sk_glwe_prepared,
279 &mut source_xa,
280 &mut source_xe,
281 scratch.borrow(),
282 );
283 let mut input_selector_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE> =
284 FheUintPrepared::alloc_from_infos(&module, &ggsw_layout);
285 input_selector_enc_prepared.prepare(&module, &input_selector_enc, &bdd_key_prepared, scratch.borrow());
286
287 module.glwe_blind_selection(
288 &mut selected_enc,
289 inputs_a_enc_vec_map,
290 &input_selector_enc_prepared,
291 0,
292 log_2_number_of_inputs,
293 scratch.borrow(),
294 );
295
296 let selected_dec = selected_enc.decrypt(&module, &sk_glwe_prepared, scratch.borrow());
297 let selected_correct = inputs_a_vec[input_selector as usize];
298 println!("Result: {} == {}", selected_dec, selected_correct);
299}
300
301fn main() {
302 #[cfg(all(feature = "enable-avx", target_arch = "x86_64"))]
303 example_bdd_arithmetic::<FFT64Avx, CGGI>();
304
305 #[cfg(not(all(feature = "enable-avx", target_arch = "x86_64")))]
306 example_bdd_arithmetic::<FFT64Ref, CGGI>();
307}