poulpy/poulpy-core/src/glwe_packing.rs
poulpy-core/src/glwe_packing.rsRaw
1use std::collections::HashMap;
2
3use poulpy_hal::{
4 api::{
5 ModuleLogN, ModuleN, ScratchAvailable, VecZnxAddInplace, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace, VecZnxBigAutomorphismInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallNegateInplace, VecZnxCopy, VecZnxDftApply, VecZnxDftBytesOf, VecZnxDftCopy, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNegateInplace, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotate, VecZnxRotateInplace, VecZnxRshInplace, VecZnxSub, VecZnxSubInplace, VecZnxSwitchRing, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes
6 },
7 layouts::{Backend, DataMut, DataRef, GaloisElement, Module, Scratch},
8};
9
10use crate::{
11 layouts::{prepared::{AutomorphismKeyPrepared, AutomorphismKeyPreparedToRef}, GGLWEInfos, GLWEAlloc, GLWEInfos, GLWEToRef, LWEInfos, GLWE}, GLWEAutomorphism, GLWEOperations, ScratchTakeCore
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, M>(module: &M, infos: &A) -> Self
42 where
43 A: GLWEInfos,
44 M: GLWEAlloc
45 {
46 Self {
47 data: GLWE::alloc_from_infos(module, infos),
48 value: false,
49 control: false,
50 }
51 }
52}
53
54impl GLWEPacker {
55 /// Instantiates a new [GLWEPacker].
56 ///
57 /// # Arguments
58 ///
59 /// * `module`: static backend FFT tables.
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 new<A, M>(module: &M, infos: &A, log_batch: usize) -> Self
66 where
67 A: GLWEInfos,
68 M: GLWEAlloc
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(module, infos)));
73 Self {
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: GLWEAlloc + GLWEAutomorphism<BE>,
95 {
96 pack_core_tmp_bytes(module, res_infos, key_infos)
97 }
98
99 pub fn galois_elements<B: Backend>(module: &Module<B>) -> Vec<i64> {
100 GLWE::trace_galois_elements(module)
101 }
102
103 /// Adds a GLWE ciphertext to the [GLWEPacker].
104 /// #Arguments
105 ///
106 /// * `module`: static backend FFT tables.
107 /// * `res`: space to append fully packed ciphertext. Only when the number
108 /// of packed ciphertexts reaches N/2^log_batch is a result written.
109 /// * `a`: ciphertext to pack. Can optionally give None to pack a 0 ciphertext.
110 /// * `auto_keys`: a [HashMap] containing the [AutomorphismKeyExec]s.
111 /// * `scratch`: scratch space of size at least [Self::tmp_bytes].
112 pub fn add<A, K, M, BE: Backend>(
113 &mut self,
114 module: &M,
115 a: Option<&A>,
116 auto_keys: &HashMap<i64, K>,
117 scratch: &mut Scratch<B>,
118 ) where
119 A: GLWEToRef,
120 K: AutomorphismKeyPreparedToRef<BE>,
121 M: GLWEAutomorphism<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<Data: DataMut, B: Backend>(&mut self, module: &Module<B>, res: &mut GLWE<Data>)
143 where
144 Module<B>: VecZnxCopy,
145 {
146 assert!(self.counter as u32 == self.accumulators[0].data.n());
147 // Copy result GLWE into res GLWE
148 res.copy(
149 module,
150 &self.accumulators[module.log_n() - self.log_batch - 1].data,
151 );
152
153 self.reset();
154 }
155}
156
157fn pack_core_tmp_bytes<R, K, M, BE: Backend>(module: &M, res_infos: &R, key_infos: &K) -> usize
158where
159 R: GLWEInfos,
160 K: GGLWEInfos,
161 M: GLWEAlloc + GLWEAutomorphism<BE>,
162{
163 combine_tmp_bytes(module, res_infos, key_infos)
164}
165
166fn pack_core<A, K, M, BE: Backend>(
167 module: &M,
168 a: Option<&A>,
169 accumulators: &mut [Accumulator],
170 i: usize,
171 auto_keys: &HashMap<i64, K>,
172 scratch: &mut Scratch<BE>,
173) where
174 A: GLWEToRef + GLWEInfos,
175 K: AutomorphismKeyPreparedToRef<BE>,
176 M: GLWEAutomorphism<BE> + ModuleLogN + VecZnxCopy,
177 Scratch<BE>: ScratchTakeCore<BE>,
178{
179 let log_n: usize = module.log_n();
180
181 if i == log_n {
182 return;
183 }
184
185 // Isolate the first accumulator
186 let (acc_prev, acc_next) = accumulators.split_at_mut(1);
187
188 // Control = true accumlator is free to overide
189 if !acc_prev[0].control {
190 let acc_mut_ref: &mut Accumulator = &mut acc_prev[0]; // from split_at_mut
191
192 // No previous value -> copies and sets flags accordingly
193 if let Some(a_ref) = a {
194 acc_mut_ref.data.copy(module, a_ref);
195 acc_mut_ref.value = true
196 } else {
197 acc_mut_ref.value = false
198 }
199 acc_mut_ref.control = true; // Able to be combined on next call
200 } else {
201 // Compresses acc_prev <- combine(acc_prev, a).
202 combine(module, &mut acc_prev[0], a, i, auto_keys, scratch);
203 acc_prev[0].control = false;
204
205 // Propagates to next accumulator
206 if acc_prev[0].value {
207 pack_core(
208 module,
209 Some(&acc_prev[0].data),
210 acc_next,
211 i + 1,
212 auto_keys,
213 scratch,
214 );
215 } else {
216 pack_core(
217 module,
218 None::<&GLWE<Vec<u8>>>,
219 acc_next,
220 i + 1,
221 auto_keys,
222 scratch,
223 );
224 }
225 }
226}
227
228fn combine_tmp_bytes<R, K, M, BE: Backend>(module: &M, res_infos: &R, key_infos: &K) -> usize
229where
230 R: GLWEInfos,
231 K: GGLWEInfos,
232 M: GLWEAlloc + GLWEAutomorphism<BE>,
233{
234 GLWE::bytes_of_from_infos(module, res_infos)
235 + (GLWE::rsh_tmp_bytes(module.n()) | module.glwe_automorphism_tmp_bytes(res_infos, res_infos, key_infos))
236}
237
238/// [combine] merges two ciphertexts together.
239fn combine<B, M, K, BE: Backend>(
240 module: &M,
241 acc: &mut Accumulator,
242 b: Option<&B>,
243 i: usize,
244 auto_keys: &HashMap<i64, K>,
245 scratch: &mut Scratch<BE>,
246) where
247 B: GLWEToRef + GLWEInfos,
248 K: AutomorphismKeyPreparedToRef<BE>,
249 M: GLWEAutomorphism<BE> + GaloisElement + VecZnxRotateInplace<BE>,
250 Scratch<BE>: ScratchTakeCore<BE>,
251{
252 let log_n: usize = acc.data.n().log2();
253 let a: &mut GLWE<Vec<u8>> = &mut acc.data;
254
255 let gal_el: i64 = if i == 0 {
256 -1
257 } else {
258 module.galois_element(1 << (i - 1))
259 };
260
261 let t: i64 = 1 << (log_n - i - 1);
262
263 // Goal is to evaluate: a = a + b*X^t + phi(a - b*X^t))
264 // We also use the identity: AUTO(a * X^t, g) = -X^t * AUTO(a, g)
265 // where t = 2^(log_n - i - 1) and g = 5^{2^(i - 1)}
266 // Different cases for wether a and/or b are zero.
267 //
268 // Implicite RSH without modulus switch, introduces extra I(X) * Q/2 on decryption.
269 // Necessary so that the scaling of the plaintext remains constant.
270 // It however is ok to do so here because coefficients are eventually
271 // either mapped to garbage or twice their value which vanishes I(X)
272 // since 2*(I(X) * Q/2) = I(X) * Q = 0 mod Q.
273 if acc.value {
274 if let Some(b) = b {
275 let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(module, a);
276
277 // a = a * X^-t
278 a.rotate_inplace(module, -t, scratch_1);
279
280 // tmp_b = a * X^-t - b
281 tmp_b.sub(module, a, b);
282 tmp_b.rsh(module, 1, scratch_1);
283
284 // a = a * X^-t + b
285 a.add_inplace(module, b);
286 a.rsh(module, 1, scratch_1);
287
288 tmp_b.normalize_inplace(module, scratch_1);
289
290 // tmp_b = phi(a * X^-t - b)
291 if let Some(key) = auto_keys.get(&gal_el) {
292 tmp_b.automorphism_inplace(module, key, scratch_1);
293 } else {
294 panic!("auto_key[{gal_el}] not found");
295 }
296
297 // a = a * X^-t + b - phi(a * X^-t - b)
298 a.sub_inplace_ab(module, &tmp_b);
299 a.normalize_inplace(module, scratch_1);
300
301 // a = a + b * X^t - phi(a * X^-t - b) * X^t
302 // = a + b * X^t - phi(a * X^-t - b) * - phi(X^t)
303 // = a + b * X^t + phi(a - b * X^t)
304 a.rotate_inplace(module, t, scratch_1);
305 } else {
306 a.rsh(module, 1, scratch);
307 // a = a + phi(a)
308 if let Some(key) = auto_keys.get(&gal_el) {
309 a.automorphism_add_inplace(module, key, scratch);
310 } else {
311 panic!("auto_key[{gal_el}] not found");
312 }
313 }
314 } else if let Some(b) = b {
315 let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(a);
316 tmp_b.rotate(module, 1 << (log_n - i - 1), b);
317 tmp_b.rsh(module, 1, scratch_1);
318
319 // a = (b* X^t - phi(b* X^t))
320 if let Some(key) = auto_keys.get(&gal_el) {
321 a.automorphism_sub_negate(module, &tmp_b, key, scratch_1);
322 } else {
323 panic!("auto_key[{gal_el}] not found");
324 }
325
326 acc.value = true;
327 }
328}
329
330/// Packs [x_0: GLWE(m_0), x_1: GLWE(m_1), ..., x_i: GLWE(m_i)]
331/// to [0: GLWE(m_0 * X^x_0 + m_1 * X^x_1 + ... + m_i * X^x_i)]
332pub fn glwe_packing<D: DataMut, ATK, B: Backend>(
333 module: &Module<B>,
334 cts: &mut HashMap<usize, &mut GLWE<D>>,
335 log_gap_out: usize,
336 auto_keys: &HashMap<i64, AutomorphismKeyPrepared<ATK, B>>,
337 scratch: &mut Scratch<B>,
338) where
339 ATK: DataRef,
340 Module<B>: VecZnxRotateInplace<B>
341 + VecZnxNormalizeInplace<B>
342 + VecZnxNormalizeTmpBytes
343 + VecZnxSwitchRing
344 + VecZnxBigAutomorphismInplace<B>
345 + VecZnxRshInplace<B>
346 + VecZnxDftCopy<B>
347 + VecZnxIdftApplyTmpA<B>
348 + VecZnxSub
349 + VecZnxAddInplace
350 + VecZnxNegateInplace
351 + VecZnxCopy
352 + VecZnxSubInplace
353 + VecZnxDftBytesOf
354 + VmpApplyDftToDftTmpBytes
355 + VecZnxBigNormalizeTmpBytes
356 + VmpApplyDftToDft<B>
357 + VmpApplyDftToDftAdd<B>
358 + VecZnxDftApply<B>
359 + VecZnxIdftApplyConsume<B>
360 + VecZnxBigAddSmallInplace<B>
361 + VecZnxBigNormalize<B>
362 + VecZnxAutomorphismInplace<B>
363 + VecZnxBigSubSmallNegateInplace<B>
364 + VecZnxRotate
365 + VecZnxNormalize<B>,
366 Scratch<B>: ScratchAvailable,
367{
368 #[cfg(debug_assertions)]
369 {
370 assert!(*cts.keys().max().unwrap() < module.n())
371 }
372
373 let log_n: usize = module.log_n();
374
375 (0..log_n - log_gap_out).for_each(|i| {
376 let t: usize = (1 << log_n).min(1 << (log_n - 1 - i));
377
378 let auto_key: &AutomorphismKeyPrepared<ATK, B> = if i == 0 {
379 auto_keys.get(&-1).unwrap()
380 } else {
381 auto_keys.get(&module.galois_element(1 << (i - 1))).unwrap()
382 };
383
384 (0..t).for_each(|j| {
385 let mut a: Option<&mut GLWE<D>> = cts.remove(&j);
386 let mut b: Option<&mut GLWE<D>> = cts.remove(&(j + t));
387
388 pack_internal(module, &mut a, &mut b, i, auto_key, scratch);
389
390 if let Some(a) = a {
391 cts.insert(j, a);
392 } else if let Some(b) = b {
393 cts.insert(j, b);
394 }
395 });
396 });
397}
398
399#[allow(clippy::too_many_arguments)]
400fn pack_internal<A: DataMut, D: DataMut, DataAK: DataRef, B: Backend>(
401 module: &Module<B>,
402 a: &mut Option<&mut GLWE<A>>,
403 b: &mut Option<&mut GLWE<D>>,
404 i: usize,
405 auto_key: &AutomorphismKeyPrepared<DataAK, B>,
406 scratch: &mut Scratch<B>,
407) where
408 Module<B>: VecZnxRotateInplace<B>
409 + VecZnxNormalizeInplace<B>
410 + VecZnxNormalizeTmpBytes
411 + VecZnxBigAutomorphismInplace<B>
412 + VecZnxRshInplace<B>
413 + VecZnxDftCopy<B>
414 + VecZnxIdftApplyTmpA<B>
415 + VecZnxSub
416 + VecZnxAddInplace
417 + VecZnxNegateInplace
418 + VecZnxCopy
419 + VecZnxSubInplace
420 + VecZnxDftBytesOf
421 + VmpApplyDftToDftTmpBytes
422 + VecZnxBigNormalizeTmpBytes
423 + VmpApplyDftToDft<B>
424 + VmpApplyDftToDftAdd<B>
425 + VecZnxDftApply<B>
426 + VecZnxIdftApplyConsume<B>
427 + VecZnxBigAddSmallInplace<B>
428 + VecZnxBigNormalize<B>
429 + VecZnxAutomorphismInplace<B>
430 + VecZnxBigSubSmallNegateInplace<B>
431 + VecZnxRotate
432 + VecZnxNormalize<B>,
433 Scratch<B>: ScratchAvailable,
434{
435 // Goal is to evaluate: a = a + b*X^t + phi(a - b*X^t))
436 // We also use the identity: AUTO(a * X^t, g) = -X^t * AUTO(a, g)
437 // where t = 2^(log_n - i - 1) and g = 5^{2^(i - 1)}
438 // Different cases for wether a and/or b are zero.
439 //
440 // Implicite RSH without modulus switch, introduces extra I(X) * Q/2 on decryption.
441 // Necessary so that the scaling of the plaintext remains constant.
442 // It however is ok to do so here because coefficients are eventually
443 // either mapped to garbage or twice their value which vanishes I(X)
444 // since 2*(I(X) * Q/2) = I(X) * Q = 0 mod Q.
445 if let Some(a) = a.as_deref_mut() {
446 let t: i64 = 1 << (a.n().log2() - i - 1);
447
448 if let Some(b) = b.as_deref_mut() {
449 let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(a);
450
451 // a = a * X^-t
452 a.rotate_inplace(module, -t, scratch_1);
453
454 // tmp_b = a * X^-t - b
455 tmp_b.sub(module, a, b);
456 tmp_b.rsh(module, 1, scratch_1);
457
458 // a = a * X^-t + b
459 a.add_inplace(module, b);
460 a.rsh(module, 1, scratch_1);
461
462 tmp_b.normalize_inplace(module, scratch_1);
463
464 // tmp_b = phi(a * X^-t - b)
465 tmp_b.automorphism_inplace(module, auto_key, scratch_1);
466
467 // a = a * X^-t + b - phi(a * X^-t - b)
468 a.sub_inplace_ab(module, &tmp_b);
469 a.normalize_inplace(module, scratch_1);
470
471 // a = a + b * X^t - phi(a * X^-t - b) * X^t
472 // = a + b * X^t - phi(a * X^-t - b) * - phi(X^t)
473 // = a + b * X^t + phi(a - b * X^t)
474 a.rotate_inplace(module, t, scratch_1);
475 } else {
476 a.rsh(module, 1, scratch);
477 // a = a + phi(a)
478 a.automorphism_add_inplace(module, auto_key, scratch);
479 }
480 } else if let Some(b) = b.as_deref_mut() {
481 let t: i64 = 1 << (b.n().log2() - i - 1);
482
483 let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(b);
484 tmp_b.rotate(module, t, b);
485 tmp_b.rsh(module, 1, scratch_1);
486
487 // a = (b* X^t - phi(b* X^t))
488 b.automorphism_sub_negate(module, &tmp_b, auto_key, scratch_1);
489 }
490}