poulpy/poulpy-core/src/glwe_packing.rs
poulpy-core/src/glwe_packing.rsRaw
1use std::collections::HashMap;
2
3use poulpy_hal::{
4 api::{
5 ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxAddInplace, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace,
6 VecZnxBigAutomorphismInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallNegateInplace, VecZnxCopy,
7 VecZnxDftAllocBytes, VecZnxDftApply, VecZnxDftCopy, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNegateInplace,
8 VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotate, VecZnxRotateInplace, VecZnxRshInplace,
9 VecZnxSub, VecZnxSubInplace, VecZnxSwitchRing, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
10 },
11 layouts::{Backend, DataMut, DataRef, Module, Scratch},
12};
13
14use crate::{
15 GLWEOperations, TakeGLWECt,
16 layouts::{GGLWEInfos, GLWECiphertext, GLWEInfos, LWEInfos, prepared::AutomorphismKeyPrepared},
17};
18
19/// [GLWEPacker] enables only the fly GLWE packing
20/// with constant memory of Log(N) ciphertexts.
21/// Main difference with usual GLWE packing is that
22/// the output is bit-reversed.
23pub struct GLWEPacker {
24 accumulators: Vec<Accumulator>,
25 log_batch: usize,
26 counter: usize,
27}
28
29/// [Accumulator] stores intermediate packing result.
30/// There are Log(N) such accumulators in a [GLWEPacker].
31struct Accumulator {
32 data: GLWECiphertext<Vec<u8>>,
33 value: bool, // Implicit flag for zero ciphertext
34 control: bool, // Can be combined with incoming value
35}
36
37impl Accumulator {
38 /// Allocates a new [Accumulator].
39 ///
40 /// #Arguments
41 ///
42 /// * `module`: static backend FFT tables.
43 /// * `base2k`: base 2 logarithm of the GLWE ciphertext in memory digit representation.
44 /// * `k`: base 2 precision of the GLWE ciphertext precision over the Torus.
45 /// * `rank`: rank of the GLWE ciphertext.
46 pub fn alloc<A>(infos: &A) -> Self
47 where
48 A: GLWEInfos,
49 {
50 Self {
51 data: GLWECiphertext::alloc(infos),
52 value: false,
53 control: false,
54 }
55 }
56}
57
58impl GLWEPacker {
59 /// Instantiates a new [GLWEPacker].
60 ///
61 /// # Arguments
62 ///
63 /// * `module`: static backend FFT tables.
64 /// * `log_batch`: packs coefficients which are multiples of X^{N/2^log_batch}.
65 /// i.e. with `log_batch=0` only the constant coefficient is packed
66 /// and N GLWE ciphertext can be packed. With `log_batch=2` all coefficients
67 /// which are multiples of X^{N/4} are packed. Meaning that N/4 ciphertexts
68 /// can be packed.
69 pub fn new<A>(infos: &A, log_batch: usize) -> Self
70 where
71 A: GLWEInfos,
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(infos)));
76 Self {
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 scratch_space<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_infos: &KEY) -> usize
94 where
95 OUT: GLWEInfos,
96 KEY: GGLWEInfos,
97 Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
98 {
99 pack_core_scratch_space(module, out_infos, key_infos)
100 }
101
102 pub fn galois_elements<B: Backend>(module: &Module<B>) -> Vec<i64> {
103 GLWECiphertext::trace_galois_elements(module)
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::scratch_space].
115 pub fn add<DataA: DataRef, DataAK: DataRef, B: Backend>(
116 &mut self,
117 module: &Module<B>,
118 a: Option<&GLWECiphertext<DataA>>,
119 auto_keys: &HashMap<i64, AutomorphismKeyPrepared<DataAK, B>>,
120 scratch: &mut Scratch<B>,
121 ) where
122 Module<B>: VecZnxDftAllocBytes
123 + VmpApplyDftToDftTmpBytes
124 + VecZnxBigNormalizeTmpBytes
125 + VmpApplyDftToDft<B>
126 + VmpApplyDftToDftAdd<B>
127 + VecZnxDftApply<B>
128 + VecZnxIdftApplyConsume<B>
129 + VecZnxBigAddSmallInplace<B>
130 + VecZnxBigNormalize<B>
131 + VecZnxCopy
132 + VecZnxRotateInplace<B>
133 + VecZnxSub
134 + VecZnxNegateInplace
135 + VecZnxRshInplace<B>
136 + VecZnxAddInplace
137 + VecZnxNormalizeInplace<B>
138 + VecZnxSubInplace
139 + VecZnxRotate
140 + VecZnxAutomorphismInplace<B>
141 + VecZnxBigSubSmallNegateInplace<B>
142 + VecZnxBigAutomorphismInplace<B>
143 + VecZnxNormalize<B>
144 + VecZnxNormalizeTmpBytes,
145 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
146 {
147 assert!(
148 (self.counter as u32) < self.accumulators[0].data.n(),
149 "Packing limit of {} reached",
150 self.accumulators[0].data.n().0 as usize >> self.log_batch
151 );
152
153 pack_core(
154 module,
155 a,
156 &mut self.accumulators,
157 self.log_batch,
158 auto_keys,
159 scratch,
160 );
161 self.counter += 1 << self.log_batch;
162 }
163
164 /// Flush result to`res`.
165 pub fn flush<Data: DataMut, B: Backend>(&mut self, module: &Module<B>, res: &mut GLWECiphertext<Data>)
166 where
167 Module<B>: VecZnxCopy,
168 {
169 assert!(self.counter as u32 == self.accumulators[0].data.n());
170 // Copy result GLWE into res GLWE
171 res.copy(
172 module,
173 &self.accumulators[module.log_n() - self.log_batch - 1].data,
174 );
175
176 self.reset();
177 }
178}
179
180fn pack_core_scratch_space<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_infos: &KEY) -> usize
181where
182 OUT: GLWEInfos,
183 KEY: GGLWEInfos,
184 Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
185{
186 combine_scratch_space(module, out_infos, key_infos)
187}
188
189fn pack_core<D: DataRef, DataAK: DataRef, B: Backend>(
190 module: &Module<B>,
191 a: Option<&GLWECiphertext<D>>,
192 accumulators: &mut [Accumulator],
193 i: usize,
194 auto_keys: &HashMap<i64, AutomorphismKeyPrepared<DataAK, B>>,
195 scratch: &mut Scratch<B>,
196) where
197 Module<B>: VecZnxDftAllocBytes
198 + VmpApplyDftToDftTmpBytes
199 + VecZnxBigNormalizeTmpBytes
200 + VmpApplyDftToDft<B>
201 + VmpApplyDftToDftAdd<B>
202 + VecZnxDftApply<B>
203 + VecZnxIdftApplyConsume<B>
204 + VecZnxBigAddSmallInplace<B>
205 + VecZnxBigNormalize<B>
206 + VecZnxCopy
207 + VecZnxRotateInplace<B>
208 + VecZnxSub
209 + VecZnxNegateInplace
210 + VecZnxRshInplace<B>
211 + VecZnxAddInplace
212 + VecZnxNormalizeInplace<B>
213 + VecZnxSubInplace
214 + VecZnxRotate
215 + VecZnxAutomorphismInplace<B>
216 + VecZnxBigSubSmallNegateInplace<B>
217 + VecZnxBigAutomorphismInplace<B>
218 + VecZnxNormalize<B>
219 + VecZnxNormalizeTmpBytes,
220 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
221{
222 let log_n: usize = module.log_n();
223
224 if i == log_n {
225 return;
226 }
227
228 // Isolate the first accumulator
229 let (acc_prev, acc_next) = accumulators.split_at_mut(1);
230
231 // Control = true accumlator is free to overide
232 if !acc_prev[0].control {
233 let acc_mut_ref: &mut Accumulator = &mut acc_prev[0]; // from split_at_mut
234
235 // No previous value -> copies and sets flags accordingly
236 if let Some(a_ref) = a {
237 acc_mut_ref.data.copy(module, a_ref);
238 acc_mut_ref.value = true
239 } else {
240 acc_mut_ref.value = false
241 }
242 acc_mut_ref.control = true; // Able to be combined on next call
243 } else {
244 // Compresses acc_prev <- combine(acc_prev, a).
245 combine(module, &mut acc_prev[0], a, i, auto_keys, scratch);
246 acc_prev[0].control = false;
247
248 // Propagates to next accumulator
249 if acc_prev[0].value {
250 pack_core(
251 module,
252 Some(&acc_prev[0].data),
253 acc_next,
254 i + 1,
255 auto_keys,
256 scratch,
257 );
258 } else {
259 pack_core(
260 module,
261 None::<&GLWECiphertext<Vec<u8>>>,
262 acc_next,
263 i + 1,
264 auto_keys,
265 scratch,
266 );
267 }
268 }
269}
270
271fn combine_scratch_space<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_infos: &KEY) -> usize
272where
273 OUT: GLWEInfos,
274 KEY: GGLWEInfos,
275 Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
276{
277 GLWECiphertext::alloc_bytes(out_infos)
278 + (GLWECiphertext::rsh_scratch_space(module.n())
279 | GLWECiphertext::automorphism_inplace_scratch_space(module, out_infos, key_infos))
280}
281
282/// [combine] merges two ciphertexts together.
283fn combine<D: DataRef, DataAK: DataRef, B: Backend>(
284 module: &Module<B>,
285 acc: &mut Accumulator,
286 b: Option<&GLWECiphertext<D>>,
287 i: usize,
288 auto_keys: &HashMap<i64, AutomorphismKeyPrepared<DataAK, B>>,
289 scratch: &mut Scratch<B>,
290) where
291 Module<B>: VecZnxDftAllocBytes
292 + VmpApplyDftToDftTmpBytes
293 + VecZnxBigNormalizeTmpBytes
294 + VmpApplyDftToDft<B>
295 + VmpApplyDftToDftAdd<B>
296 + VecZnxDftApply<B>
297 + VecZnxIdftApplyConsume<B>
298 + VecZnxBigAddSmallInplace<B>
299 + VecZnxBigNormalize<B>
300 + VecZnxCopy
301 + VecZnxRotateInplace<B>
302 + VecZnxSub
303 + VecZnxNegateInplace
304 + VecZnxRshInplace<B>
305 + VecZnxAddInplace
306 + VecZnxNormalizeInplace<B>
307 + VecZnxSubInplace
308 + VecZnxRotate
309 + VecZnxAutomorphismInplace<B>
310 + VecZnxBigSubSmallNegateInplace<B>
311 + VecZnxBigAutomorphismInplace<B>
312 + VecZnxNormalize<B>
313 + VecZnxNormalizeTmpBytes,
314 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeGLWECt,
315{
316 let log_n: usize = acc.data.n().log2();
317 let a: &mut GLWECiphertext<Vec<u8>> = &mut acc.data;
318
319 let gal_el: i64 = if i == 0 {
320 -1
321 } else {
322 module.galois_element(1 << (i - 1))
323 };
324
325 let t: i64 = 1 << (log_n - i - 1);
326
327 // Goal is to evaluate: a = a + b*X^t + phi(a - b*X^t))
328 // We also use the identity: AUTO(a * X^t, g) = -X^t * AUTO(a, g)
329 // where t = 2^(log_n - i - 1) and g = 5^{2^(i - 1)}
330 // Different cases for wether a and/or b are zero.
331 //
332 // Implicite RSH without modulus switch, introduces extra I(X) * Q/2 on decryption.
333 // Necessary so that the scaling of the plaintext remains constant.
334 // It however is ok to do so here because coefficients are eventually
335 // either mapped to garbage or twice their value which vanishes I(X)
336 // since 2*(I(X) * Q/2) = I(X) * Q = 0 mod Q.
337 if acc.value {
338 if let Some(b) = b {
339 let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(a);
340
341 // a = a * X^-t
342 a.rotate_inplace(module, -t, scratch_1);
343
344 // tmp_b = a * X^-t - b
345 tmp_b.sub(module, a, b);
346 tmp_b.rsh(module, 1, scratch_1);
347
348 // a = a * X^-t + b
349 a.add_inplace(module, b);
350 a.rsh(module, 1, scratch_1);
351
352 tmp_b.normalize_inplace(module, scratch_1);
353
354 // tmp_b = phi(a * X^-t - b)
355 if let Some(key) = auto_keys.get(&gal_el) {
356 tmp_b.automorphism_inplace(module, key, scratch_1);
357 } else {
358 panic!("auto_key[{gal_el}] not found");
359 }
360
361 // a = a * X^-t + b - phi(a * X^-t - b)
362 a.sub_inplace_ab(module, &tmp_b);
363 a.normalize_inplace(module, scratch_1);
364
365 // a = a + b * X^t - phi(a * X^-t - b) * X^t
366 // = a + b * X^t - phi(a * X^-t - b) * - phi(X^t)
367 // = a + b * X^t + phi(a - b * X^t)
368 a.rotate_inplace(module, t, scratch_1);
369 } else {
370 a.rsh(module, 1, scratch);
371 // a = a + phi(a)
372 if let Some(key) = auto_keys.get(&gal_el) {
373 a.automorphism_add_inplace(module, key, scratch);
374 } else {
375 panic!("auto_key[{gal_el}] not found");
376 }
377 }
378 } else if let Some(b) = b {
379 let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(a);
380 tmp_b.rotate(module, 1 << (log_n - i - 1), b);
381 tmp_b.rsh(module, 1, scratch_1);
382
383 // a = (b* X^t - phi(b* X^t))
384 if let Some(key) = auto_keys.get(&gal_el) {
385 a.automorphism_sub_negate(module, &tmp_b, key, scratch_1);
386 } else {
387 panic!("auto_key[{gal_el}] not found");
388 }
389
390 acc.value = true;
391 }
392}
393
394/// Packs [x_0: GLWE(m_0), x_1: GLWE(m_1), ..., x_i: GLWE(m_i)]
395/// to [0: GLWE(m_0 * X^x_0 + m_1 * X^x_1 + ... + m_i * X^x_i)]
396pub fn glwe_packing<D: DataMut, ATK, B: Backend>(
397 module: &Module<B>,
398 cts: &mut HashMap<usize, &mut GLWECiphertext<D>>,
399 log_gap_out: usize,
400 auto_keys: &HashMap<i64, AutomorphismKeyPrepared<ATK, B>>,
401 scratch: &mut Scratch<B>,
402) where
403 ATK: DataRef,
404 Module<B>: VecZnxRotateInplace<B>
405 + VecZnxNormalizeInplace<B>
406 + VecZnxNormalizeTmpBytes
407 + VecZnxSwitchRing
408 + VecZnxBigAutomorphismInplace<B>
409 + VecZnxRshInplace<B>
410 + VecZnxDftCopy<B>
411 + VecZnxIdftApplyTmpA<B>
412 + VecZnxSub
413 + VecZnxAddInplace
414 + VecZnxNegateInplace
415 + VecZnxCopy
416 + VecZnxSubInplace
417 + VecZnxDftAllocBytes
418 + VmpApplyDftToDftTmpBytes
419 + VecZnxBigNormalizeTmpBytes
420 + VmpApplyDftToDft<B>
421 + VmpApplyDftToDftAdd<B>
422 + VecZnxDftApply<B>
423 + VecZnxIdftApplyConsume<B>
424 + VecZnxBigAddSmallInplace<B>
425 + VecZnxBigNormalize<B>
426 + VecZnxAutomorphismInplace<B>
427 + VecZnxBigSubSmallNegateInplace<B>
428 + VecZnxRotate
429 + VecZnxNormalize<B>,
430 Scratch<B>: TakeVecZnx + TakeVecZnxDft<B> + ScratchAvailable,
431{
432 #[cfg(debug_assertions)]
433 {
434 assert!(*cts.keys().max().unwrap() < module.n())
435 }
436
437 let log_n: usize = module.log_n();
438
439 (0..log_n - log_gap_out).for_each(|i| {
440 let t: usize = (1 << log_n).min(1 << (log_n - 1 - i));
441
442 let auto_key: &AutomorphismKeyPrepared<ATK, B> = if i == 0 {
443 auto_keys.get(&-1).unwrap()
444 } else {
445 auto_keys.get(&module.galois_element(1 << (i - 1))).unwrap()
446 };
447
448 (0..t).for_each(|j| {
449 let mut a: Option<&mut GLWECiphertext<D>> = cts.remove(&j);
450 let mut b: Option<&mut GLWECiphertext<D>> = cts.remove(&(j + t));
451
452 pack_internal(module, &mut a, &mut b, i, auto_key, scratch);
453
454 if let Some(a) = a {
455 cts.insert(j, a);
456 } else if let Some(b) = b {
457 cts.insert(j, b);
458 }
459 });
460 });
461}
462
463#[allow(clippy::too_many_arguments)]
464fn pack_internal<A: DataMut, D: DataMut, DataAK: DataRef, B: Backend>(
465 module: &Module<B>,
466 a: &mut Option<&mut GLWECiphertext<A>>,
467 b: &mut Option<&mut GLWECiphertext<D>>,
468 i: usize,
469 auto_key: &AutomorphismKeyPrepared<DataAK, B>,
470 scratch: &mut Scratch<B>,
471) where
472 Module<B>: VecZnxRotateInplace<B>
473 + VecZnxNormalizeInplace<B>
474 + VecZnxNormalizeTmpBytes
475 + VecZnxBigAutomorphismInplace<B>
476 + VecZnxRshInplace<B>
477 + VecZnxDftCopy<B>
478 + VecZnxIdftApplyTmpA<B>
479 + VecZnxSub
480 + VecZnxAddInplace
481 + VecZnxNegateInplace
482 + VecZnxCopy
483 + VecZnxSubInplace
484 + VecZnxDftAllocBytes
485 + VmpApplyDftToDftTmpBytes
486 + VecZnxBigNormalizeTmpBytes
487 + VmpApplyDftToDft<B>
488 + VmpApplyDftToDftAdd<B>
489 + VecZnxDftApply<B>
490 + VecZnxIdftApplyConsume<B>
491 + VecZnxBigAddSmallInplace<B>
492 + VecZnxBigNormalize<B>
493 + VecZnxAutomorphismInplace<B>
494 + VecZnxBigSubSmallNegateInplace<B>
495 + VecZnxRotate
496 + VecZnxNormalize<B>,
497 Scratch<B>: TakeVecZnx + TakeVecZnxDft<B> + ScratchAvailable,
498{
499 // Goal is to evaluate: a = a + b*X^t + phi(a - b*X^t))
500 // We also use the identity: AUTO(a * X^t, g) = -X^t * AUTO(a, g)
501 // where t = 2^(log_n - i - 1) and g = 5^{2^(i - 1)}
502 // Different cases for wether a and/or b are zero.
503 //
504 // Implicite RSH without modulus switch, introduces extra I(X) * Q/2 on decryption.
505 // Necessary so that the scaling of the plaintext remains constant.
506 // It however is ok to do so here because coefficients are eventually
507 // either mapped to garbage or twice their value which vanishes I(X)
508 // since 2*(I(X) * Q/2) = I(X) * Q = 0 mod Q.
509 if let Some(a) = a.as_deref_mut() {
510 let t: i64 = 1 << (a.n().log2() - i - 1);
511
512 if let Some(b) = b.as_deref_mut() {
513 let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(a);
514
515 // a = a * X^-t
516 a.rotate_inplace(module, -t, scratch_1);
517
518 // tmp_b = a * X^-t - b
519 tmp_b.sub(module, a, b);
520 tmp_b.rsh(module, 1, scratch_1);
521
522 // a = a * X^-t + b
523 a.add_inplace(module, b);
524 a.rsh(module, 1, scratch_1);
525
526 tmp_b.normalize_inplace(module, scratch_1);
527
528 // tmp_b = phi(a * X^-t - b)
529 tmp_b.automorphism_inplace(module, auto_key, scratch_1);
530
531 // a = a * X^-t + b - phi(a * X^-t - b)
532 a.sub_inplace_ab(module, &tmp_b);
533 a.normalize_inplace(module, scratch_1);
534
535 // a = a + b * X^t - phi(a * X^-t - b) * X^t
536 // = a + b * X^t - phi(a * X^-t - b) * - phi(X^t)
537 // = a + b * X^t + phi(a - b * X^t)
538 a.rotate_inplace(module, t, scratch_1);
539 } else {
540 a.rsh(module, 1, scratch);
541 // a = a + phi(a)
542 a.automorphism_add_inplace(module, auto_key, scratch);
543 }
544 } else if let Some(b) = b.as_deref_mut() {
545 let t: i64 = 1 << (b.n().log2() - i - 1);
546
547 let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(b);
548 tmp_b.rotate(module, t, b);
549 tmp_b.rsh(module, 1, scratch_1);
550
551 // a = (b* X^t - phi(b* X^t))
552 b.automorphism_sub_negate(module, &tmp_b, auto_key, scratch_1);
553 }
554}