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