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