mirror of
https://github.com/arnaucube/poulpy.git
synced 2026-02-10 13:16:44 +01:00
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:
committed by
GitHub
parent
4da790ea6a
commit
37e13b965c
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user