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