mirror of
https://github.com/arnaucube/poulpy.git
synced 2026-02-10 13:16:44 +01:00
compiling CBT but failing tests
This commit is contained in:
Pro7ech
@@ -1,41 +1,38 @@
|
||||
use poulpy_core::layouts::{
|
||||
AutomorphismKey, AutomorphismKeyLayout, GGLWEInfos, GGSWInfos, GLWE, GLWEInfos, GLWESecret, LWEInfos, LWESecret, TensorKey,
|
||||
TensorKeyLayout,
|
||||
prepared::{GLWEAutomorphismKeyPrepared, GLWESecretPrepared, TensorKeyPrepared},
|
||||
use poulpy_core::{
|
||||
GLWEAutomorphismKeyEncryptSk, GLWETensorKeyEncryptSk, GetDistribution, ScratchTakeCore,
|
||||
layouts::{
|
||||
GGLWEInfos, GGSWInfos, GLWEAutomorphismKey, GLWEAutomorphismKeyLayout, GLWEInfos, GLWESecretPreparedFactory,
|
||||
GLWESecretToRef, GLWETensorKey, GLWETensorKeyLayout, LWEInfos, LWESecretToRef, prepared::GLWESecretPrepared,
|
||||
},
|
||||
trace_galois_elements,
|
||||
};
|
||||
use std::collections::HashMap;
|
||||
|
||||
use poulpy_hal::{
|
||||
api::{
|
||||
ScratchAvailable, SvpApplyDftToDft, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolBytesOf, SvpPrepare, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace,
|
||||
VecZnxAutomorphism, VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf, VecZnxFillUniform, VecZnxIdftApplyConsume,
|
||||
VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace,
|
||||
VecZnxSwitchRing, VmpPMatAlloc, VmpPrepare,
|
||||
},
|
||||
layouts::{Backend, Data, DataRef, Module, Scratch},
|
||||
layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
|
||||
source::Source,
|
||||
};
|
||||
|
||||
use crate::tfhe::blind_rotation::{
|
||||
BlindRotationAlgo, BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyEncryptSk, BlindRotationKeyInfos,
|
||||
BlindRotationKeyLayout, BlindRotationKeyPrepared,
|
||||
BlindRotationAlgo, BlindRotationKey, BlindRotationKeyEncryptSk, BlindRotationKeyFactory, BlindRotationKeyInfos,
|
||||
BlindRotationKeyLayout,
|
||||
};
|
||||
|
||||
pub trait CircuitBootstrappingKeyInfos {
|
||||
fn brk_infos(&self) -> BlindRotationKeyLayout;
|
||||
fn atk_infos(&self) -> AutomorphismKeyLayout;
|
||||
fn tsk_infos(&self) -> TensorKeyLayout;
|
||||
fn atk_infos(&self) -> GLWEAutomorphismKeyLayout;
|
||||
fn tsk_infos(&self) -> GLWETensorKeyLayout;
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct CircuitBootstrappingKeyLayout {
|
||||
pub layout_brk: BlindRotationKeyLayout,
|
||||
pub layout_atk: AutomorphismKeyLayout,
|
||||
pub layout_tsk: TensorKeyLayout,
|
||||
pub layout_atk: GLWEAutomorphismKeyLayout,
|
||||
pub layout_tsk: GLWETensorKeyLayout,
|
||||
}
|
||||
|
||||
impl CircuitBootstrappingKeyInfos for CircuitBootstrappingKeyLayout {
|
||||
fn atk_infos(&self) -> AutomorphismKeyLayout {
|
||||
fn atk_infos(&self) -> GLWEAutomorphismKeyLayout {
|
||||
self.layout_atk
|
||||
}
|
||||
|
||||
@@ -43,96 +40,114 @@ impl CircuitBootstrappingKeyInfos for CircuitBootstrappingKeyLayout {
|
||||
self.layout_brk
|
||||
}
|
||||
|
||||
fn tsk_infos(&self) -> TensorKeyLayout {
|
||||
fn tsk_infos(&self) -> GLWETensorKeyLayout {
|
||||
self.layout_tsk
|
||||
}
|
||||
}
|
||||
|
||||
pub trait CircuitBootstrappingKeyEncryptSk<B: Backend> {
|
||||
pub trait CircuitBootstrappingKeyEncryptSk<BRA: BlindRotationAlgo, BE: Backend> {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn encrypt_sk<DLwe, DGlwe, INFOS>(
|
||||
module: &Module<B>,
|
||||
sk_lwe: &LWESecret<DLwe>,
|
||||
sk_glwe: &GLWESecret<DGlwe>,
|
||||
cbt_infos: &INFOS,
|
||||
fn circuit_bootstrapping_key_encrypt_sk<D, S0, S1>(
|
||||
&self,
|
||||
res: &mut CircuitBootstrappingKey<D, BRA>,
|
||||
sk_lwe: &S0,
|
||||
sk_glwe: &S1,
|
||||
source_xa: &mut Source,
|
||||
source_xe: &mut Source,
|
||||
scratch: &mut Scratch<B>,
|
||||
) -> Self
|
||||
scratch: &mut Scratch<BE>,
|
||||
) where
|
||||
D: DataMut,
|
||||
S0: LWESecretToRef + GetDistribution + LWEInfos,
|
||||
S1: GLWESecretToRef + GLWEInfos + GetDistribution;
|
||||
}
|
||||
|
||||
impl<BRA: BlindRotationAlgo> CircuitBootstrappingKey<Vec<u8>, BRA> {
|
||||
pub fn alloc_from_infos<A>(infos: &A) -> Self
|
||||
where
|
||||
INFOS: CircuitBootstrappingKeyInfos,
|
||||
DLwe: DataRef,
|
||||
DGlwe: DataRef;
|
||||
A: CircuitBootstrappingKeyInfos,
|
||||
BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
|
||||
{
|
||||
let atk_infos: &GLWEAutomorphismKeyLayout = &infos.atk_infos();
|
||||
let brk_infos: &BlindRotationKeyLayout = &infos.brk_infos();
|
||||
let trk_infos: &GLWETensorKeyLayout = &infos.tsk_infos();
|
||||
let gal_els: Vec<i64> = trace_galois_elements(atk_infos.log_n(), 2 * atk_infos.n().as_usize() as i64);
|
||||
|
||||
Self {
|
||||
brk: <BlindRotationKey<Vec<u8>, BRA> as BlindRotationKeyFactory<BRA>>::blind_rotation_key_alloc(brk_infos),
|
||||
atk: gal_els
|
||||
.iter()
|
||||
.map(|&gal_el| {
|
||||
let key = GLWEAutomorphismKey::alloc_from_infos(atk_infos);
|
||||
(gal_el, key)
|
||||
})
|
||||
.collect(),
|
||||
tsk: GLWETensorKey::alloc_from_infos(trk_infos),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CircuitBootstrappingKey<D: Data, BRA: BlindRotationAlgo> {
|
||||
pub(crate) brk: BlindRotationKey<D, BRA>,
|
||||
pub(crate) tsk: TensorKey<Vec<u8>>,
|
||||
pub(crate) atk: HashMap<i64, AutomorphismKey<Vec<u8>>>,
|
||||
pub(crate) tsk: GLWETensorKey<Vec<u8>>,
|
||||
pub(crate) atk: HashMap<i64, GLWEAutomorphismKey<Vec<u8>>>,
|
||||
}
|
||||
|
||||
impl<BRA: BlindRotationAlgo, B: Backend> CircuitBootstrappingKeyEncryptSk<B> for CircuitBootstrappingKey<Vec<u8>, BRA>
|
||||
where
|
||||
BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<B>,
|
||||
Module<B>: SvpApplyDftToDft<B>
|
||||
+ VecZnxIdftApplyTmpA<B>
|
||||
+ VecZnxAddScalarInplace
|
||||
+ VecZnxDftBytesOf
|
||||
+ VecZnxBigNormalize<B>
|
||||
+ VecZnxDftApply<B>
|
||||
+ SvpApplyDftToDftInplace<B>
|
||||
+ VecZnxIdftApplyConsume<B>
|
||||
+ VecZnxNormalizeTmpBytes
|
||||
+ VecZnxFillUniform
|
||||
+ VecZnxSubInplace
|
||||
+ VecZnxAddInplace
|
||||
+ VecZnxNormalizeInplace<B>
|
||||
+ VecZnxAddNormal
|
||||
+ VecZnxNormalize<B>
|
||||
+ VecZnxSub
|
||||
+ SvpPrepare<B>
|
||||
+ VecZnxSwitchRing
|
||||
+ SvpPPolBytesOf
|
||||
+ SvpPPolAlloc<B>
|
||||
+ VecZnxAutomorphism,
|
||||
Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeSvpPPol<B> + TakeVecZnxBig<B>,
|
||||
{
|
||||
fn encrypt_sk<DLwe, DGlwe, INFOS>(
|
||||
module: &Module<B>,
|
||||
sk_lwe: &LWESecret<DLwe>,
|
||||
sk_glwe: &GLWESecret<DGlwe>,
|
||||
cbt_infos: &INFOS,
|
||||
impl<D: DataMut, BRA: BlindRotationAlgo> CircuitBootstrappingKey<D, BRA> {
|
||||
pub fn encrypt_sk<M, S0, S1, BE: Backend>(
|
||||
&mut self,
|
||||
module: &M,
|
||||
sk_lwe: &S0,
|
||||
sk_glwe: &S1,
|
||||
source_xa: &mut Source,
|
||||
source_xe: &mut Source,
|
||||
scratch: &mut Scratch<B>,
|
||||
) -> Self
|
||||
where
|
||||
INFOS: CircuitBootstrappingKeyInfos,
|
||||
DLwe: DataRef,
|
||||
DGlwe: DataRef,
|
||||
scratch: &mut Scratch<BE>,
|
||||
) where
|
||||
S0: LWESecretToRef + GetDistribution + LWEInfos,
|
||||
S1: GLWESecretToRef + GLWEInfos + GetDistribution,
|
||||
M: CircuitBootstrappingKeyEncryptSk<BRA, BE>,
|
||||
{
|
||||
assert_eq!(sk_lwe.n(), cbt_infos.brk_infos().n_lwe());
|
||||
assert_eq!(sk_glwe.n(), cbt_infos.brk_infos().n_glwe());
|
||||
assert_eq!(sk_glwe.n(), cbt_infos.atk_infos().n());
|
||||
assert_eq!(sk_glwe.n(), cbt_infos.tsk_infos().n());
|
||||
module.circuit_bootstrapping_key_encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
|
||||
}
|
||||
}
|
||||
|
||||
let atk_infos: AutomorphismKeyLayout = cbt_infos.atk_infos();
|
||||
let brk_infos: BlindRotationKeyLayout = cbt_infos.brk_infos();
|
||||
let trk_infos: TensorKeyLayout = cbt_infos.tsk_infos();
|
||||
impl<BRA: BlindRotationAlgo, BE: Backend> CircuitBootstrappingKeyEncryptSk<BRA, BE> for Module<BE>
|
||||
where
|
||||
Self: GLWETensorKeyEncryptSk<BE>
|
||||
+ BlindRotationKeyEncryptSk<BRA, BE>
|
||||
+ GLWEAutomorphismKeyEncryptSk<BE>
|
||||
+ GLWESecretPreparedFactory<BE>,
|
||||
Scratch<BE>: ScratchTakeCore<BE>,
|
||||
{
|
||||
fn circuit_bootstrapping_key_encrypt_sk<D, S0, S1>(
|
||||
&self,
|
||||
res: &mut CircuitBootstrappingKey<D, BRA>,
|
||||
sk_lwe: &S0,
|
||||
sk_glwe: &S1,
|
||||
source_xa: &mut Source,
|
||||
source_xe: &mut Source,
|
||||
scratch: &mut Scratch<BE>,
|
||||
) where
|
||||
D: DataMut,
|
||||
S0: LWESecretToRef + GetDistribution + LWEInfos,
|
||||
S1: GLWESecretToRef + GLWEInfos + GetDistribution,
|
||||
{
|
||||
let brk_infos: &BlindRotationKeyLayout = &res.brk_infos();
|
||||
let atk_infos: &GLWEAutomorphismKeyLayout = &res.atk_infos();
|
||||
let tsk_infos: &GLWETensorKeyLayout = &res.tsk_infos();
|
||||
|
||||
let mut auto_keys: HashMap<i64, AutomorphismKey<Vec<u8>>> = HashMap::new();
|
||||
let gal_els: Vec<i64> = GLWE::trace_galois_elements(module);
|
||||
gal_els.iter().for_each(|gal_el| {
|
||||
let mut key: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc_from_infos(&atk_infos);
|
||||
key.encrypt_sk(module, *gal_el, sk_glwe, source_xa, source_xe, scratch);
|
||||
auto_keys.insert(*gal_el, key);
|
||||
});
|
||||
assert_eq!(sk_lwe.n(), brk_infos.n_lwe());
|
||||
assert_eq!(sk_glwe.n(), brk_infos.n_glwe());
|
||||
assert_eq!(sk_glwe.n(), atk_infos.n());
|
||||
assert_eq!(sk_glwe.n(), tsk_infos.n());
|
||||
|
||||
let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch);
|
||||
for (p, atk) in res.atk.iter_mut() {
|
||||
atk.encrypt_sk(self, *p, sk_glwe, source_xa, source_xe, scratch);
|
||||
}
|
||||
|
||||
let mut brk: BlindRotationKey<Vec<u8>, BRA> = BlindRotationKey::<Vec<u8>, BRA>::alloc(&brk_infos);
|
||||
brk.encrypt_sk(
|
||||
module,
|
||||
let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(self, brk_infos.rank());
|
||||
|
||||
res.brk.encrypt_sk(
|
||||
self,
|
||||
&sk_glwe_prepared,
|
||||
sk_lwe,
|
||||
source_xa,
|
||||
@@ -140,27 +155,15 @@ where
|
||||
scratch,
|
||||
);
|
||||
|
||||
let mut tsk: TensorKey<Vec<u8>> = TensorKey::alloc_from_infos(&trk_infos);
|
||||
tsk.encrypt_sk(module, sk_glwe, source_xa, source_xe, scratch);
|
||||
|
||||
Self {
|
||||
brk,
|
||||
atk: auto_keys,
|
||||
tsk,
|
||||
}
|
||||
res.tsk
|
||||
.encrypt_sk(self, sk_glwe, source_xa, source_xe, scratch);
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CircuitBootstrappingKeyPrepared<D: Data, BRA: BlindRotationAlgo, B: Backend> {
|
||||
pub(crate) brk: BlindRotationKeyPrepared<D, BRA, B>,
|
||||
pub(crate) tsk: TensorKeyPrepared<Vec<u8>, B>,
|
||||
pub(crate) atk: HashMap<i64, GLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
|
||||
}
|
||||
|
||||
impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> CircuitBootstrappingKeyInfos for CircuitBootstrappingKeyPrepared<D, BRA, B> {
|
||||
fn atk_infos(&self) -> AutomorphismKeyLayout {
|
||||
impl<D: DataRef, BRA: BlindRotationAlgo> CircuitBootstrappingKeyInfos for CircuitBootstrappingKey<D, BRA> {
|
||||
fn atk_infos(&self) -> GLWEAutomorphismKeyLayout {
|
||||
let (_, atk) = self.atk.iter().next().expect("atk is empty");
|
||||
AutomorphismKeyLayout {
|
||||
GLWEAutomorphismKeyLayout {
|
||||
n: atk.n(),
|
||||
base2k: atk.base2k(),
|
||||
k: atk.k(),
|
||||
@@ -181,8 +184,8 @@ impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> CircuitBootstrappingKeyInfo
|
||||
}
|
||||
}
|
||||
|
||||
fn tsk_infos(&self) -> TensorKeyLayout {
|
||||
TensorKeyLayout {
|
||||
fn tsk_infos(&self) -> GLWETensorKeyLayout {
|
||||
GLWETensorKeyLayout {
|
||||
n: self.tsk.n(),
|
||||
base2k: self.tsk.base2k(),
|
||||
k: self.tsk.k(),
|
||||
@@ -192,22 +195,3 @@ impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> CircuitBootstrappingKeyInfo
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> PrepareAlloc<B, CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B>>
|
||||
for CircuitBootstrappingKey<D, BRA>
|
||||
where
|
||||
Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
|
||||
BlindRotationKey<D, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
|
||||
TensorKey<D>: PrepareAlloc<B, TensorKeyPrepared<Vec<u8>, B>>,
|
||||
AutomorphismKey<D>: PrepareAlloc<B, GLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
|
||||
{
|
||||
fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B> {
|
||||
let brk: BlindRotationKeyPrepared<Vec<u8>, BRA, B> = self.brk.prepare_alloc(module, scratch);
|
||||
let tsk: TensorKeyPrepared<Vec<u8>, B> = self.tsk.prepare_alloc(module, scratch);
|
||||
let mut atk: HashMap<i64, GLWEAutomorphismKeyPrepared<Vec<u8>, B>> = HashMap::new();
|
||||
for (key, value) in &self.atk {
|
||||
atk.insert(*key, value.prepare_alloc(module, scratch));
|
||||
}
|
||||
CircuitBootstrappingKeyPrepared { brk, tsk, atk }
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user