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