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