poulpy/poulpy-core/src/glwe_packer.rs
poulpy-core/src/glwe_packer.rsRaw
1use poulpy_hal::{
2 api::ModuleLogN,
3 layouts::{Backend, GaloisElement, Module, Scratch},
4};
5
6use crate::{
7 GLWEAdd, GLWEAutomorphism, GLWECopy, GLWENormalize, GLWERotate, GLWEShift, GLWESub, ScratchTakeCore,
8 glwe_trace::GLWETrace,
9 layouts::{
10 GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEAutomorphismKeyHelper, GLWEInfos, GLWEToMut, GLWEToRef, GetGaloisElement,
11 LWEInfos,
12 },
13};
14
15/// [GLWEPacker] enables only the fly GLWE packing
16/// with constant memory of Log(N) ciphertexts.
17/// Main difference with usual GLWE packing is that
18/// the output is bit-reversed.
19pub struct GLWEPacker {
20 accumulators: Vec<Accumulator>,
21 log_batch: usize,
22 counter: usize,
23}
24
25/// [Accumulator] stores intermediate packing result.
26/// There are Log(N) such accumulators in a [GLWEPacker].
27struct Accumulator {
28 data: GLWE<Vec<u8>>,
29 value: bool, // Implicit flag for zero ciphertext
30 control: bool, // Can be combined with incoming value
31}
32
33impl Accumulator {
34 /// Allocates a new [Accumulator].
35 ///
36 /// #Arguments
37 ///
38 /// * `module`: static backend FFT tables.
39 /// * `base2k`: base 2 logarithm of the GLWE ciphertext in memory digit representation.
40 /// * `k`: base 2 precision of the GLWE ciphertext precision over the Torus.
41 /// * `rank`: rank of the GLWE ciphertext.
42 pub fn alloc<A>(infos: &A) -> Self
43 where
44 A: GLWEInfos,
45 {
46 Self {
47 data: GLWE::alloc_from_infos(infos),
48 value: false,
49 control: false,
50 }
51 }
52}
53
54impl GLWEPacker {
55 /// Instantiates a new [GLWEPacker].
56 ///
57 /// # Arguments
58 ///
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 alloc<A>(infos: &A, log_batch: usize) -> Self
65 where
66 A: GLWEInfos,
67 {
68 let mut accumulators: Vec<Accumulator> = Vec::<Accumulator>::new();
69 let log_n: usize = infos.n().log2();
70 (0..log_n - log_batch).for_each(|_| accumulators.push(Accumulator::alloc(infos)));
71 GLWEPacker {
72 accumulators,
73 log_batch,
74 counter: 0,
75 }
76 }
77
78 /// Implicit reset of the internal state (to be called before a new packing procedure).
79 fn reset(&mut self) {
80 for i in 0..self.accumulators.len() {
81 self.accumulators[i].value = false;
82 self.accumulators[i].control = false;
83 }
84 self.counter = 0;
85 }
86
87 /// Number of scratch space bytes required to call [Self::add].
88 pub fn tmp_bytes<R, K, M, BE: Backend>(module: &M, res_infos: &R, key_infos: &K) -> usize
89 where
90 R: GLWEInfos,
91 K: GGLWEInfos,
92 M: GLWEPackerOps<BE>,
93 {
94 GLWE::bytes_of_from_infos(res_infos)
95 + module
96 .glwe_rsh_tmp_byte()
97 .max(module.glwe_automorphism_tmp_bytes(res_infos, res_infos, key_infos))
98 }
99
100 pub fn galois_elements<M, BE: Backend>(module: &M) -> Vec<i64>
101 where
102 M: GLWETrace<BE>,
103 {
104 module.glwe_trace_galois_elements()
105 }
106
107 /// Adds a GLWE ciphertext to the [GLWEPacker].
108 /// #Arguments
109 ///
110 /// * `module`: static backend FFT tables.
111 /// * `res`: space to append fully packed ciphertext. Only when the number
112 /// of packed ciphertexts reaches N/2^log_batch is a result written.
113 /// * `a`: ciphertext to pack. Can optionally give None to pack a 0 ciphertext.
114 /// * `auto_keys`: an implementation of [GLWEAutomorphismKeyHelper], containing [GLWEAutomorphismKeyPrepared] with index of [Self::galois_elements].
115 /// * `scratch`: scratch space of size at least [Self::tmp_bytes].
116 pub fn add<A, K, H, M, BE: Backend>(&mut self, module: &M, a: Option<&A>, auto_keys: &H, scratch: &mut Scratch<BE>)
117 where
118 A: GLWEToRef + GLWEInfos,
119 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
120 H: GLWEAutomorphismKeyHelper<K, BE>,
121 M: GLWEPackerOps<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 module.packer_add(self, a, self.log_batch, auto_keys, scratch);
131 self.counter += 1 << self.log_batch;
132 }
133
134 /// Flush result to`res`.
135 pub fn flush<R, M, BE: Backend>(&mut self, module: &M, res: &mut R)
136 where
137 R: GLWEToMut,
138 M: GLWEPackerOps<BE>,
139 {
140 assert!(self.counter as u32 == self.accumulators[0].data.n());
141 // Copy result GLWE into res GLWE
142 module.glwe_copy(
143 res,
144 &self.accumulators[module.log_n() - self.log_batch - 1].data,
145 );
146
147 self.reset();
148 }
149}
150
151impl<BE: Backend> GLWEPackerOps<BE> for Module<BE> where
152 Self: Sized
153 + ModuleLogN
154 + GLWEAutomorphism<BE>
155 + GaloisElement
156 + GLWERotate<BE>
157 + GLWESub
158 + GLWEShift<BE>
159 + GLWEAdd
160 + GLWENormalize<BE>
161 + GLWECopy
162 + GLWEAutomorphism<BE>
163 + GaloisElement
164 + GLWERotate<BE>
165 + GLWESub
166 + GLWEShift<BE>
167 + GLWEAdd
168 + GLWENormalize<BE>
169{
170}
171
172pub trait GLWEPackerOps<BE: Backend>
173where
174 Self: Sized
175 + ModuleLogN
176 + GLWEAutomorphism<BE>
177 + GaloisElement
178 + GLWERotate<BE>
179 + GLWESub
180 + GLWEShift<BE>
181 + GLWEAdd
182 + GLWENormalize<BE>
183 + GLWECopy
184 + GLWEAutomorphism<BE>
185 + GaloisElement
186 + GLWERotate<BE>
187 + GLWESub
188 + GLWEShift<BE>
189 + GLWEAdd
190 + GLWENormalize<BE>,
191{
192 fn packer_add<A, K, H>(&self, packer: &mut GLWEPacker, a: Option<&A>, i: usize, auto_keys: &H, scratch: &mut Scratch<BE>)
193 where
194 A: GLWEToRef + GLWEInfos,
195 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
196 H: GLWEAutomorphismKeyHelper<K, BE>,
197 Scratch<BE>: ScratchTakeCore<BE>,
198 {
199 pack_core(self, a, &mut packer.accumulators, i, auto_keys, scratch)
200 }
201}
202
203fn pack_core<A, K, H, M, BE: Backend>(
204 module: &M,
205 a: Option<&A>,
206 accumulators: &mut [Accumulator],
207 i: usize,
208 auto_keys: &H,
209 scratch: &mut Scratch<BE>,
210) where
211 A: GLWEToRef + GLWEInfos,
212 M: ModuleLogN
213 + GLWEAutomorphism<BE>
214 + GaloisElement
215 + GLWERotate<BE>
216 + GLWESub
217 + GLWEShift<BE>
218 + GLWEAdd
219 + GLWENormalize<BE>
220 + GLWECopy
221 + GLWEAutomorphism<BE>
222 + GaloisElement
223 + GLWERotate<BE>
224 + GLWESub
225 + GLWEShift<BE>
226 + GLWEAdd
227 + GLWENormalize<BE>,
228 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
229 H: GLWEAutomorphismKeyHelper<K, BE>,
230 Scratch<BE>: ScratchTakeCore<BE>,
231{
232 let log_n: usize = module.log_n();
233
234 if i == log_n {
235 return;
236 }
237
238 // Isolate the first accumulator
239 let (acc_prev, acc_next) = accumulators.split_at_mut(1);
240
241 // Control = true accumlator is free to overide
242 if !acc_prev[0].control {
243 let acc_mut_ref: &mut Accumulator = &mut acc_prev[0]; // from split_at_mut
244
245 // No previous value -> copies and sets flags accordingly
246 if let Some(a_ref) = a {
247 module.glwe_copy(&mut acc_mut_ref.data, a_ref);
248 acc_mut_ref.value = true
249 } else {
250 acc_mut_ref.value = false
251 }
252 acc_mut_ref.control = true; // Able to be combined on next call
253 } else {
254 // Compresses acc_prev <- combine(acc_prev, a).
255 combine(module, &mut acc_prev[0], a, i, auto_keys, scratch);
256 acc_prev[0].control = false;
257
258 // Propagates to next accumulator
259 if acc_prev[0].value {
260 pack_core(
261 module,
262 Some(&acc_prev[0].data),
263 acc_next,
264 i + 1,
265 auto_keys,
266 scratch,
267 );
268 } else {
269 pack_core(
270 module,
271 None::<&GLWE<Vec<u8>>>,
272 acc_next,
273 i + 1,
274 auto_keys,
275 scratch,
276 );
277 }
278 }
279}
280
281fn combine<B, K, H, M, BE: Backend>(
282 module: &M,
283 acc: &mut Accumulator,
284 b: Option<&B>,
285 i: usize,
286 auto_keys: &H,
287 scratch: &mut Scratch<BE>,
288) where
289 B: GLWEToRef + GLWEInfos,
290 B: GLWEToRef + GLWEInfos,
291 M: ModuleLogN
292 + GLWEAutomorphism<BE>
293 + GaloisElement
294 + GLWERotate<BE>
295 + GLWESub
296 + GLWEShift<BE>
297 + GLWEAdd
298 + GLWENormalize<BE>
299 + GLWECopy
300 + GLWEAutomorphism<BE>
301 + GaloisElement
302 + GLWERotate<BE>
303 + GLWESub
304 + GLWEShift<BE>
305 + GLWEAdd
306 + GLWENormalize<BE>,
307 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
308 H: GLWEAutomorphismKeyHelper<K, BE>,
309 Scratch<BE>: ScratchTakeCore<BE>,
310{
311 let log_n: usize = acc.data.n().log2();
312 let a: &mut GLWE<Vec<u8>> = &mut acc.data;
313
314 let gal_el: i64 = if i == 0 {
315 -1
316 } else {
317 module.galois_element(1 << (i - 1))
318 };
319
320 let t: i64 = 1 << (log_n - i - 1);
321
322 // Goal is to evaluate: a = a + b*X^t + phi(a - b*X^t))
323 // We also use the identity: AUTO(a * X^t, g) = -X^t * AUTO(a, g)
324 // where t = 2^(log_n - i - 1) and g = 5^{2^(i - 1)}
325 // Different cases for wether a and/or b are zero.
326 //
327 // Implicite RSH without modulus switch, introduces extra I(X) * Q/2 on decryption.
328 // Necessary so that the scaling of the plaintext remains constant.
329 // It however is ok to do so here because coefficients are eventually
330 // either mapped to garbage or twice their value which vanishes I(X)
331 // since 2*(I(X) * Q/2) = I(X) * Q = 0 mod Q.
332 if acc.value {
333 if let Some(b) = b {
334 let (mut tmp_b, scratch_1) = scratch.take_glwe(a);
335
336 // a = a * X^-t
337 module.glwe_rotate_inplace(-t, a, scratch_1);
338
339 // tmp_b = a * X^-t - b
340 module.glwe_sub(&mut tmp_b, a, b);
341 module.glwe_rsh(1, &mut tmp_b, scratch_1);
342
343 // a = a * X^-t + b
344 module.glwe_add_inplace(a, b);
345 module.glwe_rsh(1, a, scratch_1);
346
347 module.glwe_normalize_inplace(&mut tmp_b, scratch_1);
348
349 // tmp_b = phi(a * X^-t - b)
350 if let Some(auto_key) = auto_keys.get_automorphism_key(gal_el) {
351 module.glwe_automorphism_inplace(&mut tmp_b, auto_key, scratch_1);
352 } else {
353 panic!("auto_key[{gal_el}] not found");
354 }
355
356 // a = a * X^-t + b - phi(a * X^-t - b)
357 module.glwe_sub_inplace(a, &tmp_b);
358 module.glwe_normalize_inplace(a, scratch_1);
359
360 // a = a + b * X^t - phi(a * X^-t - b) * X^t
361 // = a + b * X^t - phi(a * X^-t - b) * - phi(X^t)
362 // = a + b * X^t + phi(a - b * X^t)
363 module.glwe_rotate_inplace(t, a, scratch_1);
364 } else {
365 module.glwe_rsh(1, a, scratch);
366 // a = a + phi(a)
367 if let Some(auto_key) = auto_keys.get_automorphism_key(gal_el) {
368 module.glwe_automorphism_add_inplace(a, auto_key, scratch);
369 } else {
370 panic!("auto_key[{gal_el}] not found");
371 }
372 }
373 } else if let Some(b) = b {
374 let (mut tmp_b, scratch_1) = scratch.take_glwe(a);
375 module.glwe_rotate(t, &mut tmp_b, b);
376 module.glwe_rsh(1, &mut tmp_b, scratch_1);
377
378 // a = (b* X^t - phi(b* X^t))
379 if let Some(auto_key) = auto_keys.get_automorphism_key(gal_el) {
380 module.glwe_automorphism_sub_negate(a, &tmp_b, auto_key, scratch_1);
381 } else {
382 panic!("auto_key[{gal_el}] not found");
383 }
384
385 acc.value = true;
386 }
387}