diff --git a/core/src/glwe_packing.rs b/core/src/glwe_packing.rs new file mode 100644 index 0000000..b391aeb --- /dev/null +++ b/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, + log_batch: usize, + counter: usize, +} + +/// [Accumulator] stores intermediate packing result. +/// There are Log(N) such accumulators in a [StreamPacker]. +struct Accumulator { + data: GLWECiphertext>, + 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, 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, log_batch: usize, basek: usize, k: usize, rank: usize) -> Self { + let mut accumulators: Vec = Vec::::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, 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) -> Vec { + 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, DataAK: AsRef<[u8]>>( + &mut self, + module: &Module, + res: &mut Vec>>, + a: Option<&GLWECiphertext>, + auto_keys: &HashMap>, + 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>( + &mut self, + module: &Module, + res: &mut Vec>>, + auto_keys: &HashMap>, + scratch: &mut Scratch, + ) { + if self.counter != 0 { + while self.counter != 0 { + self.add( + module, + res, + None::<&GLWECiphertext>>, + auto_keys, + scratch, + ); + } + } + } +} + +fn pack_core_scratch_space(module: &Module, ct_size: usize, autokey_size: usize, rank: usize) -> usize { + combine_scratch_space(module, ct_size, autokey_size, rank) +} + +fn pack_core, DataAK: AsRef<[u8]>>( + module: &Module, + a: Option<&GLWECiphertext>, + accumulators: &mut [Accumulator], + i: usize, + auto_keys: &HashMap>, + 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>>, + acc_next, + i + 1, + auto_keys, + scratch, + ); + } + } +} + +fn combine_scratch_space(module: &Module, 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, DataAK: AsRef<[u8]>>( + module: &Module, + acc: &mut Accumulator, + b: Option<&GLWECiphertext>, + i: usize, + auto_keys: &HashMap>, + scratch: &mut Scratch, +) { + let log_n: usize = module.log_n(); + let a: &mut GLWECiphertext> = &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; + } + } +} diff --git a/core/src/test_fft64/glwe_packing.rs b/core/src/test_fft64/glwe_packing.rs new file mode 100644 index 0000000..3a0641d --- /dev/null +++ b/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 = Module::::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> = SecretKey::alloc(&module, rank); + sk.fill_ternary_prob(0.5, &mut source_xs); + + let mut sk_dft: SecretKeyFourier, FFT64> = SecretKeyFourier::alloc(&module, rank); + sk_dft.dft(&module, &sk); + + let mut pt: GLWEPlaintext> = GLWEPlaintext::alloc(&module, basek, k_ct); + let mut data: Vec = 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 = StreamPacker::galois_elements(&module); + + let mut auto_keys: HashMap, FFT64>> = HashMap::new(); + gal_els.iter().for_each(|gal_el| { + let mut key: AutomorphismKey, 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> = 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>> = 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>>, + &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> = GLWEPlaintext::alloc(&module, basek, k_ct); + + println!("{}", res.len()); + + res.iter().enumerate().for_each(|(i, res_i)| { + let mut data: Vec = 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) +} diff --git a/core/src/test_fft64/mod.rs b/core/src/test_fft64/mod.rs index 912cdfd..0ccf371 100644 --- a/core/src/test_fft64/mod.rs +++ b/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;