poulpy/poulpy-core/src/glwe_packer.rs
poulpy-core/src/glwe_packer.rsRaw
1use std::collections::HashMap;
2
3use poulpy_hal::{
4 api::ModuleLogN,
5 layouts::{Backend, GaloisElement, Module, Scratch},
6};
7
8use crate::{
9 GLWEAdd, GLWEAutomorphism, GLWECopy, GLWENormalize, GLWERotate, GLWEShift, GLWESub, ScratchTakeCore,
10 glwe_trace::GLWETrace,
11 layouts::{GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEInfos, GLWEToMut, GLWEToRef, GetGaloisElement, LWEInfos},
12};
13
14/// [GLWEPacker] enables only the fly GLWE packing
15/// with constant memory of Log(N) ciphertexts.
16/// Main difference with usual GLWE packing is that
17/// the output is bit-reversed.
18pub struct GLWEPacker {
19 accumulators: Vec<Accumulator>,
20 log_batch: usize,
21 counter: usize,
22}
23
24/// [Accumulator] stores intermediate packing result.
25/// There are Log(N) such accumulators in a [GLWEPacker].
26struct Accumulator {
27 data: GLWE<Vec<u8>>,
28 value: bool, // Implicit flag for zero ciphertext
29 control: bool, // Can be combined with incoming value
30}
31
32impl Accumulator {
33 /// Allocates a new [Accumulator].
34 ///
35 /// #Arguments
36 ///
37 /// * `module`: static backend FFT tables.
38 /// * `base2k`: base 2 logarithm of the GLWE ciphertext in memory digit representation.
39 /// * `k`: base 2 precision of the GLWE ciphertext precision over the Torus.
40 /// * `rank`: rank of the GLWE ciphertext.
41 pub fn alloc<A>(infos: &A) -> Self
42 where
43 A: GLWEInfos,
44 {
45 Self {
46 data: GLWE::alloc_from_infos(infos),
47 value: false,
48 control: false,
49 }
50 }
51}
52
53impl GLWEPacker {
54 /// Instantiates a new [GLWEPacker].
55 ///
56 /// # Arguments
57 ///
58 /// * `log_batch`: packs coefficients which are multiples of X^{N/2^log_batch}.
59 /// i.e. with `log_batch=0` only the constant coefficient is packed
60 /// and N GLWE ciphertext can be packed. With `log_batch=2` all coefficients
61 /// which are multiples of X^{N/4} are packed. Meaning that N/4 ciphertexts
62 /// can be packed.
63 pub fn alloc<A>(infos: &A, log_batch: usize) -> Self
64 where
65 A: GLWEInfos,
66 {
67 let mut accumulators: Vec<Accumulator> = Vec::<Accumulator>::new();
68 let log_n: usize = infos.n().log2();
69 (0..log_n - log_batch).for_each(|_| accumulators.push(Accumulator::alloc(infos)));
70 GLWEPacker {
71 accumulators,
72 log_batch,
73 counter: 0,
74 }
75 }
76
77 /// Implicit reset of the internal state (to be called before a new packing procedure).
78 fn reset(&mut self) {
79 for i in 0..self.accumulators.len() {
80 self.accumulators[i].value = false;
81 self.accumulators[i].control = false;
82 }
83 self.counter = 0;
84 }
85
86 /// Number of scratch space bytes required to call [Self::add].
87 pub fn tmp_bytes<R, K, M, BE: Backend>(module: &M, res_infos: &R, key_infos: &K) -> usize
88 where
89 R: GLWEInfos,
90 K: GGLWEInfos,
91 M: GLWEPackerOps<BE>,
92 {
93 GLWE::bytes_of_from_infos(res_infos)
94 + module
95 .glwe_rsh_tmp_byte()
96 .max(module.glwe_automorphism_tmp_bytes(res_infos, res_infos, key_infos))
97 }
98
99 pub fn galois_elements<M, BE: Backend>(module: &M) -> Vec<i64>
100 where
101 M: GLWETrace<BE>,
102 {
103 module.glwe_trace_galois_elements()
104 }
105
106 /// Adds a GLWE ciphertext to the [GLWEPacker].
107 /// #Arguments
108 ///
109 /// * `module`: static backend FFT tables.
110 /// * `res`: space to append fully packed ciphertext. Only when the number
111 /// of packed ciphertexts reaches N/2^log_batch is a result written.
112 /// * `a`: ciphertext to pack. Can optionally give None to pack a 0 ciphertext.
113 /// * `auto_keys`: a [HashMap] containing the [AutomorphismKeyExec]s.
114 /// * `scratch`: scratch space of size at least [Self::tmp_bytes].
115 pub fn add<A, K, M, BE: Backend>(&mut self, module: &M, a: Option<&A>, auto_keys: &HashMap<i64, K>, scratch: &mut Scratch<BE>)
116 where
117 A: GLWEToRef + GLWEInfos,
118 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
119 M: GLWEPackerOps<BE>,
120 Scratch<BE>: ScratchTakeCore<BE>,
121 {
122 assert!(
123 (self.counter as u32) < self.accumulators[0].data.n(),
124 "Packing limit of {} reached",
125 self.accumulators[0].data.n().0 as usize >> self.log_batch
126 );
127
128 module.packer_add(self, a, self.log_batch, auto_keys, scratch);
129 self.counter += 1 << self.log_batch;
130 }
131
132 /// Flush result to`res`.
133 pub fn flush<R, M, BE: Backend>(&mut self, module: &M, res: &mut R)
134 where
135 R: GLWEToMut,
136 M: GLWEPackerOps<BE>,
137 {
138 assert!(self.counter as u32 == self.accumulators[0].data.n());
139 // Copy result GLWE into res GLWE
140 module.glwe_copy(
141 res,
142 &self.accumulators[module.log_n() - self.log_batch - 1].data,
143 );
144
145 self.reset();
146 }
147}
148
149impl<BE: Backend> GLWEPackerOps<BE> for Module<BE> where
150 Self: Sized
151 + ModuleLogN
152 + GLWEAutomorphism<BE>
153 + GaloisElement
154 + GLWERotate<BE>
155 + GLWESub
156 + GLWEShift<BE>
157 + GLWEAdd
158 + GLWENormalize<BE>
159 + GLWECopy
160 + GLWEAutomorphism<BE>
161 + GaloisElement
162 + GLWERotate<BE>
163 + GLWESub
164 + GLWEShift<BE>
165 + GLWEAdd
166 + GLWENormalize<BE>
167{
168}
169
170pub trait GLWEPackerOps<BE: Backend>
171where
172 Self: Sized
173 + ModuleLogN
174 + GLWEAutomorphism<BE>
175 + GaloisElement
176 + GLWERotate<BE>
177 + GLWESub
178 + GLWEShift<BE>
179 + GLWEAdd
180 + GLWENormalize<BE>
181 + GLWECopy
182 + GLWEAutomorphism<BE>
183 + GaloisElement
184 + GLWERotate<BE>
185 + GLWESub
186 + GLWEShift<BE>
187 + GLWEAdd
188 + GLWENormalize<BE>,
189{
190 fn packer_add<A, K>(
191 &self,
192 packer: &mut GLWEPacker,
193 a: Option<&A>,
194 i: usize,
195 auto_keys: &HashMap<i64, K>,
196 scratch: &mut Scratch<BE>,
197 ) where
198 A: GLWEToRef + GLWEInfos,
199 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
200 Scratch<BE>: ScratchTakeCore<BE>,
201 {
202 pack_core(self, a, &mut packer.accumulators, i, auto_keys, scratch)
203 }
204}
205
206fn pack_core<A, K, M, BE: Backend>(
207 module: &M,
208 a: Option<&A>,
209 accumulators: &mut [Accumulator],
210 i: usize,
211 auto_keys: &HashMap<i64, K>,
212 scratch: &mut Scratch<BE>,
213) where
214 A: GLWEToRef + GLWEInfos,
215 M: ModuleLogN
216 + GLWEAutomorphism<BE>
217 + GaloisElement
218 + GLWERotate<BE>
219 + GLWESub
220 + GLWEShift<BE>
221 + GLWEAdd
222 + GLWENormalize<BE>
223 + GLWECopy
224 + GLWEAutomorphism<BE>
225 + GaloisElement
226 + GLWERotate<BE>
227 + GLWESub
228 + GLWEShift<BE>
229 + GLWEAdd
230 + GLWENormalize<BE>,
231 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
232 Scratch<BE>: ScratchTakeCore<BE>,
233{
234 let log_n: usize = module.log_n();
235
236 if i == log_n {
237 return;
238 }
239
240 // Isolate the first accumulator
241 let (acc_prev, acc_next) = accumulators.split_at_mut(1);
242
243 // Control = true accumlator is free to overide
244 if !acc_prev[0].control {
245 let acc_mut_ref: &mut Accumulator = &mut acc_prev[0]; // from split_at_mut
246
247 // No previous value -> copies and sets flags accordingly
248 if let Some(a_ref) = a {
249 module.glwe_copy(&mut acc_mut_ref.data, a_ref);
250 acc_mut_ref.value = true
251 } else {
252 acc_mut_ref.value = false
253 }
254 acc_mut_ref.control = true; // Able to be combined on next call
255 } else {
256 // Compresses acc_prev <- combine(acc_prev, a).
257 combine(module, &mut acc_prev[0], a, i, auto_keys, scratch);
258 acc_prev[0].control = false;
259
260 // Propagates to next accumulator
261 if acc_prev[0].value {
262 pack_core(
263 module,
264 Some(&acc_prev[0].data),
265 acc_next,
266 i + 1,
267 auto_keys,
268 scratch,
269 );
270 } else {
271 pack_core(
272 module,
273 None::<&GLWE<Vec<u8>>>,
274 acc_next,
275 i + 1,
276 auto_keys,
277 scratch,
278 );
279 }
280 }
281}
282
283fn combine<B, K, M, BE: Backend>(
284 module: &M,
285 acc: &mut Accumulator,
286 b: Option<&B>,
287 i: usize,
288 auto_keys: &HashMap<i64, K>,
289 scratch: &mut Scratch<BE>,
290) where
291 B: GLWEToRef + GLWEInfos,
292 B: GLWEToRef + GLWEInfos,
293 M: ModuleLogN
294 + GLWEAutomorphism<BE>
295 + GaloisElement
296 + GLWERotate<BE>
297 + GLWESub
298 + GLWEShift<BE>
299 + GLWEAdd
300 + GLWENormalize<BE>
301 + GLWECopy
302 + GLWEAutomorphism<BE>
303 + GaloisElement
304 + GLWERotate<BE>
305 + GLWESub
306 + GLWEShift<BE>
307 + GLWEAdd
308 + GLWENormalize<BE>,
309 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
310 Scratch<BE>: ScratchTakeCore<BE>,
311{
312 let log_n: usize = acc.data.n().log2();
313 let a: &mut GLWE<Vec<u8>> = &mut acc.data;
314
315 let gal_el: i64 = if i == 0 {
316 -1
317 } else {
318 module.galois_element(1 << (i - 1))
319 };
320
321 let t: i64 = 1 << (log_n - i - 1);
322
323 // Goal is to evaluate: a = a + b*X^t + phi(a - b*X^t))
324 // We also use the identity: AUTO(a * X^t, g) = -X^t * AUTO(a, g)
325 // where t = 2^(log_n - i - 1) and g = 5^{2^(i - 1)}
326 // Different cases for wether a and/or b are zero.
327 //
328 // Implicite RSH without modulus switch, introduces extra I(X) * Q/2 on decryption.
329 // Necessary so that the scaling of the plaintext remains constant.
330 // It however is ok to do so here because coefficients are eventually
331 // either mapped to garbage or twice their value which vanishes I(X)
332 // since 2*(I(X) * Q/2) = I(X) * Q = 0 mod Q.
333 if acc.value {
334 if let Some(b) = b {
335 let (mut tmp_b, scratch_1) = scratch.take_glwe(a);
336
337 // a = a * X^-t
338 module.glwe_rotate_inplace(-t, a, scratch_1);
339
340 // tmp_b = a * X^-t - b
341 module.glwe_sub(&mut tmp_b, a, b);
342 module.glwe_rsh(1, &mut tmp_b, scratch_1);
343
344 // a = a * X^-t + b
345 module.glwe_add_inplace(a, b);
346 module.glwe_rsh(1, a, scratch_1);
347
348 module.glwe_normalize_inplace(&mut tmp_b, scratch_1);
349
350 // tmp_b = phi(a * X^-t - b)
351 if let Some(auto_key) = auto_keys.get(&gal_el) {
352 module.glwe_automorphism_inplace(&mut tmp_b, auto_key, scratch_1);
353 } else {
354 panic!("auto_key[{gal_el}] not found");
355 }
356
357 // a = a * X^-t + b - phi(a * X^-t - b)
358 module.glwe_sub_inplace(a, &tmp_b);
359 module.glwe_normalize_inplace(a, scratch_1);
360
361 // a = a + b * X^t - phi(a * X^-t - b) * X^t
362 // = a + b * X^t - phi(a * X^-t - b) * - phi(X^t)
363 // = a + b * X^t + phi(a - b * X^t)
364 module.glwe_rotate_inplace(t, a, scratch_1);
365 } else {
366 module.glwe_rsh(1, a, scratch);
367 // a = a + phi(a)
368 if let Some(auto_key) = auto_keys.get(&gal_el) {
369 module.glwe_automorphism_add_inplace(a, auto_key, scratch);
370 } else {
371 panic!("auto_key[{gal_el}] not found");
372 }
373 }
374 } else if let Some(b) = b {
375 let (mut tmp_b, scratch_1) = scratch.take_glwe(a);
376 module.glwe_rotate(t, &mut tmp_b, b);
377 module.glwe_rsh(1, &mut tmp_b, scratch_1);
378
379 // a = (b* X^t - phi(b* X^t))
380 if let Some(auto_key) = auto_keys.get(&gal_el) {
381 module.glwe_automorphism_sub_negate(a, &tmp_b, auto_key, scratch_1);
382 } else {
383 panic!("auto_key[{gal_el}] not found");
384 }
385
386 acc.value = true;
387 }
388}