Added support for arbitrary extended LUT

This commit is contained in:
Jean-Philippe Bossuat
2025-06-19 16:33:47 +02:00
parent 6a006b442a
commit 4c1a84d702
17 changed files with 219 additions and 96 deletions

View File

@@ -1,10 +1,15 @@
use std::time::Instant;
use backend::{MatZnxDftOps, MatZnxDftScratch, Module, ScalarZnxDftOps, Scratch, VecZnxDftOps, VecZnxOps, ZnxView, ZnxViewMut, ZnxZero, FFT64};
use backend::{
FFT64, MatZnxDftOps, MatZnxDftScratch, Module, ScalarZnxDftOps, Scratch, VecZnxDftOps, VecZnxOps, ZnxView, ZnxViewMut,
ZnxZero,
};
use itertools::izip;
use crate::{
blind_rotation::key::BlindRotationKeyCGGI, lwe::ciphertext::LWECiphertextToRef, FourierGLWESecret, GGSWCiphertext, GLWECiphertext, GLWECiphertextToMut, GLWECiphertextToRef, GLWEPlaintext, GLWESecret, Infos, LWECiphertext, LWESecret, ScratchCore
GGSWCiphertext, GLWECiphertext, GLWECiphertextToMut, Infos, LWECiphertext, ScratchCore,
blind_rotation::{key::BlindRotationKeyCGGI, lut::LookUpTable},
lwe::ciphertext::LWECiphertextToRef,
};
pub fn cggi_blind_rotate_scratch_space(
@@ -21,26 +26,24 @@ pub fn cggi_blind_rotate_scratch_space(
| GLWECiphertext::external_product_inplace_scratch_space(module, basek, k_lut, k_brk, 1, rank))
}
pub fn cggi_blind_rotate<DataRes, DataIn, DataLUT>(
pub fn cggi_blind_rotate<DataRes, DataIn>(
module: &Module<FFT64>,
res: &mut GLWECiphertext<DataRes>,
lwe: &LWECiphertext<DataIn>,
lut: &GLWEPlaintext<DataLUT>,
lut: &LookUpTable,
brk: &BlindRotationKeyCGGI<FFT64>,
scratch: &mut Scratch,
) where
DataRes: AsRef<[u8]> + AsMut<[u8]>,
DataIn: AsRef<[u8]>,
DataLUT: AsRef<[u8]>,
{
let basek = res.basek();
let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
let mut out_mut: GLWECiphertext<&mut [u8]> = res.to_mut();
let lwe_ref: LWECiphertext<&[u8]> = lwe.to_ref();
let lut_ref: GLWECiphertext<&[u8]> = lut.to_ref();
let cols = out_mut.rank()+1;
let cols: usize = out_mut.rank() + 1;
mod_switch_2n(module, &mut lwe_2n, &lwe_ref);
@@ -50,7 +53,7 @@ pub fn cggi_blind_rotate<DataRes, DataIn, DataLUT>(
out_mut.data.zero();
// Initialize out to X^{b} * LUT(X)
module.vec_znx_rotate(b, &mut out_mut.data, 0, &lut_ref.data, 0);
module.vec_znx_rotate(b, &mut out_mut.data, 0, &lut.data[0], 0);
let block_size: usize = brk.block_size();
@@ -68,37 +71,32 @@ pub fn cggi_blind_rotate<DataRes, DataIn, DataLUT>(
brk.data.chunks_exact(block_size)
)
.for_each(|(ai, ski)| {
out_mut.dft(module, &mut acc_dft);
acc_add_dft.data.zero();
izip!(ai.iter(), ski.iter())
.enumerate()
.for_each(|(i, (aii, skii))| {
izip!(ai.iter(), ski.iter()).for_each(|(aii, skii)| {
// vmp_res = DFT(acc) * BRK[i]
module.vmp_apply(&mut vmp_res.data, &acc_dft.data, &skii.data, scratch5);
// vmp_res = DFT(acc) * BRK[i]
module.vmp_apply(&mut vmp_res.data, &acc_dft.data, &skii.data, scratch5);
// DFT(X^ai -1)
xai_minus_one.zero();
xai_minus_one.at_mut(0, 0)[0] = 1;
module.vec_znx_rotate_inplace(*aii, &mut xai_minus_one, 0);
xai_minus_one.at_mut(0, 0)[0] -= 1;
module.svp_prepare(&mut xai_minus_one_dft, 0, &xai_minus_one, 0);
// DFT(X^ai -1)
xai_minus_one.zero();
xai_minus_one.at_mut(0, 0)[0] = 1;
module.vec_znx_rotate_inplace(*aii, &mut xai_minus_one, 0);
xai_minus_one.at_mut(0, 0)[0] -= 1;
module.svp_prepare(&mut xai_minus_one_dft, 0, &xai_minus_one, 0);
// DFT(X^ai -1) * (DFT(acc) * BRK[i])
(0..cols).for_each(|i|{
module.svp_apply_inplace(&mut vmp_res.data, i, &xai_minus_one_dft, 0);
module.vec_znx_dft_add_inplace(&mut acc_add_dft.data, i, &vmp_res.data, i);
});
// DFT(X^ai -1) * (DFT(acc) * BRK[i])
(0..cols).for_each(|i| {
module.svp_apply_inplace(&mut vmp_res.data, i, &xai_minus_one_dft, 0);
module.vec_znx_dft_add_inplace(&mut acc_add_dft.data, i, &vmp_res.data, i);
});
(0..cols).for_each(|i|{
});
(0..cols).for_each(|i| {
module.vec_znx_dft_add_inplace(&mut acc_dft.data, i, &acc_add_dft.data, i);
});
acc_dft.idft(module, &mut out_mut, scratch5);
acc_dft.idft(module, &mut out_mut, scratch5);
});
let duration: std::time::Duration = start.elapsed();
println!("external products: {} us", duration.as_micros());

View File

@@ -69,14 +69,17 @@ impl BlindRotationKeyCGGI<FFT64> {
}
}
#[allow(dead_code)]
pub(crate) fn rows(&self) -> usize {
self.data[0].rows()
}
#[allow(dead_code)]
pub(crate) fn k(&self) -> usize {
self.data[0].k()
}
#[allow(dead_code)]
pub(crate) fn rank(&self) -> usize {
self.data[0].rank()
}

View File

@@ -0,0 +1,84 @@
use backend::{FFT64, Module, ScratchOwned, VecZnx, VecZnxAlloc, VecZnxOps, ZnxInfos, ZnxViewMut, alloc_aligned};
pub struct LookUpTable {
pub(crate) data: Vec<VecZnx<Vec<u8>>>,
pub(crate) basek: usize,
pub(crate) k: usize,
}
impl LookUpTable {
pub fn alloc(module: &Module<FFT64>, basek: usize, k: usize, extend_factor: usize) -> Self {
let size: usize = k.div_ceil(basek);
let mut data: Vec<VecZnx<Vec<u8>>> = Vec::with_capacity(extend_factor);
(0..extend_factor).for_each(|_| {
data.push(module.new_vec_znx(1, size));
});
Self { data, basek, k }
}
pub fn set(&mut self, module: &Module<FFT64>, f: fn(i64) -> i64, message_modulus: usize) {
let basek: usize = self.basek;
// Get the number minimum limb to store the message modulus
let limbs: usize = message_modulus.div_ceil(1 << basek);
// Scaling factor
let scale: i64 = (1 << (basek * limbs - 1)).div_round(message_modulus) as i64;
// Updates function
let f_scaled = |x: i64| (f(x) % message_modulus as i64) * scale;
// If LUT size > module.n()
let domain_size: usize = self.data[0].n() * self.data.len();
let size: usize = self.k.div_ceil(self.basek);
// Equivalent to AUTO([f(0), f(1), ..., f(n-1)], -1)
let mut lut_full: VecZnx<Vec<u8>> = VecZnx::new::<i64>(domain_size, 1, size);
{
let lut_at: &mut [i64] = lut_full.at_mut(0, limbs - 1);
let start: usize = 0;
let end: usize = (domain_size).div_round(message_modulus);
let y: i64 = f_scaled(0);
(start..end).for_each(|i| {
lut_at[i] = y;
});
(1..message_modulus).for_each(|x| {
let start: usize = (x * domain_size).div_round(message_modulus);
let end: usize = ((x + 1) * domain_size).div_round(message_modulus);
let y: i64 = f_scaled(x as i64);
(start..end).for_each(|i| {
lut_at[domain_size - i] = -y;
})
});
}
// Rotates half the step to the left
let half_step: usize = domain_size.div_round(message_modulus << 1);
module.vec_znx_rotate_inplace(-(half_step as i64), &mut lut_full, 0);
let mut tmp_bytes: Vec<u8> = alloc_aligned(lut_full.n() * size_of::<i64>());
lut_full.normalize(self.basek, 0, &mut tmp_bytes);
if self.data.len() > 1 {
let mut scratch: ScratchOwned = ScratchOwned::new(module.bytes_of_vec_znx(1, size));
module.vec_znx_split(&mut self.data, 0, &lut_full, 0, scratch.borrow());
} else {
module.vec_znx_copy(&mut self.data[0], 0, &lut_full, 0);
}
}
}
pub trait DivRound {
fn div_round(self, rhs: Self) -> Self;
}
impl DivRound for usize {
#[inline]
fn div_round(self, rhs: Self) -> Self {
(self + rhs / 2) / rhs
}
}

View File

@@ -1,6 +1,7 @@
// pub mod cggi;
pub mod ccgi;
pub mod key;
pub mod lut;
#[cfg(test)]
pub mod test_fft64;

View File

@@ -1,10 +1,16 @@
use std::time::Instant;
use backend::{Encoding, Module, ScratchOwned, Stats, ZnxView, ZnxViewMut, FFT64};
use backend::{Encoding, FFT64, Module, ScratchOwned, Stats, VecZnxOps, ZnxView};
use sampling::source::Source;
use crate::{
blind_rotation::{ccgi::{cggi_blind_rotate, cggi_blind_rotate_scratch_space, mod_switch_2n}, key::BlindRotationKeyCGGI}, lwe::{ciphertext::{LWECiphertextToMut, LWECiphertextToRef}, LWEPlaintext}, FourierGLWESecret, GLWECiphertext, GLWEOps, GLWEPlaintext, GLWESecret, Infos, LWECiphertext, LWESecret
FourierGLWESecret, GLWECiphertext, GLWEPlaintext, GLWESecret, Infos, LWECiphertext, LWESecret,
blind_rotation::{
ccgi::{cggi_blind_rotate, cggi_blind_rotate_scratch_space, mod_switch_2n},
key::BlindRotationKeyCGGI,
lut::LookUpTable,
},
lwe::{LWEPlaintext, ciphertext::LWECiphertextToRef},
};
#[test]
@@ -21,6 +27,8 @@ fn blind_rotation() {
let rank: usize = 1;
let block_size: usize = 7;
let message_modulus: usize = 64;
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]);
@@ -32,12 +40,14 @@ fn blind_rotation() {
let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
sk_lwe.fill_binary_block(block_size, &mut source_xs);
let mut scratch: ScratchOwned = ScratchOwned::new(BlindRotationKeyCGGI::generate_from_sk_scratch_space(
&module, basek, k_brk, rank,
) | cggi_blind_rotate_scratch_space(&module, basek, k_lut, k_brk, rows_brk, rank));
let mut scratch: ScratchOwned = ScratchOwned::new(
BlindRotationKeyCGGI::generate_from_sk_scratch_space(&module, basek, k_brk, rank)
| cggi_blind_rotate_scratch_space(&module, basek, k_lut, k_brk, rows_brk, rank),
);
let start: Instant = Instant::now();
let mut brk: BlindRotationKeyCGGI<FFT64> = BlindRotationKeyCGGI::allocate(&module, n_lwe, basek, k_brk, rows_brk, rank);
brk.generate_from_sk(
&module,
&sk_glwe_dft,
@@ -47,6 +57,7 @@ fn blind_rotation() {
3.2,
scratch.borrow(),
);
let duration: std::time::Duration = start.elapsed();
println!("brk-gen: {} ms", duration.as_millis());
@@ -67,44 +78,48 @@ fn blind_rotation() {
println!("{}", pt_lwe.data);
let mut lut: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_lut);
lut.data.at_mut(0, 0)[0] = 0;
(1..module.n()).for_each(|i|{
lut.data.at_mut(0, 0)[i] = - ((module.n() as i64 - i as i64 - 1)<<(basek - module.log_n() - 1));
});
fn lut_fn(x: i64) -> i64 {
2 * x + 1
}
let mut lut: LookUpTable = LookUpTable::alloc(&module, basek, k_lut, 1);
lut.set(&module, lut_fn, message_modulus);
let mut res: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_lut, rank);
let start: Instant = Instant::now();
(0..1).for_each(|i|{
(0..1).for_each(|_| {
cggi_blind_rotate(&module, &mut res, &lwe, &lut, &brk, scratch.borrow());
});
let duration: std::time::Duration = start.elapsed();
println!("blind-rotate: {} ms", duration.as_millis());
let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_lut);
res.decrypt(&module , &mut pt_have, &sk_glwe_dft, scratch.borrow());
res.decrypt(&module, &mut pt_have, &sk_glwe_dft, scratch.borrow());
println!("pt_have: {}", pt_have.data);
let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
mod_switch_2n(&module, &mut lwe_2n, &lwe.to_ref());
let pt_want: i64 = (lwe_2n[0] + lwe_2n[1..].iter().zip(sk_lwe.data.at(0, 0)).map(|(x, y)| x * y).sum::<i64>()) % (module.n() as i64 * 2);
let pt_want: i64 = (lwe_2n[0]
+ lwe_2n[1..]
.iter()
.zip(sk_lwe.data.at(0, 0))
.map(|(x, y)| x * y)
.sum::<i64>())
% (module.n() as i64 * 2);
lut.rotate_inplace(&module, pt_want);
module.vec_znx_rotate_inplace(pt_want, &mut lut.data[0], 0);
lut.sub_inplace_ab(&module, &pt_have);
println!("pt_want: {}", lut.data[0]);
let noise: f64 = lut.data.std(0, basek);
module.vec_znx_sub_ab_inplace(&mut lut.data[0], 0, &pt_have.data, 0);
let noise: f64 = lut.data[0].std(0, basek);
println!("noise: {}", noise);
}