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,5 +1,6 @@
use poulpy_core::layouts::{
GGLWEAutomorphismKey, GGLWETensorKey, GLWECiphertext, GLWESecret, LWESecret,
GGLWEAutomorphismKey, GGLWEAutomorphismKeyLayout, GGLWELayoutInfos, GGLWETensorKey, GGLWETensorKeyLayout, GGSWInfos,
GLWECiphertext, GLWEInfos, GLWESecret, LWEInfos, LWESecret,
prepared::{GGLWEAutomorphismKeyPrepared, GGLWETensorKeyPrepared, GLWESecretPrepared, PrepareAlloc},
};
use std::collections::HashMap;
@@ -9,7 +10,7 @@ use poulpy_hal::{
ScratchAvailable, SvpApplyDftToDft, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare, TakeScalarZnx,
TakeSvpPPol, TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace,
VecZnxAutomorphism, VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxFillUniform, VecZnxIdftApplyConsume,
VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace,
VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace,
VecZnxSwitchRing, VmpPMatAlloc, VmpPrepare,
},
layouts::{Backend, Data, DataRef, Module, Scratch},
@@ -17,27 +18,49 @@ use poulpy_hal::{
};
use crate::tfhe::blind_rotation::{
BlindRotationAlgo, BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyEncryptSk, BlindRotationKeyPrepared,
BlindRotationAlgo, BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyEncryptSk, BlindRotationKeyInfos,
BlindRotationKeyLayout, BlindRotationKeyPrepared,
};
pub trait CircuitBootstrappingKeyInfos {
fn layout_brk(&self) -> BlindRotationKeyLayout;
fn layout_atk(&self) -> GGLWEAutomorphismKeyLayout;
fn layout_tsk(&self) -> GGLWETensorKeyLayout;
}
pub struct CircuitBootstrappingKeyLayout {
pub layout_brk: BlindRotationKeyLayout,
pub layout_atk: GGLWEAutomorphismKeyLayout,
pub layout_tsk: GGLWETensorKeyLayout,
}
impl CircuitBootstrappingKeyInfos for CircuitBootstrappingKeyLayout {
fn layout_atk(&self) -> GGLWEAutomorphismKeyLayout {
self.layout_atk
}
fn layout_brk(&self) -> BlindRotationKeyLayout {
self.layout_brk
}
fn layout_tsk(&self) -> GGLWETensorKeyLayout {
self.layout_tsk
}
}
pub trait CircuitBootstrappingKeyEncryptSk<B: Backend> {
#[allow(clippy::too_many_arguments)]
fn encrypt_sk<DLwe, DGlwe>(
fn encrypt_sk<DLwe, DGlwe, INFOS>(
module: &Module<B>,
basek: usize,
sk_lwe: &LWESecret<DLwe>,
sk_glwe: &GLWESecret<DGlwe>,
k_brk: usize,
rows_brk: usize,
k_trace: usize,
rows_trace: usize,
k_tsk: usize,
rows_tsk: usize,
cbt_infos: &INFOS,
source_xa: &mut Source,
source_xe: &mut Source,
scratch: &mut Scratch<B>,
) -> Self
where
INFOS: CircuitBootstrappingKeyInfos,
DLwe: DataRef,
DGlwe: DataRef;
}
@@ -61,7 +84,7 @@ where
+ VecZnxIdftApplyConsume<B>
+ VecZnxNormalizeTmpBytes
+ VecZnxFillUniform
+ VecZnxSubABInplace
+ VecZnxSubInplace
+ VecZnxAddInplace
+ VecZnxNormalizeInplace<B>
+ VecZnxAddNormal
@@ -74,46 +97,41 @@ where
+ VecZnxAutomorphism,
Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeSvpPPol<B> + TakeVecZnxBig<B>,
{
fn encrypt_sk<DLwe, DGlwe>(
fn encrypt_sk<DLwe, DGlwe, INFOS>(
module: &Module<B>,
basek: usize,
sk_lwe: &LWESecret<DLwe>,
sk_glwe: &GLWESecret<DGlwe>,
k_brk: usize,
rows_brk: usize,
k_trace: usize,
rows_trace: usize,
k_tsk: usize,
rows_tsk: usize,
cbt_infos: &INFOS,
source_xa: &mut Source,
source_xe: &mut Source,
scratch: &mut Scratch<B>,
) -> Self
where
INFOS: CircuitBootstrappingKeyInfos,
DLwe: DataRef,
DGlwe: DataRef,
Module<B>:,
{
assert_eq!(sk_lwe.n(), cbt_infos.layout_brk().n_lwe());
assert_eq!(sk_glwe.n(), cbt_infos.layout_brk().n_glwe());
assert_eq!(sk_glwe.n(), cbt_infos.layout_atk().n());
assert_eq!(sk_glwe.n(), cbt_infos.layout_tsk().n());
let atk_infos: GGLWEAutomorphismKeyLayout = cbt_infos.layout_atk();
let brk_infos: BlindRotationKeyLayout = cbt_infos.layout_brk();
let trk_infos: GGLWETensorKeyLayout = cbt_infos.layout_tsk();
let mut auto_keys: HashMap<i64, GGLWEAutomorphismKey<Vec<u8>>> = HashMap::new();
let gal_els: Vec<i64> = GLWECiphertext::trace_galois_elements(module);
gal_els.iter().for_each(|gal_el| {
let mut key: GGLWEAutomorphismKey<Vec<u8>> =
GGLWEAutomorphismKey::alloc(sk_glwe.n(), basek, k_trace, rows_trace, 1, sk_glwe.rank());
let mut key: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(&atk_infos);
key.encrypt_sk(module, *gal_el, sk_glwe, source_xa, source_xe, scratch);
auto_keys.insert(*gal_el, key);
});
let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch);
let mut brk: BlindRotationKey<Vec<u8>, BRA> = BlindRotationKey::<Vec<u8>, BRA>::alloc(
sk_glwe.n(),
sk_lwe.n(),
basek,
k_brk,
rows_brk,
sk_glwe.rank(),
);
let mut brk: BlindRotationKey<Vec<u8>, BRA> = BlindRotationKey::<Vec<u8>, BRA>::alloc(&brk_infos);
brk.encrypt_sk(
module,
&sk_glwe_prepared,
@@ -123,7 +141,7 @@ where
scratch,
);
let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(sk_glwe.n(), basek, k_tsk, rows_tsk, 1, sk_glwe.rank());
let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(&trk_infos);
tsk.encrypt_sk(module, sk_glwe, source_xa, source_xe, scratch);
Self {
@@ -140,6 +158,42 @@ pub struct CircuitBootstrappingKeyPrepared<D: Data, BRA: BlindRotationAlgo, B: B
pub(crate) atk: HashMap<i64, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
}
impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> CircuitBootstrappingKeyInfos for CircuitBootstrappingKeyPrepared<D, BRA, B> {
fn layout_atk(&self) -> GGLWEAutomorphismKeyLayout {
let (_, atk) = self.atk.iter().next().expect("atk is empty");
GGLWEAutomorphismKeyLayout {
n: atk.n(),
base2k: atk.base2k(),
k: atk.k(),
rows: atk.rows(),
digits: atk.digits(),
rank: atk.rank(),
}
}
fn layout_brk(&self) -> BlindRotationKeyLayout {
BlindRotationKeyLayout {
n_glwe: self.brk.n_glwe(),
n_lwe: self.brk.n_lwe(),
base2k: self.brk.base2k(),
k: self.brk.k(),
rows: self.brk.rows(),
rank: self.brk.rank(),
}
}
fn layout_tsk(&self) -> GGLWETensorKeyLayout {
GGLWETensorKeyLayout {
n: self.tsk.n(),
base2k: self.tsk.base2k(),
k: self.tsk.k(),
rows: self.tsk.rows(),
digits: self.tsk.digits(),
rank: self.tsk.rank(),
}
}
}
impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> PrepareAlloc<B, CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B>>
for CircuitBootstrappingKey<D, BRA>
where