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::{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: GLWEPacking<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: GLWEPacking<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 pack_core(
129 module,
130 a,
131 &mut self.accumulators,
132 self.log_batch,
133 auto_keys,
134 scratch,
135 );
136 self.counter += 1 << self.log_batch;
137 }
138
139 /// Flush result to`res`.
140 pub fn flush<R, M, BE: Backend>(&mut self, module: &M, res: &mut R)
141 where
142 R: GLWEToMut,
143 M: GLWEPacking<BE>,
144 {
145 assert!(self.counter as u32 == self.accumulators[0].data.n());
146 // Copy result GLWE into res GLWE
147 module.glwe_copy(
148 res,
149 &self.accumulators[module.log_n() - self.log_batch - 1].data,
150 );
151
152 self.reset();
153 }
154}
155
156impl<BE: Backend> GLWEPacking<BE> for Module<BE> where
157 Self: GLWEAutomorphism<BE>
158 + GaloisElement
159 + ModuleLogN
160 + GLWERotate<BE>
161 + GLWESub
162 + GLWEShift<BE>
163 + GLWEAdd
164 + GLWENormalize<BE>
165 + GLWECopy
166{
167}
168
169pub trait GLWEPacking<BE: Backend>
170where
171 Self: GLWEAutomorphism<BE>
172 + GaloisElement
173 + ModuleLogN
174 + GLWERotate<BE>
175 + GLWESub
176 + GLWEShift<BE>
177 + GLWEAdd
178 + GLWENormalize<BE>
179 + GLWECopy,
180{
181 /// Packs [x_0: GLWE(m_0), x_1: GLWE(m_1), ..., x_i: GLWE(m_i)]
182 /// to [0: GLWE(m_0 * X^x_0 + m_1 * X^x_1 + ... + m_i * X^x_i)]
183 fn glwe_pack<R, K>(
184 &self,
185 cts: &mut HashMap<usize, &mut R>,
186 log_gap_out: usize,
187 keys: &HashMap<i64, K>,
188 scratch: &mut Scratch<BE>,
189 ) where
190 R: GLWEToMut + GLWEToRef + GLWEInfos,
191 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
192 Scratch<BE>: ScratchTakeCore<BE>,
193 {
194 #[cfg(debug_assertions)]
195 {
196 assert!(*cts.keys().max().unwrap() < self.n())
197 }
198
199 let log_n: usize = self.log_n();
200
201 for i in 0..(log_n - log_gap_out) {
202 let t: usize = (1 << log_n).min(1 << (log_n - 1 - i));
203
204 let key: &K = if i == 0 {
205 keys.get(&-1).unwrap()
206 } else {
207 keys.get(&self.galois_element(1 << (i - 1))).unwrap()
208 };
209
210 for j in 0..t {
211 let mut a: Option<&mut R> = cts.remove(&j);
212 let mut b: Option<&mut R> = cts.remove(&(j + t));
213
214 pack_internal(self, &mut a, &mut b, i, key, scratch);
215
216 if let Some(a) = a {
217 cts.insert(j, a);
218 } else if let Some(b) = b {
219 cts.insert(j, b);
220 }
221 }
222 }
223 }
224}
225
226fn pack_core<A, K, M, BE: Backend>(
227 module: &M,
228 a: Option<&A>,
229 accumulators: &mut [Accumulator],
230 i: usize,
231 auto_keys: &HashMap<i64, K>,
232 scratch: &mut Scratch<BE>,
233) where
234 A: GLWEToRef + GLWEInfos,
235 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
236 M: ModuleLogN
237 + GLWEAutomorphism<BE>
238 + GaloisElement
239 + GLWERotate<BE>
240 + GLWESub
241 + GLWEShift<BE>
242 + GLWEAdd
243 + GLWENormalize<BE>
244 + GLWECopy,
245 Scratch<BE>: ScratchTakeCore<BE>,
246{
247 let log_n: usize = module.log_n();
248
249 if i == log_n {
250 return;
251 }
252
253 // Isolate the first accumulator
254 let (acc_prev, acc_next) = accumulators.split_at_mut(1);
255
256 // Control = true accumlator is free to overide
257 if !acc_prev[0].control {
258 let acc_mut_ref: &mut Accumulator = &mut acc_prev[0]; // from split_at_mut
259
260 // No previous value -> copies and sets flags accordingly
261 if let Some(a_ref) = a {
262 module.glwe_copy(&mut acc_mut_ref.data, a_ref);
263 acc_mut_ref.value = true
264 } else {
265 acc_mut_ref.value = false
266 }
267 acc_mut_ref.control = true; // Able to be combined on next call
268 } else {
269 // Compresses acc_prev <- combine(acc_prev, a).
270 combine(module, &mut acc_prev[0], a, i, auto_keys, scratch);
271 acc_prev[0].control = false;
272
273 // Propagates to next accumulator
274 if acc_prev[0].value {
275 pack_core(
276 module,
277 Some(&acc_prev[0].data),
278 acc_next,
279 i + 1,
280 auto_keys,
281 scratch,
282 );
283 } else {
284 pack_core(
285 module,
286 None::<&GLWE<Vec<u8>>>,
287 acc_next,
288 i + 1,
289 auto_keys,
290 scratch,
291 );
292 }
293 }
294}
295
296/// [combine] merges two ciphertexts together.
297fn combine<B, M, K, BE: Backend>(
298 module: &M,
299 acc: &mut Accumulator,
300 b: Option<&B>,
301 i: usize,
302 auto_keys: &HashMap<i64, K>,
303 scratch: &mut Scratch<BE>,
304) where
305 B: GLWEToRef + GLWEInfos,
306 M: GLWEAutomorphism<BE> + GaloisElement + GLWERotate<BE> + GLWESub + GLWEShift<BE> + GLWEAdd + GLWENormalize<BE>,
307 B: GLWEToRef + GLWEInfos,
308 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
309 Scratch<BE>: ScratchTakeCore<BE>,
310{
311 let log_n: usize = acc.data.n().log2();
312 let a: &mut GLWE<Vec<u8>> = &mut acc.data;
313
314 let gal_el: i64 = if i == 0 {
315 -1
316 } else {
317 module.galois_element(1 << (i - 1))
318 };
319
320 let t: i64 = 1 << (log_n - i - 1);
321
322 // Goal is to evaluate: a = a + b*X^t + phi(a - b*X^t))
323 // We also use the identity: AUTO(a * X^t, g) = -X^t * AUTO(a, g)
324 // where t = 2^(log_n - i - 1) and g = 5^{2^(i - 1)}
325 // Different cases for wether a and/or b are zero.
326 //
327 // Implicite RSH without modulus switch, introduces extra I(X) * Q/2 on decryption.
328 // Necessary so that the scaling of the plaintext remains constant.
329 // It however is ok to do so here because coefficients are eventually
330 // either mapped to garbage or twice their value which vanishes I(X)
331 // since 2*(I(X) * Q/2) = I(X) * Q = 0 mod Q.
332 if acc.value {
333 if let Some(b) = b {
334 let (mut tmp_b, scratch_1) = scratch.take_glwe(module, a);
335
336 // a = a * X^-t
337 module.glwe_rotate_inplace(-t, a, scratch_1);
338
339 // tmp_b = a * X^-t - b
340 module.glwe_sub(&mut tmp_b, a, b);
341 module.glwe_rsh(1, &mut tmp_b, scratch_1);
342
343 // a = a * X^-t + b
344 module.glwe_add_inplace(a, b);
345 module.glwe_rsh(1, a, scratch_1);
346
347 module.glwe_normalize_inplace(&mut tmp_b, scratch_1);
348
349 // tmp_b = phi(a * X^-t - b)
350 if let Some(auto_key) = auto_keys.get(&gal_el) {
351 module.glwe_automorphism_inplace(&mut tmp_b, auto_key, scratch_1);
352 } else {
353 panic!("auto_key[{gal_el}] not found");
354 }
355
356 // a = a * X^-t + b - phi(a * X^-t - b)
357 module.glwe_sub_inplace(a, &tmp_b);
358 module.glwe_normalize_inplace(a, scratch_1);
359
360 // a = a + b * X^t - phi(a * X^-t - b) * X^t
361 // = a + b * X^t - phi(a * X^-t - b) * - phi(X^t)
362 // = a + b * X^t + phi(a - b * X^t)
363 module.glwe_rotate_inplace(t, a, scratch_1);
364 } else {
365 module.glwe_rsh(1, a, scratch);
366 // a = a + phi(a)
367 if let Some(auto_key) = auto_keys.get(&gal_el) {
368 module.glwe_automorphism_add_inplace(a, auto_key, scratch);
369 } else {
370 panic!("auto_key[{gal_el}] not found");
371 }
372 }
373 } else if let Some(b) = b {
374 let (mut tmp_b, scratch_1) = scratch.take_glwe(module, a);
375 module.glwe_rotate(t, &mut tmp_b, b);
376 module.glwe_rsh(1, &mut tmp_b, scratch_1);
377
378 // a = (b* X^t - phi(b* X^t))
379 if let Some(auto_key) = auto_keys.get(&gal_el) {
380 module.glwe_automorphism_sub_negate(a, &tmp_b, auto_key, scratch_1);
381 } else {
382 panic!("auto_key[{gal_el}] not found");
383 }
384
385 acc.value = true;
386 }
387}
388
389#[allow(clippy::too_many_arguments)]
390fn pack_internal<M, A, B, K, BE: Backend>(
391 module: &M,
392 a: &mut Option<&mut A>,
393 b: &mut Option<&mut B>,
394 i: usize,
395 auto_key: &K,
396 scratch: &mut Scratch<BE>,
397) where
398 M: GLWEAutomorphism<BE> + GLWERotate<BE> + GLWESub + GLWEShift<BE> + GLWEAdd + GLWENormalize<BE>,
399 A: GLWEToMut + GLWEToRef + GLWEInfos,
400 B: GLWEToMut + GLWEToRef + GLWEInfos,
401 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
402 Scratch<BE>: ScratchTakeCore<BE>,
403{
404 // Goal is to evaluate: a = a + b*X^t + phi(a - b*X^t))
405 // We also use the identity: AUTO(a * X^t, g) = -X^t * AUTO(a, g)
406 // where t = 2^(log_n - i - 1) and g = 5^{2^(i - 1)}
407 // Different cases for wether a and/or b are zero.
408 //
409 // Implicite RSH without modulus switch, introduces extra I(X) * Q/2 on decryption.
410 // Necessary so that the scaling of the plaintext remains constant.
411 // It however is ok to do so here because coefficients are eventually
412 // either mapped to garbage or twice their value which vanishes I(X)
413 // since 2*(I(X) * Q/2) = I(X) * Q = 0 mod Q.
414 if let Some(a) = a.as_deref_mut() {
415 let t: i64 = 1 << (a.n().log2() - i - 1);
416
417 if let Some(b) = b.as_deref_mut() {
418 let (mut tmp_b, scratch_1) = scratch.take_glwe(module, a);
419
420 // a = a * X^-t
421 module.glwe_rotate_inplace(-t, a, scratch_1);
422
423 // tmp_b = a * X^-t - b
424 module.glwe_sub(&mut tmp_b, a, b);
425 module.glwe_rsh(1, &mut tmp_b, scratch_1);
426
427 // a = a * X^-t + b
428 module.glwe_add_inplace(a, b);
429 module.glwe_rsh(1, a, scratch_1);
430
431 module.glwe_normalize_inplace(&mut tmp_b, scratch_1);
432
433 // tmp_b = phi(a * X^-t - b)
434 module.glwe_automorphism_inplace(&mut tmp_b, auto_key, scratch_1);
435
436 // a = a * X^-t + b - phi(a * X^-t - b)
437 module.glwe_sub_inplace(a, &tmp_b);
438 module.glwe_normalize_inplace(a, scratch_1);
439
440 // a = a + b * X^t - phi(a * X^-t - b) * X^t
441 // = a + b * X^t - phi(a * X^-t - b) * - phi(X^t)
442 // = a + b * X^t + phi(a - b * X^t)
443 module.glwe_rotate_inplace(t, a, scratch_1);
444 } else {
445 module.glwe_rsh(1, a, scratch);
446 // a = a + phi(a)
447 module.glwe_automorphism_add_inplace(a, auto_key, scratch);
448 }
449 } else if let Some(b) = b.as_deref_mut() {
450 let t: i64 = 1 << (b.n().log2() - i - 1);
451
452 let (mut tmp_b, scratch_1) = scratch.take_glwe(module, b);
453 module.glwe_rotate(t, &mut tmp_b, b);
454 module.glwe_rsh(1, &mut tmp_b, scratch_1);
455
456 // a = (b* X^t - phi(b* X^t))
457 module.glwe_automorphism_sub_negate(b, &tmp_b, auto_key, scratch_1);
458 }
459}