mirror of
https://github.com/arnaucube/poulpy.git
synced 2026-02-10 05:06:44 +01:00
Added support for packing GLWE
This commit is contained in:
Jean-Philippe Bossuat
317
core/src/glwe_packing.rs
Normal file
317
core/src/glwe_packing.rs
Normal file
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
158
core/src/test_fft64/glwe_packing.rs
Normal file
158
core/src/test_fft64/glwe_packing.rs
Normal file
@@ -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)
|
||||
}
|
||||
@@ -3,6 +3,6 @@ mod gglwe;
|
||||
mod ggsw;
|
||||
mod glwe;
|
||||
mod glwe_fourier;
|
||||
mod glwe_packing;
|
||||
mod tensor_key;
|
||||
|
||||
mod trace;
|
||||
|
||||
Reference in New Issue
Block a user