poulpy/poulpy-core/src/glwe_packing.rs
poulpy-core/src/glwe_packing.rsRaw
1use std::collections::HashMap;
2
3use poulpy_hal::{
4 api::{
5 DFT, IDFTConsume, ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxAddInplace, VecZnxAutomorphismInplace,
6 VecZnxBigAddSmallInplace, VecZnxBigAutomorphismInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
7 VecZnxBigSubSmallBInplace, VecZnxCopy, VecZnxDftAllocBytes, VecZnxNegateInplace, VecZnxNormalizeInplace, VecZnxRotate,
8 VecZnxRotateInplace, VecZnxRshInplace, VecZnxSub, VecZnxSubABInplace, VmpApply, VmpApplyAdd, VmpApplyTmpBytes,
9 },
10 layouts::{Backend, DataMut, DataRef, Module, Scratch},
11};
12
13use crate::{
14 GLWEOperations, TakeGLWECt,
15 layouts::{GLWECiphertext, Infos, prepared::GGLWEAutomorphismKeyPrepared},
16};
17
18/// [GLWEPacker] enables only the fly GLWE packing
19/// with constant memory of Log(N) ciphertexts.
20/// Main difference with usual GLWE packing is that
21/// the output is bit-reversed.
22pub struct GLWEPacker {
23 accumulators: Vec<Accumulator>,
24 log_batch: usize,
25 counter: usize,
26}
27
28/// [Accumulator] stores intermediate packing result.
29/// There are Log(N) such accumulators in a [GLWEPacker].
30struct Accumulator {
31 data: GLWECiphertext<Vec<u8>>,
32 value: bool, // Implicit flag for zero ciphertext
33 control: bool, // Can be combined with incoming value
34}
35
36impl Accumulator {
37 /// Allocates a new [Accumulator].
38 ///
39 /// #Arguments
40 ///
41 /// * `module`: static backend FFT tables.
42 /// * `basek`: base 2 logarithm of the GLWE ciphertext in memory digit representation.
43 /// * `k`: base 2 precision of the GLWE ciphertext precision over the Torus.
44 /// * `rank`: rank of the GLWE ciphertext.
45 pub fn alloc(n: usize, basek: usize, k: usize, rank: usize) -> Self {
46 Self {
47 data: GLWECiphertext::alloc(n, basek, k, rank),
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 /// * `basek`: base 2 logarithm of the GLWE ciphertext in memory digit representation.
66 /// * `k`: base 2 precision of the GLWE ciphertext precision over the Torus.
67 /// * `rank`: rank of the GLWE ciphertext.
68 pub fn new(n: usize, log_batch: usize, basek: usize, k: usize, rank: usize) -> Self {
69 let mut accumulators: Vec<Accumulator> = Vec::<Accumulator>::new();
70 let log_n: usize = (usize::BITS - (n - 1).leading_zeros()) as _;
71 (0..log_n - log_batch).for_each(|_| accumulators.push(Accumulator::alloc(n, basek, k, rank)));
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 scratch_space<B: Backend>(
90 module: &Module<B>,
91 basek: usize,
92 ct_k: usize,
93 k_ksk: usize,
94 digits: usize,
95 rank: usize,
96 ) -> usize
97 where
98 Module<B>: VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxBigNormalizeTmpBytes,
99 {
100 pack_core_scratch_space(module, basek, ct_k, k_ksk, digits, rank)
101 }
102
103 pub fn galois_elements<B: Backend>(module: &Module<B>) -> Vec<i64> {
104 GLWECiphertext::trace_galois_elements(module)
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`: a [HashMap] containing the [AutomorphismKeyExec]s.
115 /// * `scratch`: scratch space of size at least [Self::scratch_space].
116 pub fn add<DataA: DataRef, DataAK: DataRef, B: Backend>(
117 &mut self,
118 module: &Module<B>,
119 a: Option<&GLWECiphertext<DataA>>,
120 auto_keys: &HashMap<i64, GGLWEAutomorphismKeyPrepared<DataAK, B>>,
121 scratch: &mut Scratch<B>,
122 ) where
123 Module<B>: VecZnxDftAllocBytes
124 + VmpApplyTmpBytes
125 + VecZnxBigNormalizeTmpBytes
126 + VmpApply<B>
127 + VmpApplyAdd<B>
128 + DFT<B>
129 + IDFTConsume<B>
130 + VecZnxBigAddSmallInplace<B>
131 + VecZnxBigNormalize<B>
132 + VecZnxCopy
133 + VecZnxRotateInplace
134 + VecZnxSub
135 + VecZnxNegateInplace
136 + VecZnxRshInplace
137 + VecZnxAddInplace
138 + VecZnxNormalizeInplace<B>
139 + VecZnxSubABInplace
140 + VecZnxRotate
141 + VecZnxAutomorphismInplace
142 + VecZnxBigSubSmallBInplace<B>
143 + VecZnxBigAutomorphismInplace<B>,
144 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
145 {
146 assert!(
147 self.counter < self.accumulators[0].data.n(),
148 "Packing limit of {} reached",
149 self.accumulators[0].data.n() >> self.log_batch
150 );
151
152 pack_core(
153 module,
154 a,
155 &mut self.accumulators,
156 self.log_batch,
157 auto_keys,
158 scratch,
159 );
160 self.counter += 1 << self.log_batch;
161 }
162
163 /// Flush result to`res`.
164 pub fn flush<Data: DataMut, B: Backend>(&mut self, module: &Module<B>, res: &mut GLWECiphertext<Data>)
165 where
166 Module<B>: VecZnxCopy,
167 {
168 assert!(self.counter == self.accumulators[0].data.n());
169 // Copy result GLWE into res GLWE
170 res.copy(
171 module,
172 &self.accumulators[module.log_n() - self.log_batch - 1].data,
173 );
174
175 self.reset();
176 }
177}
178
179fn pack_core_scratch_space<B: Backend>(
180 module: &Module<B>,
181 basek: usize,
182 ct_k: usize,
183 k_ksk: usize,
184 digits: usize,
185 rank: usize,
186) -> usize
187where
188 Module<B>: VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxBigNormalizeTmpBytes,
189{
190 combine_scratch_space(module, basek, ct_k, k_ksk, digits, rank)
191}
192
193fn pack_core<D: DataRef, DataAK: DataRef, B: Backend>(
194 module: &Module<B>,
195 a: Option<&GLWECiphertext<D>>,
196 accumulators: &mut [Accumulator],
197 i: usize,
198 auto_keys: &HashMap<i64, GGLWEAutomorphismKeyPrepared<DataAK, B>>,
199 scratch: &mut Scratch<B>,
200) where
201 Module<B>: VecZnxDftAllocBytes
202 + VmpApplyTmpBytes
203 + VecZnxBigNormalizeTmpBytes
204 + VmpApply<B>
205 + VmpApplyAdd<B>
206 + DFT<B>
207 + IDFTConsume<B>
208 + VecZnxBigAddSmallInplace<B>
209 + VecZnxBigNormalize<B>
210 + VecZnxCopy
211 + VecZnxRotateInplace
212 + VecZnxSub
213 + VecZnxNegateInplace
214 + VecZnxRshInplace
215 + VecZnxAddInplace
216 + VecZnxNormalizeInplace<B>
217 + VecZnxSubABInplace
218 + VecZnxRotate
219 + VecZnxAutomorphismInplace
220 + VecZnxBigSubSmallBInplace<B>
221 + VecZnxBigAutomorphismInplace<B>,
222 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
223{
224 let log_n: usize = module.log_n();
225
226 if i == log_n {
227 return;
228 }
229
230 // Isolate the first accumulator
231 let (acc_prev, acc_next) = accumulators.split_at_mut(1);
232
233 // Control = true accumlator is free to overide
234 if !acc_prev[0].control {
235 let acc_mut_ref: &mut Accumulator = &mut acc_prev[0]; // from split_at_mut
236
237 // No previous value -> copies and sets flags accordingly
238 if let Some(a_ref) = a {
239 acc_mut_ref.data.copy(module, a_ref);
240 acc_mut_ref.value = true
241 } else {
242 acc_mut_ref.value = false
243 }
244 acc_mut_ref.control = true; // Able to be combined on next call
245 } else {
246 // Compresses acc_prev <- combine(acc_prev, a).
247 combine(module, &mut acc_prev[0], a, i, auto_keys, scratch);
248 acc_prev[0].control = false;
249
250 // Propagates to next accumulator
251 if acc_prev[0].value {
252 pack_core(
253 module,
254 Some(&acc_prev[0].data),
255 acc_next,
256 i + 1,
257 auto_keys,
258 scratch,
259 );
260 } else {
261 pack_core(
262 module,
263 None::<&GLWECiphertext<Vec<u8>>>,
264 acc_next,
265 i + 1,
266 auto_keys,
267 scratch,
268 );
269 }
270 }
271}
272
273fn combine_scratch_space<B: Backend>(
274 module: &Module<B>,
275 basek: usize,
276 ct_k: usize,
277 k_ksk: usize,
278 digits: usize,
279 rank: usize,
280) -> usize
281where
282 Module<B>: VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxBigNormalizeTmpBytes,
283{
284 GLWECiphertext::bytes_of(module.n(), basek, ct_k, rank)
285 + (GLWECiphertext::rsh_scratch_space(module.n())
286 | GLWECiphertext::automorphism_scratch_space(module, basek, ct_k, ct_k, k_ksk, digits, rank))
287}
288
289/// [combine] merges two ciphertexts together.
290fn combine<D: DataRef, DataAK: DataRef, B: Backend>(
291 module: &Module<B>,
292 acc: &mut Accumulator,
293 b: Option<&GLWECiphertext<D>>,
294 i: usize,
295 auto_keys: &HashMap<i64, GGLWEAutomorphismKeyPrepared<DataAK, B>>,
296 scratch: &mut Scratch<B>,
297) where
298 Module<B>: VecZnxDftAllocBytes
299 + VmpApplyTmpBytes
300 + VecZnxBigNormalizeTmpBytes
301 + VmpApply<B>
302 + VmpApplyAdd<B>
303 + DFT<B>
304 + IDFTConsume<B>
305 + VecZnxBigAddSmallInplace<B>
306 + VecZnxBigNormalize<B>
307 + VecZnxCopy
308 + VecZnxRotateInplace
309 + VecZnxSub
310 + VecZnxNegateInplace
311 + VecZnxRshInplace
312 + VecZnxAddInplace
313 + VecZnxNormalizeInplace<B>
314 + VecZnxSubABInplace
315 + VecZnxRotate
316 + VecZnxAutomorphismInplace
317 + VecZnxBigSubSmallBInplace<B>
318 + VecZnxBigAutomorphismInplace<B>,
319 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
320{
321 let n: usize = acc.data.n();
322 let log_n: usize = (u64::BITS - (n - 1).leading_zeros()) as _;
323 let a: &mut GLWECiphertext<Vec<u8>> = &mut acc.data;
324 let basek: usize = a.basek();
325 let k: usize = a.k();
326 let rank: usize = a.rank();
327
328 let gal_el: i64 = if i == 0 {
329 -1
330 } else {
331 module.galois_element(1 << (i - 1))
332 };
333
334 let t: i64 = 1 << (log_n - i - 1);
335
336 // Goal is to evaluate: a = a + b*X^t + phi(a - b*X^t))
337 // We also use the identity: AUTO(a * X^t, g) = -X^t * AUTO(a, g)
338 // where t = 2^(log_n - i - 1) and g = 5^{2^(i - 1)}
339 // Different cases for wether a and/or b are zero.
340 //
341 // Implicite RSH without modulus switch, introduces extra I(X) * Q/2 on decryption.
342 // Necessary so that the scaling of the plaintext remains constant.
343 // It however is ok to do so here because coefficients are eventually
344 // either mapped to garbage or twice their value which vanishes I(X)
345 // since 2*(I(X) * Q/2) = I(X) * Q = 0 mod Q.
346 if acc.value {
347 if let Some(b) = b {
348 let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(n, basek, k, rank);
349
350 // a = a * X^-t
351 a.rotate_inplace(module, -t);
352
353 // tmp_b = a * X^-t - b
354 tmp_b.sub(module, a, b);
355 tmp_b.rsh(module, 1);
356
357 // a = a * X^-t + b
358 a.add_inplace(module, b);
359 a.rsh(module, 1);
360
361 tmp_b.normalize_inplace(module, scratch_1);
362
363 // tmp_b = phi(a * X^-t - b)
364 if let Some(key) = auto_keys.get(&gal_el) {
365 tmp_b.automorphism_inplace(module, key, scratch_1);
366 } else {
367 panic!("auto_key[{}] not found", gal_el);
368 }
369
370 // a = a * X^-t + b - phi(a * X^-t - b)
371 a.sub_inplace_ab(module, &tmp_b);
372 a.normalize_inplace(module, scratch_1);
373
374 // a = a + b * X^t - phi(a * X^-t - b) * X^t
375 // = a + b * X^t - phi(a * X^-t - b) * - phi(X^t)
376 // = a + b * X^t + phi(a - b * X^t)
377 a.rotate_inplace(module, t);
378 } else {
379 a.rsh(module, 1);
380 // a = a + phi(a)
381 if let Some(key) = auto_keys.get(&gal_el) {
382 a.automorphism_add_inplace(module, key, scratch);
383 } else {
384 panic!("auto_key[{}] not found", gal_el);
385 }
386 }
387 } else if let Some(b) = b {
388 let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(n, basek, k, rank);
389 tmp_b.rotate(module, 1 << (log_n - i - 1), b);
390 tmp_b.rsh(module, 1);
391
392 // a = (b* X^t - phi(b* X^t))
393 if let Some(key) = auto_keys.get(&gal_el) {
394 a.automorphism_sub_ba(module, &tmp_b, key, scratch_1);
395 } else {
396 panic!("auto_key[{}] not found", gal_el);
397 }
398
399 acc.value = true;
400 }
401}