poulpy/poulpy-schemes/examples/max_array.rs
poulpy-schemes/examples/max_array.rsRaw
1use std::collections::HashMap;
2
3use poulpy_core::{
4 GLWECopy, GLWEDecrypt, GLWEEncryptSk, GLWEExternalProduct, LWEEncryptSk, ScratchTakeCore,
5 layouts::{
6 Base2K, Degree, Dnum, Dsize, GGLWEToGGSWKeyLayout, GGSWLayout, GGSWPreparedFactory, GLWEAutomorphismKeyLayout,
7 GLWELayout, GLWESecret, GLWESecretPrepared, GLWESecretPreparedFactory, GLWESwitchingKeyLayout, GLWEToLWEKeyLayout,
8 GLWEToMut, GLWEToRef, 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 BDDKey, BDDKeyEncryptSk, BDDKeyLayout, BDDKeyPrepared, BDDKeyPreparedFactory, ExecuteBDDCircuit2WTo1W, FheUint,
19 FheUintPrepare, FheUintPrepared, GLWEBlindSelection, Sltu,
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// to compute the maximum of an array of integers.
33
34fn example_max_array<BE: Backend, BRA: BlindRotationAlgo>()
35where
36 Module<BE>: ModuleNew<BE>
37 + ModuleN
38 + GLWESecretPreparedFactory<BE>
39 + GLWEExternalProduct<BE>
40 + GLWEDecrypt<BE>
41 + LWEEncryptSk<BE>
42 + GGSWPreparedFactory<BE>
43 + GLWEEncryptSk<BE>
44 + VecZnxRotateInplace<BE>
45 + BDDKeyEncryptSk<BRA, BE>
46 + BDDKeyPreparedFactory<BRA, BE>
47 + FheUintPrepare<BRA, BE>
48 + ExecuteBDDCircuit2WTo1W<BE>
49 + GLWEBlindSelection<u32, BE>,
50 BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>, // TODO find a way to remove this bound or move it to CBT KEY
51 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
52 Scratch<BE>: ScratchTakeCore<BE>,
53{
54 ////////// Parameter Selection
55 const N_GLWE: u32 = 1024;
56 const N_LWE: u32 = 567;
57 const BINARY_BLOCK_SIZE: u32 = 7;
58 const BASE2K: u32 = 17;
59 const RANK: u32 = 1;
60
61 // GLWE layout, used to generate GLWE Ciphertexts, keys, switching keys, etc
62 let glwe_layout = GLWELayout {
63 n: Degree(N_GLWE),
64 base2k: Base2K(BASE2K),
65 k: TorusPrecision(2 * BASE2K),
66 rank: Rank(RANK),
67 };
68
69 // Used to generate GGSW Ciphertexts
70 let ggsw_layout = GGSWLayout {
71 n: Degree(N_GLWE),
72 base2k: Base2K(BASE2K),
73 k: TorusPrecision(3 * BASE2K),
74 rank: Rank(RANK),
75 dnum: Dnum(3),
76 dsize: Dsize(1),
77 };
78
79 // Used to generate CBT Keys
80 let cbt_layout = CircuitBootstrappingKeyLayout {
81 brk_layout: BlindRotationKeyLayout {
82 n_glwe: Degree(N_GLWE),
83 n_lwe: Degree(N_LWE),
84 base2k: Base2K(BASE2K),
85 k: TorusPrecision(4 * BASE2K),
86 dnum: Dnum(4),
87 rank: Rank(RANK),
88 },
89 atk_layout: GLWEAutomorphismKeyLayout {
90 n: Degree(N_GLWE),
91 base2k: Base2K(BASE2K),
92 k: TorusPrecision(4 * BASE2K),
93 dnum: Dnum(4),
94 dsize: Dsize(1),
95 rank: Rank(RANK),
96 },
97 tsk_layout: GGLWEToGGSWKeyLayout {
98 n: Degree(N_GLWE),
99 base2k: Base2K(BASE2K),
100 k: TorusPrecision(4 * BASE2K),
101 dnum: Dnum(4),
102 dsize: Dsize(1),
103 rank: Rank(RANK),
104 },
105 };
106
107 // Used to generate BDD Keys, for the arithmetic operations
108 let bdd_layout = BDDKeyLayout {
109 cbt_layout: cbt_layout,
110 ks_glwe_layout: Some(GLWESwitchingKeyLayout {
111 n: Degree(N_GLWE),
112 base2k: Base2K(BASE2K),
113 k: TorusPrecision(4 * BASE2K),
114 dnum: Dnum(4),
115 dsize: Dsize(1),
116 rank_in: Rank(RANK),
117 rank_out: Rank(1),
118 }),
119 ks_lwe_layout: GLWEToLWEKeyLayout {
120 n: Degree(N_GLWE),
121 base2k: Base2K(BASE2K),
122 k: TorusPrecision(4 * BASE2K),
123 rank_in: Rank(RANK),
124 dnum: Dnum(4),
125 },
126 };
127
128 let module = Module::<BE>::new(N_GLWE as u64);
129
130 // Secret key sampling source
131 let mut source_xs: Source = Source::new([1u8; 32]);
132
133 // Public randomness sampling source
134 let mut source_xa: Source = Source::new([1u8; 32]);
135
136 // Noise sampling source
137 let mut source_xe: Source = Source::new([1u8; 32]);
138
139 // Scratch space (4MB)
140 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
141
142 ////////// Key Generation and Preparation
143 // Generating the GLWE and LWE key
144 let mut sk_glwe = GLWESecret::alloc_from_infos(&glwe_layout);
145 sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
146
147 let mut sk_lwe = LWESecret::alloc(Degree(N_LWE));
148 sk_lwe.fill_binary_block(BINARY_BLOCK_SIZE as usize, &mut source_xs);
149
150 // Preparing the private keys
151 let mut sk_glwe_prepared = GLWESecretPrepared::alloc_from_infos(&module, &glwe_layout);
152 sk_glwe_prepared.prepare(&module, &sk_glwe);
153
154 // Creating the public BDD Key
155 // This key is required to prepare all Fhe Integers for operations,
156 // and for performing the operations themselves
157 let mut bdd_key: BDDKey<Vec<u8>, BRA> = BDDKey::alloc_from_infos(&bdd_layout);
158 bdd_key.encrypt_sk(
159 &module,
160 &sk_lwe,
161 &sk_glwe,
162 &mut source_xa,
163 &mut source_xe,
164 scratch.borrow(),
165 );
166
167 ////////// Input Encryption
168 // Encrypting the inputs
169 let mut rng = rand::rng();
170 let inputs: Vec<u32> = (0..3).map(|_| rng.random_range(0..u32::MAX - 1)).collect();
171
172 let mut inputs_enc: Vec<FheUint<Vec<u8>, u32>> = Vec::new();
173 for input in &inputs {
174 let mut next_input = FheUint::alloc_from_infos(&glwe_layout);
175 next_input.encrypt_sk(
176 &module,
177 *input,
178 &sk_glwe_prepared,
179 &mut source_xa,
180 &mut source_xe,
181 scratch.borrow(),
182 );
183 inputs_enc.push(next_input);
184 }
185
186 //////// Homomorphic computation starts here ////////
187
188 // Preparing the BDD Key
189 // The BDD key must be prepared once before any operation is performed
190 let mut bdd_key_prepared: BDDKeyPrepared<Vec<u8>, BRA, BE> = BDDKeyPrepared::alloc_from_infos(&module, &bdd_layout);
191 bdd_key_prepared.prepare(&module, &bdd_key, scratch.borrow());
192
193 let mut max_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_layout);
194 max_enc.encrypt_sk(
195 &module,
196 0,
197 &sk_glwe_prepared,
198 &mut source_xa,
199 &mut source_xe,
200 scratch.borrow(),
201 );
202 // Copy of max_enc for the HashMap
203 let mut max_enc_copy: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_layout);
204
205 // Allocating the intermediate ciphertext c_enc
206 let mut compare_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_layout);
207 let mut compare_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE> = FheUintPrepared::alloc_from_infos(&module, &ggsw_layout);
208
209 for input_i in inputs_enc.iter_mut() {
210 let mut max_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE> = FheUintPrepared::alloc_from_infos(&module, &ggsw_layout);
211 max_enc_prepared.prepare(&module, &max_enc, &bdd_key_prepared, scratch.borrow());
212
213 let mut input_i_enc_prepared: FheUintPrepared<Vec<u8>, u32, BE> =
214 FheUintPrepared::alloc_from_infos(&module, &ggsw_layout);
215 input_i_enc_prepared.prepare(&module, &input_i, &bdd_key_prepared, scratch.borrow());
216
217 // b = (input_i < max)
218 compare_enc.sltu(
219 &module,
220 &input_i_enc_prepared,
221 &max_enc_prepared,
222 &bdd_key_prepared,
223 scratch.borrow(),
224 );
225
226 compare_enc_prepared.prepare(&module, &compare_enc, &bdd_key_prepared, scratch.borrow());
227
228 module.glwe_copy(&mut max_enc_copy.to_mut(), &max_enc.to_ref());
229
230 let cts = HashMap::from([(0, input_i), (1, &mut max_enc_copy)]);
231
232 <Module<BE> as GLWEBlindSelection<u32, BE>>::glwe_blind_selection(
233 &module,
234 &mut max_enc,
235 cts,
236 &compare_enc_prepared,
237 0,
238 1,
239 scratch.borrow(),
240 );
241 }
242
243 //////// Homomorphic computation ends here ////////
244
245 // Decrypting the result
246 let result_dec = max_enc.decrypt(&module, &sk_glwe_prepared, scratch.borrow());
247
248 // result = max of inputs
249 let result_correct = inputs.iter().max().unwrap();
250 println!("Result: {} == {}", result_dec, result_correct);
251}
252
253fn main() {
254 #[cfg(all(feature = "enable-avx", target_arch = "x86_64"))]
255 example_max_array::<FFT64Avx, CGGI>();
256
257 #[cfg(not(all(feature = "enable-avx", target_arch = "x86_64")))]
258 example_max_array::<FFT64Ref, CGGI>();
259}