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