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