arnaucube/poulpy
1
0
Fork 0
mirror of https://github.com/arnaucube/poulpy.git synced 2026-02-10 21:26:41 +01:00
Code Issues Packages Projects Releases Wiki Activity

Add cross-basek normalization (#90)

Browse Source 
* added cross_basek_normalization

* updated method signatures to take layouts

* fixed cross-base normalization

fix #91
fix #93
This commit is contained in:
Jean-Philippe Bossuat
2025-09-30 14:40:10 +02:00
committed by GitHub
parent 4da790ea6a
commit 37e13b965c
216 changed files with 12481 additions and 7745 deletions
Download Patch File Download Diff File Expand all files Collapse all files

219
poulpy-core/src/keyswitching/ggsw_ct.rs
View File

@@ -1,101 +1,115 @@
use poulpy_hal::{
api::{
ScratchAvailable, TakeVecZnxBig, TakeVecZnxDft, VecZnxBigAddSmallInplace, VecZnxBigAllocBytes, VecZnxBigNormalize,
VecZnxBigNormalizeTmpBytes, VecZnxCopy, VecZnxDftAddInplace, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxDftCopy,
VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNormalizeTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd,
VmpApplyDftToDftTmpBytes,
ScratchAvailable, TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, VecZnxBigAddSmallInplace, VecZnxBigAllocBytes,
VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxCopy, VecZnxDftAddInplace, VecZnxDftAllocBytes, VecZnxDftApply,
VecZnxDftCopy, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft,
VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
},
layouts::{Backend, DataMut, DataRef, Module, Scratch, VecZnx, VmpPMat, ZnxInfos},
};
use crate::{
layouts::{
GGLWECiphertext, GGSWCiphertext, GLWECiphertext, Infos,
GGLWECiphertext, GGLWELayoutInfos, GGSWCiphertext, GGSWInfos, GLWECiphertext, GLWEInfos, LWEInfos,
prepared::{GGLWESwitchingKeyPrepared, GGLWETensorKeyPrepared},
},
operations::GLWEOperations,
};
impl GGSWCiphertext<Vec<u8>> {
pub(crate) fn expand_row_scratch_space<B: Backend>(
module: &Module<B>,
basek: usize,
self_k: usize,
k_tsk: usize,
digits: usize,
rank: usize,
) -> usize
pub(crate) fn expand_row_scratch_space<B: Backend, OUT, TSK>(module: &Module<B>, out_infos: &OUT, tsk_infos: &TSK) -> usize
where
OUT: GGSWInfos,
TSK: GGLWELayoutInfos,
Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigAllocBytes + VecZnxNormalizeTmpBytes,
{
let tsk_size: usize = k_tsk.div_ceil(basek);
let self_size_out: usize = self_k.div_ceil(basek);
let self_size_in: usize = self_size_out.div_ceil(digits);
let tsk_size: usize = tsk_infos.k().div_ceil(tsk_infos.base2k()) as usize;
let size_in: usize = out_infos
.k()
.div_ceil(tsk_infos.base2k())
.div_ceil(tsk_infos.digits().into()) as usize;
let tmp_dft_i: usize = module.vec_znx_dft_alloc_bytes(rank + 1, tsk_size);
let tmp_a: usize = module.vec_znx_dft_alloc_bytes(1, self_size_in);
let tmp_dft_i: usize = module.vec_znx_dft_alloc_bytes((tsk_infos.rank_out() + 1).into(), tsk_size);
let tmp_a: usize = module.vec_znx_dft_alloc_bytes(1, size_in);
let vmp: usize = module.vmp_apply_dft_to_dft_tmp_bytes(
self_size_out,
self_size_in,
self_size_in,
rank,
rank,
tsk_size,
size_in,
size_in,
(tsk_infos.rank_in()).into(), // Verify if rank+1
(tsk_infos.rank_out()).into(), // Verify if rank+1
tsk_size,
);
let tmp_idft: usize = module.vec_znx_big_alloc_bytes(1, tsk_size);
let norm: usize = module.vec_znx_normalize_tmp_bytes();
tmp_dft_i + ((tmp_a + vmp) | (tmp_idft + norm))
}
#[allow(clippy::too_many_arguments)]
pub fn keyswitch_scratch_space<B: Backend>(
pub fn keyswitch_scratch_space<B: Backend, OUT, IN, KEY, TSK>(
module: &Module<B>,
basek: usize,
k_out: usize,
k_in: usize,
k_ksk: usize,
digits_ksk: usize,
k_tsk: usize,
digits_tsk: usize,
rank: usize,
out_infos: &OUT,
in_infos: &IN,
apply_infos: &KEY,
tsk_infos: &TSK,
) -> usize
where
OUT: GGSWInfos,
IN: GGSWInfos,
KEY: GGLWELayoutInfos,
TSK: GGLWELayoutInfos,
Module<B>: VecZnxDftAllocBytes
+ VmpApplyDftToDftTmpBytes
+ VecZnxBigAllocBytes
+ VecZnxNormalizeTmpBytes
+ VecZnxBigNormalizeTmpBytes,
{
let out_size: usize = k_out.div_ceil(basek);
let res_znx: usize = VecZnx::alloc_bytes(module.n(), rank + 1, out_size);
let ci_dft: usize = module.vec_znx_dft_alloc_bytes(rank + 1, out_size);
let ks: usize = GLWECiphertext::keyswitch_scratch_space(module, basek, k_out, k_in, k_ksk, digits_ksk, rank, rank);
let expand_rows: usize = GGSWCiphertext::expand_row_scratch_space(module, basek, k_out, k_tsk, digits_tsk, rank);
let res_dft: usize = module.vec_znx_dft_alloc_bytes(rank + 1, out_size);
res_znx + ci_dft + (ks | expand_rows | res_dft)
#[cfg(debug_assertions)]
{
assert_eq!(apply_infos.rank_in(), apply_infos.rank_out());
assert_eq!(tsk_infos.rank_in(), tsk_infos.rank_out());
assert_eq!(apply_infos.rank_in(), tsk_infos.rank_in());
}
let rank: usize = apply_infos.rank_out().into();
let size_out: usize = out_infos.k().div_ceil(out_infos.base2k()) as usize;
let res_znx: usize = VecZnx::alloc_bytes(module.n(), rank + 1, size_out);
let ci_dft: usize = module.vec_znx_dft_alloc_bytes(rank + 1, size_out);
let ks: usize = GLWECiphertext::keyswitch_scratch_space(module, out_infos, in_infos, apply_infos);
let expand_rows: usize = GGSWCiphertext::expand_row_scratch_space(module, out_infos, tsk_infos);
let res_dft: usize = module.vec_znx_dft_alloc_bytes(rank + 1, size_out);
if in_infos.base2k() == tsk_infos.base2k() {
res_znx + ci_dft + (ks | expand_rows | res_dft)
} else {
let a_conv: usize = VecZnx::alloc_bytes(
module.n(),
1,
out_infos.k().div_ceil(tsk_infos.base2k()) as usize,
) + module.vec_znx_normalize_tmp_bytes();
res_znx + ci_dft + (a_conv | ks | expand_rows | res_dft)
}
}
#[allow(clippy::too_many_arguments)]
pub fn keyswitch_inplace_scratch_space<B: Backend>(
pub fn keyswitch_inplace_scratch_space<B: Backend, OUT, KEY, TSK>(
module: &Module<B>,
basek: usize,
k_out: usize,
k_ksk: usize,
digits_ksk: usize,
k_tsk: usize,
digits_tsk: usize,
rank: usize,
out_infos: &OUT,
apply_infos: &KEY,
tsk_infos: &TSK,
) -> usize
where
OUT: GGSWInfos,
KEY: GGLWELayoutInfos,
TSK: GGLWELayoutInfos,
Module<B>: VecZnxDftAllocBytes
+ VmpApplyDftToDftTmpBytes
+ VecZnxBigAllocBytes
+ VecZnxNormalizeTmpBytes
+ VecZnxBigNormalizeTmpBytes,
{
GGSWCiphertext::keyswitch_scratch_space(
module, basek, k_out, k_out, k_ksk, digits_ksk, k_tsk, digits_tsk, rank,
)
GGSWCiphertext::keyswitch_scratch_space(module, out_infos, out_infos, apply_infos, tsk_infos)
}
}
@@ -120,18 +134,21 @@ impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
+ VmpApplyDftToDftAdd<B>
+ VecZnxDftAddInplace<B>
+ VecZnxBigNormalize<B>
+ VecZnxIdftApplyTmpA<B>,
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B>,
+ VecZnxIdftApplyTmpA<B>
+ VecZnxNormalize<B>,
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeVecZnx,
{
#[cfg(debug_assertions)]
{
assert_eq!(self.rank(), a.rank());
use crate::layouts::{GLWEInfos, LWEInfos};
assert_eq!(self.rank(), a.rank_out());
assert_eq!(self.rows(), a.rows());
assert_eq!(self.n(), module.n());
assert_eq!(a.n(), module.n());
assert_eq!(tsk.n(), module.n());
assert_eq!(self.n(), module.n() as u32);
assert_eq!(a.n(), module.n() as u32);
assert_eq!(tsk.n(), module.n() as u32);
}
(0..self.rows()).for_each(|row_i| {
(0..self.rows().into()).for_each(|row_i| {
self.at_mut(row_i, 0).copy(module, &a.at(row_i, 0));
});
self.expand_row(module, tsk, scratch);
@@ -159,10 +176,11 @@ impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
+ VecZnxNormalizeTmpBytes
+ VecZnxDftCopy<B>
+ VecZnxDftAddInplace<B>
+ VecZnxIdftApplyTmpA<B>,
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B>,
+ VecZnxIdftApplyTmpA<B>
+ VecZnxNormalize<B>,
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeVecZnx,
{
(0..lhs.rows()).for_each(|row_i| {
(0..lhs.rows().into()).for_each(|row_i| {
// Key-switch column 0, i.e.
// col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0s0' + a1s1' + a2s2') + M[i], a0, a1, a2)
self.at_mut(row_i, 0)
@@ -192,10 +210,11 @@ impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
+ VecZnxNormalizeTmpBytes
+ VecZnxDftCopy<B>
+ VecZnxDftAddInplace<B>
+ VecZnxIdftApplyTmpA<B>,
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B>,
+ VecZnxIdftApplyTmpA<B>
+ VecZnxNormalize<B>,
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeVecZnx,
{
(0..self.rows()).for_each(|row_i| {
(0..self.rows().into()).for_each(|row_i| {
// Key-switch column 0, i.e.
// col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0s0' + a1s1' + a2s2') + M[i], a0, a1, a2)
self.at_mut(row_i, 0)
@@ -220,34 +239,41 @@ impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
+ VmpApplyDftToDftAdd<B>
+ VecZnxDftAddInplace<B>
+ VecZnxBigNormalize<B>
+ VecZnxIdftApplyTmpA<B>,
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B>,
+ VecZnxIdftApplyTmpA<B>
+ VecZnxNormalize<B>,
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeVecZnx,
{
assert!(
scratch.available()
>= GGSWCiphertext::expand_row_scratch_space(
module,
self.basek(),
self.k(),
tsk.k(),
tsk.digits(),
tsk.rank()
)
);
let basek_in: usize = self.base2k().into();
let basek_tsk: usize = tsk.base2k().into();
let n: usize = self.n();
let rank: usize = self.rank();
assert!(scratch.available() >= GGSWCiphertext::expand_row_scratch_space(module, self, tsk));
let n: usize = self.n().into();
let rank: usize = self.rank().into();
let cols: usize = rank + 1;
// Keyswitch the j-th row of the col 0
(0..self.rows()).for_each(|row_i| {
// Pre-compute DFT of (a0, a1, a2)
let (mut ci_dft, scratch_1) = scratch.take_vec_znx_dft(n, cols, self.size());
(0..cols).for_each(|i| {
module.vec_znx_dft_apply(1, 0, &mut ci_dft, i, &self.at(row_i, 0).data, i);
});
let a_size: usize = (self.size() * basek_in).div_ceil(basek_tsk);
(1..cols).for_each(|col_j| {
// Keyswitch the j-th row of the col 0
for row_i in 0..self.rows().into() {
let a = &self.at(row_i, 0).data;
// Pre-compute DFT of (a0, a1, a2)
let (mut ci_dft, scratch_1) = scratch.take_vec_znx_dft(n, cols, a_size);
if basek_in == basek_tsk {
for i in 0..cols {
module.vec_znx_dft_apply(1, 0, &mut ci_dft, i, a, i);
}
} else {
let (mut a_conv, scratch_2) = scratch_1.take_vec_znx(n, 1, a_size);
for i in 0..cols {
module.vec_znx_normalize(basek_tsk, &mut a_conv, 0, basek_in, a, i, scratch_2);
module.vec_znx_dft_apply(1, 0, &mut ci_dft, i, &a_conv, 0);
}
}
for col_j in 1..cols {
// Example for rank 3:
//
// Note: M is a vector (m, Bm, B^2m, B^3m, ...), so each column is
@@ -268,7 +294,7 @@ impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
// col 2: (-(c0s0 + c1s1 + c2s2) , c0 , c1 + M[i], c2 )
// col 3: (-(d0s0 + d1s1 + d2s2) , d0 , d1 , d2 + M[i])
let digits: usize = tsk.digits();
let digits: usize = tsk.digits().into();
let (mut tmp_dft_i, scratch_2) = scratch_1.take_vec_znx_dft(n, cols, tsk.size());
let (mut tmp_a, scratch_3) = scratch_2.take_vec_znx_dft(n, 1, ci_dft.size().div_ceil(digits));
@@ -285,11 +311,11 @@ impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
// a2 * (-(h0s0 + h1s1 + h1s2) + s0s2, h0, h1, h2) = (-(a2h0s0 + a2h1s1 + a2h1s2) + a2s0s2, a2h0, a2h1, a2h2)
// =
// (-(x0s0 + x1s1 + x2s2) + s0(a0s0 + a1s1 + a2s2), x0, x1, x2)
(1..cols).for_each(|col_i| {
for col_i in 1..cols {
let pmat: &VmpPMat<DataTsk, B> = &tsk.at(col_i - 1, col_j - 1).key.data; // Selects Enc(s[i]s[j])
// Extracts a[i] and multipies with Enc(s[i]s[j])
(0..digits).for_each(|di| {
for di in 0..digits {
tmp_a.set_size((ci_dft.size() + di) / digits);
// Small optimization for digits > 2
@@ -307,8 +333,8 @@ impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
} else {
module.vmp_apply_dft_to_dft_add(&mut tmp_dft_i, &tmp_a, pmat, di, scratch_3);
}
});
});
}
}
}
// Adds -(sum a[i] * s[i]) + m) on the i-th column of tmp_idft_i
@@ -322,18 +348,19 @@ impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
// (-(x0s0 + x1s1 + x2s2), x0 + M[i], x1, x2)
module.vec_znx_dft_add_inplace(&mut tmp_dft_i, col_j, &ci_dft, 0);
let (mut tmp_idft, scratch_3) = scratch_2.take_vec_znx_big(n, 1, tsk.size());
(0..cols).for_each(|i| {
for i in 0..cols {
module.vec_znx_idft_apply_tmpa(&mut tmp_idft, 0, &mut tmp_dft_i, i);
module.vec_znx_big_normalize(
self.basek(),
basek_in,
&mut self.at_mut(row_i, col_j).data,
i,
basek_tsk,
&tmp_idft,
0,
scratch_3,
);
});
})
})
}
}
}
}
}