Added support for packing GLWE

core/src/glwe_packing.rs

@@ -0,0 +1,317 @@
+use crate::{automorphism::AutomorphismKey, elem::Infos, glwe_ciphertext::GLWECiphertext, glwe_ops::GLWEOps};
+use std::collections::HashMap;
+
+use backend::{FFT64, Module, Scratch, VecZnxAlloc};
+
+/// [StreamPacker] enables only the fly GLWE packing
+/// with constant memory of Log(N) ciphertexts.
+/// Main difference with usual GLWE packing is that
+/// the output is bit-reversed.
+pub struct StreamPacker {
+    accumulators: Vec<Accumulator>,
+    log_batch: usize,
+    counter: usize,
+}
+
+/// [Accumulator] stores intermediate packing result.
+/// There are Log(N) such accumulators in a [StreamPacker].
+struct Accumulator {
+    data: GLWECiphertext<Vec<u8>>,
+    value: bool,   // Implicit flag for zero ciphertext
+    control: bool, // Can be combined with incoming value
+}
+
+impl Accumulator {
+    /// Allocates a new [Accumulator].
+    ///
+    /// #Arguments
+    ///
+    /// * `module`: static backend FFT tables.
+    /// * `basek`: base 2 logarithm of the GLWE ciphertext in memory digit representation.
+    /// * `k`: base 2 precision of the GLWE ciphertext precision over the Torus.
+    /// * `rank`: rank of the GLWE ciphertext.
+    pub fn alloc(module: &Module<FFT64>, basek: usize, k: usize, rank: usize) -> Self {
+        Self {
+            data: GLWECiphertext::alloc(module, basek, k, rank),
+            value: false,
+            control: false,
+        }
+    }
+}
+
+impl StreamPacker {
+    /// Instantiates a new [StreamPacker].
+    ///
+    /// #Arguments
+    ///
+    /// * `module`: static backend FFT tables.
+    /// * `log_batch`: packs coefficients which are multiples of X^{N/2^log_batch}.
+    ///                i.e. with `log_batch=0` only the constant coefficient is packed
+    ///                and N GLWE ciphertext can be packed. With `log_batch=2` all coefficients
+    ///                which are multiples of X^{N/4} are packed. Meaning that N/4 ciphertexts
+    ///                can be packed.
+    /// * `basek`: base 2 logarithm of the GLWE ciphertext in memory digit representation.
+    /// * `k`: base 2 precision of the GLWE ciphertext precision over the Torus.
+    /// * `rank`: rank of the GLWE ciphertext.
+    pub fn new(module: &Module<FFT64>, log_batch: usize, basek: usize, k: usize, rank: usize) -> Self {
+        let mut accumulators: Vec<Accumulator> = Vec::<Accumulator>::new();
+        let log_n: usize = module.log_n();
+        (0..log_n - log_batch).for_each(|_| accumulators.push(Accumulator::alloc(module, basek, k, rank)));
+        Self {
+            accumulators: accumulators,
+            log_batch,
+            counter: 0,
+        }
+    }
+
+    /// Implicit reset of the internal state (to be called before a new packing procedure).
+    pub fn reset(&mut self) {
+        for i in 0..self.accumulators.len() {
+            self.accumulators[i].value = false;
+            self.accumulators[i].control = false;
+        }
+        self.counter = 0;
+    }
+
+    /// Number of scratch space bytes required to call [Self::add].
+    pub fn scratch_space(module: &Module<FFT64>, ct_size: usize, autokey_size: usize, rank: usize) -> usize {
+        pack_core_scratch_space(module, ct_size, autokey_size, rank)
+    }
+
+    pub fn galois_elements(module: &Module<FFT64>) -> Vec<i64> {
+        GLWECiphertext::trace_galois_elements(module)
+    }
+
+    /// Adds a GLWE ciphertext to the [StreamPacker]. And propagates
+    /// intermediate results among the [Accumulator]s.
+    ///
+    /// #Arguments
+    ///
+    /// * `module`: static backend FFT tables.
+    /// * `res`: space to append fully packed ciphertext. Only when the number
+    ///          of packed ciphertexts reaches N/2^log_batch is a result written.
+    /// * `a`: ciphertext to pack. Can optionally give None to pack a 0 ciphertext.
+    /// * `auto_keys`: a [HashMap] containing the [AutomorphismKey]s.
+    /// * `scratch`: scratch space of size at least [Self::add_scratch_space].
+    pub fn add<DataA: AsRef<[u8]>, DataAK: AsRef<[u8]>>(
+        &mut self,
+        module: &Module<FFT64>,
+        res: &mut Vec<GLWECiphertext<Vec<u8>>>,
+        a: Option<&GLWECiphertext<DataA>>,
+        auto_keys: &HashMap<i64, AutomorphismKey<DataAK, FFT64>>,
+        scratch: &mut Scratch,
+    ) {
+        pack_core(
+            module,
+            a,
+            &mut self.accumulators,
+            self.log_batch,
+            auto_keys,
+            scratch,
+        );
+        self.counter += 1 << self.log_batch;
+        if self.counter == module.n() {
+            res.push(
+                self.accumulators[module.log_n() - self.log_batch - 1]
+                    .data
+                    .clone(),
+            );
+            self.reset();
+        }
+    }
+
+    /// Flushes all accumlators and appends the result to `res`.
+    pub fn flush<DataAK: AsRef<[u8]>>(
+        &mut self,
+        module: &Module<FFT64>,
+        res: &mut Vec<GLWECiphertext<Vec<u8>>>,
+        auto_keys: &HashMap<i64, AutomorphismKey<DataAK, FFT64>>,
+        scratch: &mut Scratch,
+    ) {
+        if self.counter != 0 {
+            while self.counter != 0 {
+                self.add(
+                    module,
+                    res,
+                    None::<&GLWECiphertext<Vec<u8>>>,
+                    auto_keys,
+                    scratch,
+                );
+            }
+        }
+    }
+}
+
+fn pack_core_scratch_space(module: &Module<FFT64>, ct_size: usize, autokey_size: usize, rank: usize) -> usize {
+    combine_scratch_space(module, ct_size, autokey_size, rank)
+}
+
+fn pack_core<D: AsRef<[u8]>, DataAK: AsRef<[u8]>>(
+    module: &Module<FFT64>,
+    a: Option<&GLWECiphertext<D>>,
+    accumulators: &mut [Accumulator],
+    i: usize,
+    auto_keys: &HashMap<i64, AutomorphismKey<DataAK, FFT64>>,
+    scratch: &mut Scratch,
+) {
+    let log_n: usize = module.log_n();
+
+    if i == log_n {
+        return;
+    }
+
+    // Isolate the first accumulator
+    let (acc_prev, acc_next) = accumulators.split_at_mut(1);
+
+    // Control = true accumlator is free to overide
+    if !acc_prev[0].control {
+        let acc_mut_ref: &mut Accumulator = &mut acc_prev[0]; // from split_at_mut
+
+        // No previous value -> copies and sets flags accordingly
+        if let Some(a_ref) = a {
+            acc_mut_ref.data.copy(module, a_ref);
+            acc_mut_ref.value = true
+        } else {
+            acc_mut_ref.value = false
+        }
+        acc_mut_ref.control = true; // Able to be combined on next call
+    } else {
+        // Compresses acc_prev <- combine(acc_prev, a).
+        combine(module, &mut acc_prev[0], a, i, auto_keys, scratch);
+        acc_prev[0].control = false;
+
+        // Propagates to next accumulator
+        if acc_prev[0].value {
+            pack_core(
+                module,
+                Some(&acc_prev[0].data),
+                acc_next,
+                i + 1,
+                auto_keys,
+                scratch,
+            );
+        } else {
+            pack_core(
+                module,
+                None::<&GLWECiphertext<Vec<u8>>>,
+                acc_next,
+                i + 1,
+                auto_keys,
+                scratch,
+            );
+        }
+    }
+}
+
+fn combine_scratch_space(module: &Module<FFT64>, ct_size: usize, autokey_size: usize, rank: usize) -> usize {
+    2 * module.bytes_of_vec_znx(rank + 1, ct_size)
+        + (GLWECiphertext::rsh_scratch_space(module)
+            | GLWECiphertext::automorphism_scratch_space(module, ct_size, rank, ct_size, autokey_size))
+}
+
+/// [combine] merges two ciphertexts together.
+fn combine<D: AsRef<[u8]>, DataAK: AsRef<[u8]>>(
+    module: &Module<FFT64>,
+    acc: &mut Accumulator,
+    b: Option<&GLWECiphertext<D>>,
+    i: usize,
+    auto_keys: &HashMap<i64, AutomorphismKey<DataAK, FFT64>>,
+    scratch: &mut Scratch,
+) {
+    let log_n: usize = module.log_n();
+    let a: &mut GLWECiphertext<Vec<u8>> = &mut acc.data;
+    let basek: usize = a.basek();
+    let k: usize = a.k();
+    let rank: usize = a.rank();
+    let cols: usize = rank + 1;
+    let size: usize = a.size();
+
+    let gal_el: i64;
+
+    if i == 0 {
+        gal_el = -1;
+    } else {
+        gal_el = module.galois_element(1 << (i - 1))
+    }
+
+    // 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))
+    // Different cases for wether a and/or b are zero.
+    if acc.value {
+        // Implicite RSH without modulus switch, introduces extra I(X) * Q/2 on decryption.
+        // Necessary so that the scaling of the plaintext remains constant.
+        // It however is ok to do so here because coefficients are eventually
+        // either mapped to garbage or twice their value which vanishes I(X)
+        // since 2*(I(X) * Q/2) = I(X) * Q = 0 mod Q.
+        a.rsh(1, scratch);
+
+        if let Some(b) = b {
+            let (tmp_b_data, scratch_1) = scratch.tmp_vec_znx(module, cols, size);
+            let mut tmp_b: GLWECiphertext<&mut [u8]> = GLWECiphertext {
+                data: tmp_b_data,
+                k: k,
+                basek: basek,
+            };
+
+            {
+                let (tmp_a_data, scratch_2) = scratch_1.tmp_vec_znx(module, cols, size);
+                let mut tmp_a: GLWECiphertext<&mut [u8]> = GLWECiphertext {
+                    data: tmp_a_data,
+                    k: k,
+                    basek: basek,
+                };
+
+                // tmp_a = b * X^t
+                tmp_a.rotate(module, 1 << (log_n - i - 1), b);
+
+                // tmp_a >>= 1
+                tmp_a.rsh(1, scratch_2);
+
+                // tmp_b = a - b*X^t
+                tmp_b.sub(module, a, &tmp_a);
+                tmp_b.normalize_inplace(module, scratch_2);
+
+                // a = a + b * X^t
+                a.add_inplace(module, &tmp_a);
+            }
+
+            // tmp_b = phi(a - b * X^t)
+            if let Some(key) = auto_keys.get(&gal_el) {
+                tmp_b.automorphism_inplace(module, key, scratch_1);
+            } else {
+                panic!("auto_key[{}] not found", gal_el);
+            }
+
+            // a = a + b*X^t + phi(a - b*X^t))
+            a.add_inplace(module, &tmp_b);
+            a.normalize_inplace(module, scratch_1);
+        } else {
+            // a = a + phi(a)
+            if let Some(key) = auto_keys.get(&gal_el) {
+                a.automorphism_add_inplace(module, key, scratch);
+            } else {
+                panic!("auto_key[{}] not found", gal_el);
+            }
+        }
+    } else {
+        if let Some(b) = b {
+            let (tmp_b_data, scratch_1) = scratch.tmp_vec_znx(module, cols, size);
+            let mut tmp_b: GLWECiphertext<&mut [u8]> = GLWECiphertext {
+                data: tmp_b_data,
+                k: k,
+                basek: basek,
+            };
+
+            tmp_b.rotate(module, 1 << (log_n - i - 1), b);
+            tmp_b.rsh(1, scratch_1);
+
+            // a = (b* X^t - phi(b* X^t))
+            if let Some(key) = auto_keys.get(&gal_el) {
+                a.automorphism_sub_ba::<&mut [u8], _>(module, &tmp_b, key, scratch_1);
+            } else {
+                panic!("auto_key[{}] not found", gal_el);
+            }
+
+            acc.value = true;
+        }
+    }
+}

