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