Add cross-basek normalization (#90)

* 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

View File

@@ -1,31 +1,46 @@
use poulpy_hal::{
api::{
ScratchAvailable, TakeVecZnxDft, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume, VmpApplyDftToDft, VmpApplyDftToDftAdd,
VmpApplyDftToDftTmpBytes,
ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft,
VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
},
layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxView, ZnxViewMut, ZnxZero},
};
use crate::{
TakeGLWECt,
layouts::{GLWECiphertext, Infos, LWECiphertext, prepared::GLWEToLWESwitchingKeyPrepared},
layouts::{
GGLWELayoutInfos, GLWECiphertext, GLWECiphertextLayout, GLWEInfos, LWECiphertext, LWEInfos, Rank,
prepared::GLWEToLWESwitchingKeyPrepared,
},
};
impl LWECiphertext<Vec<u8>> {
pub fn from_glwe_scratch_space<B: Backend>(
pub fn from_glwe_scratch_space<B: Backend, OUT, IN, KEY>(
module: &Module<B>,
basek: usize,
k_lwe: usize,
k_glwe: usize,
k_ksk: usize,
rank: usize,
lwe_infos: &OUT,
glwe_infos: &IN,
key_infos: &KEY,
) -> usize
where
Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes,
OUT: LWEInfos,
IN: GLWEInfos,
KEY: GGLWELayoutInfos,
Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
{
GLWECiphertext::bytes_of(module.n(), basek, k_lwe, 1)
+ GLWECiphertext::keyswitch_scratch_space(module, basek, k_lwe, k_glwe, k_ksk, 1, rank, 1)
let glwe_layout: GLWECiphertextLayout = GLWECiphertextLayout {
n: module.n().into(),
base2k: lwe_infos.base2k(),
k: lwe_infos.k(),
rank: Rank(1),
};
GLWECiphertext::alloc_bytes_with(
module.n().into(),
lwe_infos.base2k(),
lwe_infos.k(),
1u32.into(),
) + GLWECiphertext::keyswitch_scratch_space(module, &glwe_layout, glwe_infos, key_infos)
}
}
@@ -34,10 +49,11 @@ impl<DLwe: DataMut> LWECiphertext<DLwe> {
#[cfg(debug_assertions)]
{
assert!(self.n() <= a.n());
assert!(self.base2k() == a.base2k());
}
let min_size: usize = self.size().min(a.size());
let n: usize = self.n();
let n: usize = self.n().into();
self.data.zero();
(0..min_size).for_each(|i| {
@@ -64,15 +80,26 @@ impl<DLwe: DataMut> LWECiphertext<DLwe> {
+ VecZnxDftApply<B>
+ VecZnxIdftApplyConsume<B>
+ VecZnxBigAddSmallInplace<B>
+ VecZnxBigNormalize<B>,
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeGLWECt,
+ VecZnxBigNormalize<B>
+ VecZnxNormalize<B>
+ VecZnxNormalizeTmpBytes,
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeGLWECt + TakeVecZnx,
{
#[cfg(debug_assertions)]
{
assert_eq!(self.basek(), a.basek());
assert_eq!(a.n(), ks.n());
assert_eq!(a.n(), module.n() as u32);
assert_eq!(ks.n(), module.n() as u32);
assert!(self.n() <= module.n() as u32);
}
let (mut tmp_glwe, scratch_1) = scratch.take_glwe_ct(a.n(), a.basek(), self.k(), 1);
let glwe_layout: GLWECiphertextLayout = GLWECiphertextLayout {
n: module.n().into(),
base2k: self.base2k(),
k: self.k(),
rank: Rank(1),
};
let (mut tmp_glwe, scratch_1) = scratch.take_glwe_ct(&glwe_layout);
tmp_glwe.keyswitch(module, a, &ks.0, scratch_1);
self.sample_extract(&tmp_glwe);
}

View File

@@ -1,31 +1,46 @@
use poulpy_hal::{
api::{
ScratchAvailable, TakeVecZnxDft, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume, VmpApplyDftToDft, VmpApplyDftToDftAdd,
VmpApplyDftToDftTmpBytes,
ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft,
VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
},
layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxView, ZnxViewMut, ZnxZero},
layouts::{Backend, DataMut, DataRef, Module, Scratch, VecZnx, ZnxView, ZnxViewMut, ZnxZero},
};
use crate::{
TakeGLWECt,
layouts::{GLWECiphertext, Infos, LWECiphertext, prepared::LWEToGLWESwitchingKeyPrepared},
layouts::{
GGLWELayoutInfos, GLWECiphertext, GLWECiphertextLayout, GLWEInfos, LWECiphertext, LWEInfos,
prepared::LWEToGLWESwitchingKeyPrepared,
},
};
impl GLWECiphertext<Vec<u8>> {
pub fn from_lwe_scratch_space<B: Backend>(
pub fn from_lwe_scratch_space<B: Backend, OUT, IN, KEY>(
module: &Module<B>,
basek: usize,
k_lwe: usize,
k_glwe: usize,
k_ksk: usize,
rank: usize,
glwe_infos: &OUT,
lwe_infos: &IN,
key_infos: &KEY,
) -> usize
where
Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes,
OUT: GLWEInfos,
IN: LWEInfos,
KEY: GGLWELayoutInfos,
Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
{
GLWECiphertext::keyswitch_scratch_space(module, basek, k_glwe, k_lwe, k_ksk, 1, 1, rank)
+ GLWECiphertext::bytes_of(module.n(), basek, k_lwe, 1)
let ct: usize = GLWECiphertext::alloc_bytes_with(
module.n().into(),
key_infos.base2k(),
lwe_infos.k().max(glwe_infos.k()),
1u32.into(),
);
let ks: usize = GLWECiphertext::keyswitch_inplace_scratch_space(module, glwe_infos, key_infos);
if lwe_infos.base2k() == key_infos.base2k() {
ct + ks
} else {
let a_conv = VecZnx::alloc_bytes(module.n(), 1, lwe_infos.size()) + module.vec_znx_normalize_tmp_bytes();
ct + a_conv + ks
}
}
}
@@ -47,25 +62,68 @@ impl<D: DataMut> GLWECiphertext<D> {
+ VecZnxDftApply<B>
+ VecZnxIdftApplyConsume<B>
+ VecZnxBigAddSmallInplace<B>
+ VecZnxBigNormalize<B>,
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeGLWECt,
+ VecZnxBigNormalize<B>
+ VecZnxNormalize<B>
+ VecZnxNormalizeTmpBytes,
Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeGLWECt + TakeVecZnx,
{
#[cfg(debug_assertions)]
{
assert!(lwe.n() <= self.n());
assert_eq!(self.basek(), self.basek());
assert_eq!(self.n(), module.n() as u32);
assert_eq!(ksk.n(), module.n() as u32);
assert!(lwe.n() <= module.n() as u32);
}
let (mut glwe, scratch_1) = scratch.take_glwe_ct(ksk.n(), lwe.basek(), lwe.k(), 1);
let (mut glwe, scratch_1) = scratch.take_glwe_ct(&GLWECiphertextLayout {
n: ksk.n(),
base2k: ksk.base2k(),
k: lwe.k(),
rank: 1u32.into(),
});
glwe.data.zero();
let n_lwe: usize = lwe.n();
let n_lwe: usize = lwe.n().into();
(0..lwe.size()).for_each(|i| {
let data_lwe: &[i64] = lwe.data.at(0, i);
glwe.data.at_mut(0, i)[0] = data_lwe[0];
glwe.data.at_mut(1, i)[..n_lwe].copy_from_slice(&data_lwe[1..]);
});
if lwe.base2k() == ksk.base2k() {
for i in 0..lwe.size() {
let data_lwe: &[i64] = lwe.data.at(0, i);
glwe.data.at_mut(0, i)[0] = data_lwe[0];
glwe.data.at_mut(1, i)[..n_lwe].copy_from_slice(&data_lwe[1..]);
}
} else {
let (mut a_conv, scratch_2) = scratch_1.take_vec_znx(module.n(), 1, lwe.size());
a_conv.zero();
for j in 0..lwe.size() {
let data_lwe: &[i64] = lwe.data.at(0, j);
a_conv.at_mut(0, j)[0] = data_lwe[0]
}
module.vec_znx_normalize(
ksk.base2k().into(),
&mut glwe.data,
0,
lwe.base2k().into(),
&a_conv,
0,
scratch_2,
);
a_conv.zero();
for j in 0..lwe.size() {
let data_lwe: &[i64] = lwe.data.at(0, j);
a_conv.at_mut(0, j)[..n_lwe].copy_from_slice(&data_lwe[1..]);
}
module.vec_znx_normalize(
ksk.base2k().into(),
&mut glwe.data,
1,
lwe.base2k().into(),
&a_conv,
0,
scratch_2,
);
}
self.keyswitch(module, &glwe, &ksk.0, scratch_1);
}