poulpy_core/
glwe_packer.rs

1use poulpy_hal::{
2    api::ModuleLogN,
3    layouts::{Backend, GaloisElement, Module, Scratch},
4};
5
6use crate::{
7    GLWEAdd, GLWEAutomorphism, GLWECopy, GLWENormalize, GLWERotate, GLWEShift, GLWESub, ScratchTakeCore,
8    glwe_trace::GLWETrace,
9    layouts::{
10        GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEAutomorphismKeyHelper, GLWEInfos, GLWEToMut, GLWEToRef, GetGaloisElement,
11        LWEInfos,
12    },
13};
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: GLWE<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    /// * `base2k`: 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<A>(infos: &A) -> Self
43    where
44        A: GLWEInfos,
45    {
46        Self {
47            data: GLWE::alloc_from_infos(infos),
48            value: false,
49            control: false,
50        }
51    }
52}
53
54impl GLWEPacker {
55    /// Instantiates a new [GLWEPacker].
56    ///
57    /// # Arguments
58    ///
59    /// * `log_batch`: packs coefficients which are multiples of X^{N/2^log_batch}.
60    ///   i.e. with `log_batch=0` only the constant coefficient is packed
61    ///   and N GLWE ciphertext can be packed. With `log_batch=2` all coefficients
62    ///   which are multiples of X^{N/4} are packed. Meaning that N/4 ciphertexts
63    ///   can be packed.
64    pub fn alloc<A>(infos: &A, log_batch: usize) -> Self
65    where
66        A: GLWEInfos,
67    {
68        let mut accumulators: Vec<Accumulator> = Vec::<Accumulator>::new();
69        let log_n: usize = infos.n().log2();
70        (0..log_n - log_batch).for_each(|_| accumulators.push(Accumulator::alloc(infos)));
71        GLWEPacker {
72            accumulators,
73            log_batch,
74            counter: 0,
75        }
76    }
77
78    /// Implicit reset of the internal state (to be called before a new packing procedure).
79    fn reset(&mut self) {
80        for i in 0..self.accumulators.len() {
81            self.accumulators[i].value = false;
82            self.accumulators[i].control = false;
83        }
84        self.counter = 0;
85    }
86
87    /// Number of scratch space bytes required to call [Self::add].
88    pub fn tmp_bytes<R, K, M, BE: Backend>(module: &M, res_infos: &R, key_infos: &K) -> usize
89    where
90        R: GLWEInfos,
91        K: GGLWEInfos,
92        M: GLWEPackerOps<BE>,
93    {
94        GLWE::bytes_of_from_infos(res_infos)
95            + module
96                .glwe_rsh_tmp_byte()
97                .max(module.glwe_automorphism_tmp_bytes(res_infos, res_infos, key_infos))
98    }
99
100    pub fn galois_elements<M, BE: Backend>(module: &M) -> Vec<i64>
101    where
102        M: GLWETrace<BE>,
103    {
104        module.glwe_trace_galois_elements()
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`: an implementation of [GLWEAutomorphismKeyHelper], containing [GLWEAutomorphismKeyPrepared] with index of [Self::galois_elements].
115    /// * `scratch`: scratch space of size at least [Self::tmp_bytes].
116    pub fn add<A, K, H, M, BE: Backend>(&mut self, module: &M, a: Option<&A>, auto_keys: &H, scratch: &mut Scratch<BE>)
117    where
118        A: GLWEToRef + GLWEInfos,
119        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
120        H: GLWEAutomorphismKeyHelper<K, BE>,
121        M: GLWEPackerOps<BE>,
122        Scratch<BE>: ScratchTakeCore<BE>,
123    {
124        assert!(
125            (self.counter as u32) < self.accumulators[0].data.n(),
126            "Packing limit of {} reached",
127            self.accumulators[0].data.n().0 as usize >> self.log_batch
128        );
129
130        module.packer_add(self, a, self.log_batch, auto_keys, scratch);
131        self.counter += 1 << self.log_batch;
132    }
133
134    /// Flush result to`res`.
135    pub fn flush<R, M, BE: Backend>(&mut self, module: &M, res: &mut R, scratch: &mut Scratch<BE>)
136    where
137        R: GLWEToMut + GLWEInfos,
138        M: GLWEPackerOps<BE>,
139        Scratch<BE>: ScratchTakeCore<BE>,
140    {
141        assert!(self.counter as u32 == self.accumulators[0].data.n());
142
143        let out: &GLWE<Vec<u8>> = &self.accumulators[module.log_n() - self.log_batch - 1].data;
144
145        if out.base2k() == res.base2k() {
146            module.glwe_copy(res, out)
147        } else {
148            module.glwe_normalize(res, out, scratch);
149        }
150
151        self.reset();
152    }
153}
154
155impl<BE: Backend> GLWEPackerOps<BE> for Module<BE> where
156    Self: Sized
157        + ModuleLogN
158        + GLWEAutomorphism<BE>
159        + GaloisElement
160        + GLWERotate<BE>
161        + GLWESub
162        + GLWEShift<BE>
163        + GLWEAdd
164        + GLWENormalize<BE>
165        + GLWECopy
166        + GLWEAutomorphism<BE>
167        + GaloisElement
168        + GLWERotate<BE>
169        + GLWESub
170        + GLWEShift<BE>
171        + GLWEAdd
172        + GLWENormalize<BE>
173{
174}
175
176pub trait GLWEPackerOps<BE: Backend>
177where
178    Self: Sized
179        + ModuleLogN
180        + GLWEAutomorphism<BE>
181        + GaloisElement
182        + GLWERotate<BE>
183        + GLWESub
184        + GLWEShift<BE>
185        + GLWEAdd
186        + GLWENormalize<BE>
187        + GLWECopy
188        + GLWEAutomorphism<BE>
189        + GaloisElement
190        + GLWERotate<BE>
191        + GLWESub
192        + GLWEShift<BE>
193        + GLWEAdd
194        + GLWENormalize<BE>,
195{
196    fn packer_add<A, K, H>(&self, packer: &mut GLWEPacker, a: Option<&A>, i: usize, auto_keys: &H, scratch: &mut Scratch<BE>)
197    where
198        A: GLWEToRef + GLWEInfos,
199        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
200        H: GLWEAutomorphismKeyHelper<K, BE>,
201        Scratch<BE>: ScratchTakeCore<BE>,
202    {
203        pack_core(self, a, &mut packer.accumulators, i, auto_keys, scratch)
204    }
205}
206
207fn pack_core<A, K, H, M, BE: Backend>(
208    module: &M,
209    a: Option<&A>,
210    accumulators: &mut [Accumulator],
211    i: usize,
212    auto_keys: &H,
213    scratch: &mut Scratch<BE>,
214) where
215    A: GLWEToRef + GLWEInfos,
216    M: ModuleLogN
217        + GLWEAutomorphism<BE>
218        + GaloisElement
219        + GLWERotate<BE>
220        + GLWESub
221        + GLWEShift<BE>
222        + GLWEAdd
223        + GLWENormalize<BE>
224        + GLWECopy
225        + GLWEAutomorphism<BE>
226        + GaloisElement
227        + GLWERotate<BE>
228        + GLWESub
229        + GLWEShift<BE>
230        + GLWEAdd
231        + GLWENormalize<BE>,
232    K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
233    H: GLWEAutomorphismKeyHelper<K, BE>,
234    Scratch<BE>: ScratchTakeCore<BE>,
235{
236    let log_n: usize = module.log_n();
237
238    if i == log_n {
239        return;
240    }
241
242    // Isolate the first accumulator
243    let (acc_prev, acc_next) = accumulators.split_at_mut(1);
244
245    // Control = true accumlator is free to overide
246    if !acc_prev[0].control {
247        let acc_mut_ref: &mut Accumulator = &mut acc_prev[0]; // from split_at_mut
248
249        // No previous value -> copies and sets flags accordingly
250        if let Some(a_ref) = a {
251            if a_ref.base2k() == acc_mut_ref.data.base2k() {
252                module.glwe_copy(&mut acc_mut_ref.data, a_ref);
253            } else {
254                module.glwe_normalize(&mut acc_mut_ref.data, a_ref, scratch);
255            }
256            acc_mut_ref.value = true
257        } else {
258            acc_mut_ref.value = false
259        }
260        acc_mut_ref.control = true; // Able to be combined on next call
261    } else {
262        // Compresses acc_prev <- combine(acc_prev, a).
263        combine(module, &mut acc_prev[0], a, i, auto_keys, scratch);
264        acc_prev[0].control = false;
265
266        // Propagates to next accumulator
267        if acc_prev[0].value {
268            pack_core(
269                module,
270                Some(&acc_prev[0].data),
271                acc_next,
272                i + 1,
273                auto_keys,
274                scratch,
275            );
276        } else {
277            pack_core(
278                module,
279                None::<&GLWE<Vec<u8>>>,
280                acc_next,
281                i + 1,
282                auto_keys,
283                scratch,
284            );
285        }
286    }
287}
288
289fn combine<B, K, H, M, BE: Backend>(
290    module: &M,
291    acc: &mut Accumulator,
292    b: Option<&B>,
293    i: usize,
294    auto_keys: &H,
295    scratch: &mut Scratch<BE>,
296) where
297    B: GLWEToRef + GLWEInfos,
298    B: GLWEToRef + GLWEInfos,
299    M: ModuleLogN
300        + GLWEAutomorphism<BE>
301        + GaloisElement
302        + GLWERotate<BE>
303        + GLWESub
304        + GLWEShift<BE>
305        + GLWEAdd
306        + GLWENormalize<BE>
307        + GLWECopy
308        + GLWEAutomorphism<BE>
309        + GaloisElement
310        + GLWERotate<BE>
311        + GLWESub
312        + GLWEShift<BE>
313        + GLWEAdd
314        + GLWENormalize<BE>,
315    K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
316    H: GLWEAutomorphismKeyHelper<K, BE>,
317    Scratch<BE>: ScratchTakeCore<BE>,
318{
319    let log_n: usize = acc.data.n().log2();
320    let a: &mut GLWE<Vec<u8>> = &mut acc.data;
321
322    let gal_el: i64 = if i == 0 {
323        -1
324    } else {
325        module.galois_element(1 << (i - 1))
326    };
327
328    let t: i64 = 1 << (log_n - i - 1);
329
330    // Goal is to evaluate: a = a + b*X^t + phi(a - b*X^t))
331    // We also use the identity: AUTO(a * X^t, g) = -X^t * AUTO(a, g)
332    // where t = 2^(log_n - i - 1) and g = 5^{2^(i - 1)}
333    // Different cases for wether a and/or b are zero.
334    //
335    // Implicite RSH without modulus switch, introduces extra I(X) * Q/2 on decryption.
336    // Necessary so that the scaling of the plaintext remains constant.
337    // It however is ok to do so here because coefficients are eventually
338    // either mapped to garbage or twice their value which vanishes I(X)
339    // since 2*(I(X) * Q/2) = I(X) * Q = 0 mod Q.
340    if acc.value {
341        if let Some(b) = b {
342            let (mut tmp, scratch_1) = scratch.take_glwe(a);
343
344            // a = a * X^-t
345            module.glwe_rotate_inplace(-t, a, scratch_1);
346
347            // tmp_b = a * X^-t - b
348            module.glwe_sub(&mut tmp, a, b);
349            module.glwe_rsh(1, &mut tmp, scratch_1);
350            // a = a * X^-t + b
351            module.glwe_add_inplace(a, b);
352
353            module.glwe_rsh(1, a, scratch_1);
354            module.glwe_normalize_inplace(&mut tmp, scratch_1);
355
356            // tmp_b = phi(a * X^-t - b)
357            if let Some(auto_key) = auto_keys.get_automorphism_key(gal_el) {
358                module.glwe_automorphism_inplace(&mut tmp, auto_key, scratch_1);
359            } else {
360                panic!("auto_key[{gal_el}] not found");
361            }
362
363            // a = a * X^-t + b - phi(a * X^-t - b)
364            module.glwe_sub_inplace(a, &tmp);
365            module.glwe_normalize_inplace(a, scratch_1);
366
367            // a = a + b * X^t - phi(a * X^-t - b) * X^t
368            //   = a + b * X^t - phi(a * X^-t - b) * - phi(X^t)
369            //   = a + b * X^t + phi(a - b * X^t)
370            module.glwe_rotate_inplace(t, a, scratch_1);
371        } else {
372            module.glwe_rsh(1, a, scratch);
373            // a = a + phi(a)
374            if let Some(auto_key) = auto_keys.get_automorphism_key(gal_el) {
375                module.glwe_automorphism_add_inplace(a, auto_key, scratch);
376            } else {
377                panic!("auto_key[{gal_el}] not found");
378            }
379        }
380    } else if let Some(b) = b {
381        let (mut tmp_b, scratch_1) = scratch.take_glwe(a);
382        module.glwe_rotate(t, &mut tmp_b, b);
383        module.glwe_rsh(1, &mut tmp_b, scratch_1);
384
385        // a = (b* X^t - phi(b* X^t))
386        if let Some(auto_key) = auto_keys.get_automorphism_key(gal_el) {
387            module.glwe_automorphism_sub_negate(a, &tmp_b, auto_key, scratch_1);
388        } else {
389            panic!("auto_key[{gal_el}] not found");
390        }
391
392        acc.value = true;
393    }
394}
poulpy_core/glwe_packer.rs

poulpy_core/
glwe_packer.rs