core/src/test_fft64/glwe_packing.rs

@@ -0,0 +1,158 @@
+use crate::{
+    automorphism::AutomorphismKey,
+    glwe_ciphertext::GLWECiphertext,
+    glwe_ops::GLWEOps,
+    glwe_packing::StreamPacker,
+    glwe_plaintext::GLWEPlaintext,
+    keys::{SecretKey, SecretKeyFourier},
+};
+use std::collections::HashMap;
+
+use backend::{Encoding, FFT64, Module, ScratchOwned, Stats};
+use sampling::source::Source;
+use std::time::Instant;
+
+#[test]
+fn packing() {
+    let log_n: usize = 5;
+    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
+
+    let mut source_xs: Source = Source::new([0u8; 32]);
+    let mut source_xe: Source = Source::new([0u8; 32]);
+    let mut source_xa: Source = Source::new([0u8; 32]);
+
+    let basek: usize = 18;
+    let k_ct: usize = 36;
+    let k_auto_key: usize = k_ct + basek;
+    let k_pt: usize = 18;
+    let rank: usize = 3;
+    let rows: usize = (k_ct + basek - 1) / basek;
+    let sigma: f64 = 3.2;
+    let ct_size: usize = rows;
+    let auto_key_size: usize = (k_auto_key + basek - 1) / basek;
+
+    let mut scratch: ScratchOwned = ScratchOwned::new(
+        GLWECiphertext::encrypt_sk_scratch_space(&module, ct_size)
+            | GLWECiphertext::decrypt_scratch_space(&module, ct_size)
+            | AutomorphismKey::generate_from_sk_scratch_space(&module, rank, auto_key_size)
+            | StreamPacker::scratch_space(&module, ct_size, auto_key_size, rank),
+    );
+
+    let mut sk: SecretKey<Vec<u8>> = SecretKey::alloc(&module, rank);
+    sk.fill_ternary_prob(0.5, &mut source_xs);
+
+    let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::alloc(&module, rank);
+    sk_dft.dft(&module, &sk);
+
+    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
+    let mut data: Vec<i64> = vec![0i64; module.n()];
+    data.iter_mut().enumerate().for_each(|(i, x)| {
+        *x = i as i64;
+    });
+    pt.data.encode_vec_i64(0, basek, k_pt, &data, 32);
+
+    let gal_els: Vec<i64> = StreamPacker::galois_elements(&module);
+
+    let mut auto_keys: HashMap<i64, AutomorphismKey<Vec<u8>, FFT64>> = HashMap::new();
+    gal_els.iter().for_each(|gal_el| {
+        let mut key: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::alloc(&module, basek, k_auto_key, rows, rank);
+        key.generate_from_sk(
+            &module,
+            *gal_el,
+            &sk,
+            &mut source_xa,
+            &mut source_xe,
+            sigma,
+            scratch.borrow(),
+        );
+        auto_keys.insert(*gal_el, key);
+    });
+
+    let log_batch: usize = 0;
+
+    let mut packer: StreamPacker = StreamPacker::new(&module, log_batch, basek, k_ct, rank);
+
+    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_ct, rank);
+
+    ct.encrypt_sk(
+        &module,
+        &pt,
+        &sk_dft,
+        &mut source_xa,
+        &mut source_xe,
+        sigma,
+        scratch.borrow(),
+    );
+
+    let mut res: Vec<GLWECiphertext<Vec<u8>>> = Vec::new();
+
+    let start = Instant::now();
+    (0..module.n() >> log_batch).for_each(|i| {
+        println!("pt {}", pt.data);
+        ct.encrypt_sk(
+            &module,
+            &pt,
+            &sk_dft,
+            &mut source_xa,
+            &mut source_xe,
+            sigma,
+            scratch.borrow(),
+        );
+
+        pt.rotate_inplace(&module, -(1 << log_batch)); // X^-batch * pt
+
+        if reverse_bits_msb(i, log_n as u32) % 5 == 0 {
+            packer.add(&module, &mut res, Some(&ct), &auto_keys, scratch.borrow());
+        } else {
+            packer.add(
+                &module,
+                &mut res,
+                None::<&GLWECiphertext<Vec<u8>>>,
+                &auto_keys,
+                scratch.borrow(),
+            )
+        }
+    });
+    let duration = start.elapsed();
+    println!("Elapsed time: {} ms", duration.as_millis());
+
+    packer.flush(&module, &mut res, &auto_keys, scratch.borrow());
+    packer.reset();
+
+    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
+
+    println!("{}", res.len());
+
+    res.iter().enumerate().for_each(|(i, res_i)| {
+        let mut data: Vec<i64> = vec![0i64; module.n()];
+        data.iter_mut().enumerate().for_each(|(i, x)| {
+            if i % 5 == 0 {
+                *x = reverse_bits_msb(i, log_n as u32) as i64;
+            }
+        });
+        pt_want.data.encode_vec_i64(0, basek, k_pt, &data, 32);
+
+        res_i.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
+
+        println!("{}", pt.data);
+
+        if i & 1 == 0 {
+            pt.sub_inplace_ab(&module, &pt_want);
+        } else {
+            pt.add_inplace(&module, &pt_want);
+        }
+
+        let noise_have = pt.data.std(0, basek).log2();
+        println!("noise_have: {}", noise_have);
+        assert!(
+            noise_have < -((k_ct - basek) as f64),
+            "noise: {}",
+            noise_have
+        );
+    });
+}
+
+#[inline(always)]
+fn reverse_bits_msb(x: usize, n: u32) -> usize {
+    x.reverse_bits() >> (usize::BITS - n)
+}

core/src/test_fft64/mod.rs

@@ -3,6 +3,6 @@ mod gglwe;
 mod ggsw;
 mod glwe;
 mod glwe_fourier;
+mod glwe_packing;
 mod tensor_key;
-
 mod trace;