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