mirror of
https://github.com/arnaucube/poulpy.git
synced 2026-02-10 13:16:44 +01:00
Add BDD Arithmetic (#98)
* Added some circuit, evaluation + some layouts * Refactor + memory reduction * Rows -> Dnum, Digits -> Dsize * fix #96 + glwe_packing (indirectly CBT) * clippy
This commit is contained in:
Jean-Philippe Bossuat
committed by
GitHub
GitHub
parent
37e13b965c
commit
6357a05509
241
poulpy-schemes/src/tfhe/bdd_arithmetic/key.rs
Normal file
241
poulpy-schemes/src/tfhe/bdd_arithmetic/key.rs
Normal file
@@ -0,0 +1,241 @@
|
||||
#[cfg(test)]
|
||||
use crate::tfhe::bdd_arithmetic::FheUintBlocksPrepDebug;
|
||||
use crate::tfhe::{
|
||||
bdd_arithmetic::{FheUintBlocks, FheUintBlocksPrep, UnsignedInteger},
|
||||
blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyEncryptSk},
|
||||
circuit_bootstrapping::{
|
||||
CircuitBootstrappingKey, CircuitBootstrappingKeyEncryptSk, CircuitBootstrappingKeyLayout,
|
||||
CircuitBootstrappingKeyPrepared, CirtuitBootstrappingExecute,
|
||||
},
|
||||
};
|
||||
use poulpy_core::{
|
||||
TakeGGSW, TakeGLWECt,
|
||||
layouts::{
|
||||
GLWESecret, GLWEToLWEKey, GLWEToLWEKeyLayout, LWECiphertext, LWESecret,
|
||||
prepared::{GLWEToLWESwitchingKeyPrepared, Prepare, PrepareAlloc},
|
||||
},
|
||||
};
|
||||
use poulpy_hal::{
|
||||
api::{
|
||||
ScratchAvailable, SvpApplyDftToDft, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare, TakeScalarZnx,
|
||||
TakeSvpPPol, TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace,
|
||||
VecZnxAutomorphism, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
|
||||
VecZnxDftAllocBytes, VecZnxDftApply, VecZnxFillUniform, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNormalize,
|
||||
VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotate, VecZnxSub, VecZnxSubInplace, VecZnxSwitchRing,
|
||||
VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes, VmpPrepare,
|
||||
},
|
||||
layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
|
||||
source::Source,
|
||||
};
|
||||
|
||||
pub trait BDDKeyInfos {
|
||||
fn cbt_infos(&self) -> CircuitBootstrappingKeyLayout;
|
||||
fn ks_infos(&self) -> GLWEToLWEKeyLayout;
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct BDDKeyLayout {
|
||||
pub cbt: CircuitBootstrappingKeyLayout,
|
||||
pub ks: GLWEToLWEKeyLayout,
|
||||
}
|
||||
|
||||
impl BDDKeyInfos for BDDKeyLayout {
|
||||
fn cbt_infos(&self) -> CircuitBootstrappingKeyLayout {
|
||||
self.cbt
|
||||
}
|
||||
|
||||
fn ks_infos(&self) -> GLWEToLWEKeyLayout {
|
||||
self.ks
|
||||
}
|
||||
}
|
||||
|
||||
pub struct BDDKey<CBT, LWE, BRA>
|
||||
where
|
||||
CBT: Data,
|
||||
LWE: Data,
|
||||
BRA: BlindRotationAlgo,
|
||||
{
|
||||
cbt: CircuitBootstrappingKey<CBT, BRA>,
|
||||
ks: GLWEToLWEKey<LWE>,
|
||||
}
|
||||
|
||||
impl<BRA: BlindRotationAlgo> BDDKey<Vec<u8>, Vec<u8>, BRA> {
|
||||
pub fn encrypt_sk<DLwe, DGlwe, A, BE: Backend>(
|
||||
module: &Module<BE>,
|
||||
sk_lwe: &LWESecret<DLwe>,
|
||||
sk_glwe: &GLWESecret<DGlwe>,
|
||||
infos: &A,
|
||||
source_xa: &mut Source,
|
||||
source_xe: &mut Source,
|
||||
scratch: &mut Scratch<BE>,
|
||||
) -> Self
|
||||
where
|
||||
A: BDDKeyInfos,
|
||||
DLwe: DataRef,
|
||||
DGlwe: DataRef,
|
||||
BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<BE>,
|
||||
Module<BE>: SvpApplyDftToDft<BE>
|
||||
+ VecZnxIdftApplyTmpA<BE>
|
||||
+ VecZnxAddScalarInplace
|
||||
+ VecZnxDftAllocBytes
|
||||
+ VecZnxBigNormalize<BE>
|
||||
+ VecZnxDftApply<BE>
|
||||
+ SvpApplyDftToDftInplace<BE>
|
||||
+ VecZnxIdftApplyConsume<BE>
|
||||
+ VecZnxNormalizeTmpBytes
|
||||
+ VecZnxFillUniform
|
||||
+ VecZnxSubInplace
|
||||
+ VecZnxAddInplace
|
||||
+ VecZnxNormalizeInplace<BE>
|
||||
+ VecZnxAddNormal
|
||||
+ VecZnxNormalize<BE>
|
||||
+ VecZnxSub
|
||||
+ SvpPrepare<BE>
|
||||
+ VecZnxSwitchRing
|
||||
+ SvpPPolAllocBytes
|
||||
+ SvpPPolAlloc<BE>
|
||||
+ VecZnxAutomorphism
|
||||
+ VecZnxAutomorphismInplace<BE>,
|
||||
Scratch<BE>: TakeVecZnxDft<BE> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeSvpPPol<BE> + TakeVecZnxBig<BE>,
|
||||
{
|
||||
let mut ks: GLWEToLWEKey<Vec<u8>> = GLWEToLWEKey::alloc(&infos.ks_infos());
|
||||
ks.encrypt_sk(module, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
|
||||
|
||||
Self {
|
||||
cbt: CircuitBootstrappingKey::encrypt_sk(
|
||||
module,
|
||||
sk_lwe,
|
||||
sk_glwe,
|
||||
&infos.cbt_infos(),
|
||||
source_xa,
|
||||
source_xe,
|
||||
scratch,
|
||||
),
|
||||
ks,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct BDDKeyPrepared<CBT, LWE, BRA, BE>
|
||||
where
|
||||
CBT: Data,
|
||||
LWE: Data,
|
||||
BRA: BlindRotationAlgo,
|
||||
BE: Backend,
|
||||
{
|
||||
cbt: CircuitBootstrappingKeyPrepared<CBT, BRA, BE>,
|
||||
ks: GLWEToLWESwitchingKeyPrepared<LWE, BE>,
|
||||
}
|
||||
|
||||
impl<CBT: DataMut, LWE: DataMut, BRA: BlindRotationAlgo, BE: Backend> PrepareAlloc<BE, BDDKeyPrepared<CBT, LWE, BRA, BE>>
|
||||
for BDDKey<CBT, LWE, BRA>
|
||||
where
|
||||
CircuitBootstrappingKey<CBT, BRA>: PrepareAlloc<BE, CircuitBootstrappingKeyPrepared<CBT, BRA, BE>>,
|
||||
GLWEToLWEKey<LWE>: PrepareAlloc<BE, GLWEToLWESwitchingKeyPrepared<LWE, BE>>,
|
||||
{
|
||||
fn prepare_alloc(&self, module: &Module<BE>, scratch: &mut Scratch<BE>) -> BDDKeyPrepared<CBT, LWE, BRA, BE> {
|
||||
BDDKeyPrepared {
|
||||
cbt: self.cbt.prepare_alloc(module, scratch),
|
||||
ks: self.ks.prepare_alloc(module, scratch),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub trait FheUintPrepare<BE: Backend, OUT, IN> {
|
||||
fn prepare(&self, module: &Module<BE>, out: &mut OUT, bits: &IN, scratch: &mut Scratch<BE>);
|
||||
}
|
||||
|
||||
impl<CBT, OUT, IN, LWE, BRA, BE, T> FheUintPrepare<BE, FheUintBlocksPrep<OUT, BE, T>, FheUintBlocks<IN, T>>
|
||||
for BDDKeyPrepared<CBT, LWE, BRA, BE>
|
||||
where
|
||||
T: UnsignedInteger,
|
||||
CBT: DataRef,
|
||||
OUT: DataMut,
|
||||
IN: DataRef,
|
||||
LWE: DataRef,
|
||||
BRA: BlindRotationAlgo,
|
||||
BE: Backend,
|
||||
Module<BE>: VmpPrepare<BE>
|
||||
+ VecZnxRotate
|
||||
+ VecZnxDftAllocBytes
|
||||
+ VmpApplyDftToDftTmpBytes
|
||||
+ VecZnxBigNormalizeTmpBytes
|
||||
+ VmpApplyDftToDft<BE>
|
||||
+ VmpApplyDftToDftAdd<BE>
|
||||
+ VecZnxDftApply<BE>
|
||||
+ VecZnxIdftApplyConsume<BE>
|
||||
+ VecZnxBigAddSmallInplace<BE>
|
||||
+ VecZnxBigNormalize<BE>
|
||||
+ VecZnxNormalize<BE>
|
||||
+ VecZnxNormalizeTmpBytes,
|
||||
Scratch<BE>: ScratchAvailable + TakeVecZnxDft<BE> + TakeGLWECt + TakeVecZnx + TakeGGSW,
|
||||
CircuitBootstrappingKeyPrepared<CBT, BRA, BE>: CirtuitBootstrappingExecute<BE>,
|
||||
{
|
||||
fn prepare(
|
||||
&self,
|
||||
module: &Module<BE>,
|
||||
out: &mut FheUintBlocksPrep<OUT, BE, T>,
|
||||
bits: &FheUintBlocks<IN, T>,
|
||||
scratch: &mut Scratch<BE>,
|
||||
) {
|
||||
#[cfg(debug_assertions)]
|
||||
{
|
||||
assert_eq!(out.blocks.len(), bits.blocks.len());
|
||||
}
|
||||
let mut lwe: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(&bits.blocks[0]); //TODO: add TakeLWE
|
||||
let (mut tmp_ggsw, scratch_1) = scratch.take_ggsw(out);
|
||||
for (dst, src) in out.blocks.iter_mut().zip(bits.blocks.iter()) {
|
||||
lwe.from_glwe(module, src, &self.ks, scratch_1);
|
||||
self.cbt
|
||||
.execute_to_constant(module, &mut tmp_ggsw, &lwe, 1, 1, scratch_1);
|
||||
dst.prepare(module, &tmp_ggsw, scratch_1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
impl<CBT, OUT, IN, LWE, BRA, BE, T> FheUintPrepare<BE, FheUintBlocksPrepDebug<OUT, T>, FheUintBlocks<IN, T>>
|
||||
for BDDKeyPrepared<CBT, LWE, BRA, BE>
|
||||
where
|
||||
T: UnsignedInteger,
|
||||
CBT: DataRef,
|
||||
OUT: DataMut,
|
||||
IN: DataRef,
|
||||
LWE: DataRef,
|
||||
BRA: BlindRotationAlgo,
|
||||
BE: Backend,
|
||||
Module<BE>: VmpPrepare<BE>
|
||||
+ VecZnxRotate
|
||||
+ VecZnxDftAllocBytes
|
||||
+ VmpApplyDftToDftTmpBytes
|
||||
+ VecZnxBigNormalizeTmpBytes
|
||||
+ VmpApplyDftToDft<BE>
|
||||
+ VmpApplyDftToDftAdd<BE>
|
||||
+ VecZnxDftApply<BE>
|
||||
+ VecZnxIdftApplyConsume<BE>
|
||||
+ VecZnxBigAddSmallInplace<BE>
|
||||
+ VecZnxBigNormalize<BE>
|
||||
+ VecZnxNormalize<BE>
|
||||
+ VecZnxNormalizeTmpBytes,
|
||||
Scratch<BE>: ScratchAvailable + TakeVecZnxDft<BE> + TakeGLWECt + TakeVecZnx + TakeGGSW,
|
||||
CircuitBootstrappingKeyPrepared<CBT, BRA, BE>: CirtuitBootstrappingExecute<BE>,
|
||||
{
|
||||
fn prepare(
|
||||
&self,
|
||||
module: &Module<BE>,
|
||||
out: &mut FheUintBlocksPrepDebug<OUT, T>,
|
||||
bits: &FheUintBlocks<IN, T>,
|
||||
scratch: &mut Scratch<BE>,
|
||||
) {
|
||||
#[cfg(debug_assertions)]
|
||||
{
|
||||
assert_eq!(out.blocks.len(), bits.blocks.len());
|
||||
}
|
||||
let mut lwe: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(&bits.blocks[0]); //TODO: add TakeLWE
|
||||
for (dst, src) in out.blocks.iter_mut().zip(bits.blocks.iter()) {
|
||||
lwe.from_glwe(module, src, &self.ks, scratch);
|
||||
self.cbt
|
||||
.execute_to_constant(module, dst, &lwe, 1, 1, scratch);
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user