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, scratch: &mut Scratch<BE>)
136 where
137 R: GLWEToMut + GLWEInfos,
138 M: GLWEPackerOps<BE>,
139 Scratch<BE>: ScratchTakeCore<BE>,
140 {
141 assert!(self.counter as u32 == self.accumulators[0].data.n());
142
143 let out: &GLWE<Vec<u8>> = &self.accumulators[module.log_n() - self.log_batch - 1].data;
144
145 if out.base2k() == res.base2k() {
146 module.glwe_copy(res, out)
147 } else {
148 module.glwe_normalize(res, out, scratch);
149 }
150
151 self.reset();
152 }
153}
154
155impl<BE: Backend> GLWEPackerOps<BE> for Module<BE> where
156 Self: Sized
157 + ModuleLogN
158 + GLWEAutomorphism<BE>
159 + GaloisElement
160 + GLWERotate<BE>
161 + GLWESub
162 + GLWEShift<BE>
163 + GLWEAdd
164 + GLWENormalize<BE>
165 + GLWECopy
166 + GLWEAutomorphism<BE>
167 + GaloisElement
168 + GLWERotate<BE>
169 + GLWESub
170 + GLWEShift<BE>
171 + GLWEAdd
172 + GLWENormalize<BE>
173{
174}
175
176pub trait GLWEPackerOps<BE: Backend>
177where
178 Self: Sized
179 + ModuleLogN
180 + GLWEAutomorphism<BE>
181 + GaloisElement
182 + GLWERotate<BE>
183 + GLWESub
184 + GLWEShift<BE>
185 + GLWEAdd
186 + GLWENormalize<BE>
187 + GLWECopy
188 + GLWEAutomorphism<BE>
189 + GaloisElement
190 + GLWERotate<BE>
191 + GLWESub
192 + GLWEShift<BE>
193 + GLWEAdd
194 + GLWENormalize<BE>,
195{
196 fn packer_add<A, K, H>(&self, packer: &mut GLWEPacker, a: Option<&A>, i: usize, auto_keys: &H, scratch: &mut Scratch<BE>)
197 where
198 A: GLWEToRef + GLWEInfos,
199 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
200 H: GLWEAutomorphismKeyHelper<K, BE>,
201 Scratch<BE>: ScratchTakeCore<BE>,
202 {
203 pack_core(self, a, &mut packer.accumulators, i, auto_keys, scratch)
204 }
205}
206
207fn pack_core<A, K, H, M, BE: Backend>(
208 module: &M,
209 a: Option<&A>,
210 accumulators: &mut [Accumulator],
211 i: usize,
212 auto_keys: &H,
213 scratch: &mut Scratch<BE>,
214) where
215 A: GLWEToRef + GLWEInfos,
216 M: ModuleLogN
217 + GLWEAutomorphism<BE>
218 + GaloisElement
219 + GLWERotate<BE>
220 + GLWESub
221 + GLWEShift<BE>
222 + GLWEAdd
223 + GLWENormalize<BE>
224 + GLWECopy
225 + GLWEAutomorphism<BE>
226 + GaloisElement
227 + GLWERotate<BE>
228 + GLWESub
229 + GLWEShift<BE>
230 + GLWEAdd
231 + GLWENormalize<BE>,
232 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
233 H: GLWEAutomorphismKeyHelper<K, BE>,
234 Scratch<BE>: ScratchTakeCore<BE>,
235{
236 let log_n: usize = module.log_n();
237
238 if i == log_n {
239 return;
240 }
241
242 // Isolate the first accumulator
243 let (acc_prev, acc_next) = accumulators.split_at_mut(1);
244
245 // Control = true accumlator is free to overide
246 if !acc_prev[0].control {
247 let acc_mut_ref: &mut Accumulator = &mut acc_prev[0]; // from split_at_mut
248
249 // No previous value -> copies and sets flags accordingly
250 if let Some(a_ref) = a {
251 if a_ref.base2k() == acc_mut_ref.data.base2k() {
252 module.glwe_copy(&mut acc_mut_ref.data, a_ref);
253 } else {
254 module.glwe_normalize(&mut acc_mut_ref.data, a_ref, scratch);
255 }
256 acc_mut_ref.value = true
257 } else {
258 acc_mut_ref.value = false
259 }
260 acc_mut_ref.control = true; // Able to be combined on next call
261 } else {
262 // Compresses acc_prev <- combine(acc_prev, a).
263 combine(module, &mut acc_prev[0], a, i, auto_keys, scratch);
264 acc_prev[0].control = false;
265
266 // Propagates to next accumulator
267 if acc_prev[0].value {
268 pack_core(module, Some(&acc_prev[0].data), acc_next, i + 1, auto_keys, scratch);
269 } else {
270 pack_core(module, None::<&GLWE<Vec<u8>>>, acc_next, i + 1, auto_keys, scratch);
271 }
272 }
273}
274
275fn combine<B, K, H, M, BE: Backend>(
276 module: &M,
277 acc: &mut Accumulator,
278 b: Option<&B>,
279 i: usize,
280 auto_keys: &H,
281 scratch: &mut Scratch<BE>,
282) where
283 B: GLWEToRef + GLWEInfos,
284 B: GLWEToRef + GLWEInfos,
285 M: ModuleLogN
286 + GLWEAutomorphism<BE>
287 + GaloisElement
288 + GLWERotate<BE>
289 + GLWESub
290 + GLWEShift<BE>
291 + GLWEAdd
292 + GLWENormalize<BE>
293 + GLWECopy
294 + GLWEAutomorphism<BE>
295 + GaloisElement
296 + GLWERotate<BE>
297 + GLWESub
298 + GLWEShift<BE>
299 + GLWEAdd
300 + GLWENormalize<BE>,
301 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
302 H: GLWEAutomorphismKeyHelper<K, BE>,
303 Scratch<BE>: ScratchTakeCore<BE>,
304{
305 let log_n: usize = acc.data.n().log2();
306 let a: &mut GLWE<Vec<u8>> = &mut acc.data;
307
308 let gal_el: i64 = if i == 0 { -1 } else { module.galois_element(1 << (i - 1)) };
309
310 let t: i64 = 1 << (log_n - i - 1);
311
312 // Goal is to evaluate: a = a + b*X^t + phi(a - b*X^t))
313 // We also use the identity: AUTO(a * X^t, g) = -X^t * AUTO(a, g)
314 // where t = 2^(log_n - i - 1) and g = 5^{2^(i - 1)}
315 // Different cases for wether a and/or b are zero.
316 //
317 // Implicite RSH without modulus switch, introduces extra I(X) * Q/2 on decryption.
318 // Necessary so that the scaling of the plaintext remains constant.
319 // It however is ok to do so here because coefficients are eventually
320 // either mapped to garbage or twice their value which vanishes I(X)
321 // since 2*(I(X) * Q/2) = I(X) * Q = 0 mod Q.
322 if acc.value {
323 if let Some(b) = b {
324 let (mut tmp, scratch_1) = scratch.take_glwe(a);
325
326 // a = a * X^-t
327 module.glwe_rotate_inplace(-t, a, scratch_1);
328
329 // tmp_b = a * X^-t - b
330 module.glwe_sub(&mut tmp, a, b);
331 module.glwe_rsh(1, &mut tmp, scratch_1);
332 // a = a * X^-t + b
333 module.glwe_add_inplace(a, b);
334
335 module.glwe_rsh(1, a, scratch_1);
336 module.glwe_normalize_inplace(&mut tmp, scratch_1);
337
338 // tmp_b = phi(a * X^-t - b)
339 if let Some(auto_key) = auto_keys.get_automorphism_key(gal_el) {
340 module.glwe_automorphism_inplace(&mut tmp, auto_key, scratch_1);
341 } else {
342 panic!("auto_key[{gal_el}] not found");
343 }
344
345 // a = a * X^-t + b - phi(a * X^-t - b)
346 module.glwe_sub_inplace(a, &tmp);
347 module.glwe_normalize_inplace(a, scratch_1);
348
349 // a = a + b * X^t - phi(a * X^-t - b) * X^t
350 // = a + b * X^t - phi(a * X^-t - b) * - phi(X^t)
351 // = a + b * X^t + phi(a - b * X^t)
352 module.glwe_rotate_inplace(t, a, scratch_1);
353 } else {
354 module.glwe_rsh(1, a, scratch);
355 // a = a + phi(a)
356 if let Some(auto_key) = auto_keys.get_automorphism_key(gal_el) {
357 module.glwe_automorphism_add_inplace(a, auto_key, scratch);
358 } else {
359 panic!("auto_key[{gal_el}] not found");
360 }
361 }
362 } else if let Some(b) = b {
363 let (mut tmp_b, scratch_1) = scratch.take_glwe(a);
364 module.glwe_rotate(t, &mut tmp_b, b);
365 module.glwe_rsh(1, &mut tmp_b, scratch_1);
366
367 // a = (b* X^t - phi(b* X^t))
368 if let Some(auto_key) = auto_keys.get_automorphism_key(gal_el) {
369 module.glwe_automorphism_sub_negate(a, &tmp_b, auto_key, scratch_1);
370 } else {
371 panic!("auto_key[{gal_el}] not found");
372 }
373
374 acc.value = true;
375 }
376}