diff --git a/poulpy-core/src/layouts/prepared/gglwe_atk.rs b/poulpy-core/src/layouts/prepared/gglwe_atk.rs index e350aaf..fadb9f3 100644 --- a/poulpy-core/src/layouts/prepared/gglwe_atk.rs +++ b/poulpy-core/src/layouts/prepared/gglwe_atk.rs @@ -5,7 +5,7 @@ use poulpy_hal::{ use crate::layouts::{ Base2K, Degree, Dnum, Dsize, GGLWEAutomorphismKey, GGLWEInfos, GLWEInfos, LWEInfos, Rank, TorusPrecision, - prepared::{GGLWESwitchingKeyPrepared, Prepare, PrepareAlloc}, + prepared::{GGLWESwitchingKeyPrepared, Prepare, PrepareAlloc, PrepareScratchSpace}, }; #[derive(PartialEq, Eq)] @@ -110,6 +110,15 @@ impl GGLWEAutomorphismKeyPrepared, B> { } } +impl PrepareScratchSpace for GGLWEAutomorphismKeyPrepared +where + GGLWESwitchingKeyPrepared: PrepareScratchSpace, +{ + fn prepare_scratch_space(&self, module: &Module, infos: &A) -> usize { + self.key.prepare_scratch_space(module, infos) + } +} + impl Prepare> for GGLWEAutomorphismKeyPrepared where Module: VmpPrepare, diff --git a/poulpy-core/src/layouts/prepared/gglwe_ct.rs b/poulpy-core/src/layouts/prepared/gglwe_ct.rs index eeea3e8..18b437f 100644 --- a/poulpy-core/src/layouts/prepared/gglwe_ct.rs +++ b/poulpy-core/src/layouts/prepared/gglwe_ct.rs @@ -1,12 +1,12 @@ use poulpy_hal::{ - api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare}, + api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare, VmpPrepareTmpBytes}, layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, VmpPMat, ZnxInfos}, oep::VmpPMatAllocBytesImpl, }; use crate::layouts::{ Base2K, BuildError, Degree, Dnum, Dsize, GGLWECiphertext, GGLWEInfos, GLWEInfos, LWEInfos, Rank, TorusPrecision, - prepared::{Prepare, PrepareAlloc}, + prepared::{Prepare, PrepareAlloc, PrepareScratchSpace}, }; #[derive(PartialEq, Eq)] @@ -260,6 +260,20 @@ impl GGLWECiphertextPrepared, B> { } } +impl PrepareScratchSpace for GGLWECiphertextPrepared +where + Module: VmpPrepareTmpBytes, +{ + fn prepare_scratch_space(&self, module: &Module, infos: &A) -> usize { + module.vmp_prepare_tmp_bytes( + infos.dnum().into(), + infos.rank_in().into(), + (infos.rank() + 1).into(), + infos.size(), + ) + } +} + impl Prepare> for GGLWECiphertextPrepared where Module: VmpPrepare, diff --git a/poulpy-core/src/layouts/prepared/gglwe_ksk.rs b/poulpy-core/src/layouts/prepared/gglwe_ksk.rs index 5445c88..495666f 100644 --- a/poulpy-core/src/layouts/prepared/gglwe_ksk.rs +++ b/poulpy-core/src/layouts/prepared/gglwe_ksk.rs @@ -5,7 +5,7 @@ use poulpy_hal::{ use crate::layouts::{ Base2K, Degree, Dnum, Dsize, GGLWEInfos, GGLWESwitchingKey, GLWEInfos, LWEInfos, Rank, TorusPrecision, - prepared::{GGLWECiphertextPrepared, Prepare, PrepareAlloc}, + prepared::{GGLWECiphertextPrepared, Prepare, PrepareAlloc, PrepareScratchSpace}, }; #[derive(PartialEq, Eq)] @@ -115,6 +115,15 @@ impl GGLWESwitchingKeyPrepared, B> { } } +impl PrepareScratchSpace for GGLWESwitchingKeyPrepared +where + GGLWECiphertextPrepared: PrepareScratchSpace, +{ + fn prepare_scratch_space(&self, module: &Module, infos: &A) -> usize { + self.key.prepare_scratch_space(module, infos) + } +} + impl Prepare> for GGLWESwitchingKeyPrepared where Module: VmpPrepare, diff --git a/poulpy-core/src/layouts/prepared/gglwe_tsk.rs b/poulpy-core/src/layouts/prepared/gglwe_tsk.rs index 2efaeff..fd2700a 100644 --- a/poulpy-core/src/layouts/prepared/gglwe_tsk.rs +++ b/poulpy-core/src/layouts/prepared/gglwe_tsk.rs @@ -5,7 +5,7 @@ use poulpy_hal::{ use crate::layouts::{ Base2K, Degree, Dnum, Dsize, GGLWEInfos, GGLWETensorKey, GLWEInfos, LWEInfos, Rank, TorusPrecision, - prepared::{GGLWESwitchingKeyPrepared, Prepare, PrepareAlloc}, + prepared::{GGLWESwitchingKeyPrepared, Prepare, PrepareAlloc, PrepareScratchSpace}, }; #[derive(PartialEq, Eq)] @@ -151,6 +151,15 @@ impl GGLWETensorKeyPrepared { } } +impl PrepareScratchSpace for GGLWETensorKeyPrepared +where + GGLWESwitchingKeyPrepared: PrepareScratchSpace, +{ + fn prepare_scratch_space(&self, module: &Module, infos: &A) -> usize { + self.keys[0].prepare_scratch_space(module, infos) + } +} + impl Prepare> for GGLWETensorKeyPrepared where Module: VmpPrepare, diff --git a/poulpy-core/src/layouts/prepared/ggsw_ct.rs b/poulpy-core/src/layouts/prepared/ggsw_ct.rs index 117ae1e..ccdf310 100644 --- a/poulpy-core/src/layouts/prepared/ggsw_ct.rs +++ b/poulpy-core/src/layouts/prepared/ggsw_ct.rs @@ -1,12 +1,12 @@ use poulpy_hal::{ - api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare}, + api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare, VmpPrepareTmpBytes}, layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, VmpPMat, VmpPMatToRef, ZnxInfos}, oep::VmpPMatAllocBytesImpl, }; use crate::layouts::{ Base2K, BuildError, Degree, Dnum, Dsize, GGSWCiphertext, GGSWInfos, GLWEInfos, LWEInfos, Rank, TorusPrecision, - prepared::{Prepare, PrepareAlloc}, + prepared::{Prepare, PrepareAlloc, PrepareScratchSpace}, }; #[derive(PartialEq, Eq)] @@ -258,6 +258,20 @@ impl GGSWCiphertextPrepared { } } +impl PrepareScratchSpace for GGSWCiphertextPrepared +where + Module: VmpPrepareTmpBytes, +{ + fn prepare_scratch_space(&self, module: &Module, infos: &A) -> usize { + module.vmp_prepare_tmp_bytes( + infos.dnum().into(), + (infos.rank() + 1).into(), + (infos.rank() + 1).into(), + infos.size(), + ) + } +} + impl Prepare> for GGSWCiphertextPrepared where Module: VmpPrepare, diff --git a/poulpy-core/src/layouts/prepared/glwe_pk.rs b/poulpy-core/src/layouts/prepared/glwe_pk.rs index 1ccf982..d06fc30 100644 --- a/poulpy-core/src/layouts/prepared/glwe_pk.rs +++ b/poulpy-core/src/layouts/prepared/glwe_pk.rs @@ -8,7 +8,7 @@ use crate::{ dist::Distribution, layouts::{ Base2K, BuildError, Degree, GLWEInfos, GLWEPublicKey, LWEInfos, Rank, TorusPrecision, - prepared::{Prepare, PrepareAlloc}, + prepared::{Prepare, PrepareAlloc, PrepareScratchSpace}, }, }; @@ -180,6 +180,12 @@ where } } +impl PrepareScratchSpace for GLWEPublicKeyPrepared { + fn prepare_scratch_space(&self, _module: &Module, _infos: &A) -> usize { + 0 + } +} + impl Prepare> for GLWEPublicKeyPrepared where Module: VecZnxDftApply, diff --git a/poulpy-core/src/layouts/prepared/glwe_sk.rs b/poulpy-core/src/layouts/prepared/glwe_sk.rs index 604d617..7aa0fd9 100644 --- a/poulpy-core/src/layouts/prepared/glwe_sk.rs +++ b/poulpy-core/src/layouts/prepared/glwe_sk.rs @@ -1,13 +1,13 @@ use poulpy_hal::{ api::{SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare}, - layouts::{Backend, Data, DataMut, DataRef, Module, SvpPPol, ZnxInfos}, + layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, SvpPPol, ZnxInfos}, }; use crate::{ dist::Distribution, layouts::{ Base2K, Degree, GLWEInfos, GLWESecret, LWEInfos, Rank, TorusPrecision, - prepared::{Prepare, PrepareAlloc}, + prepared::{Prepare, PrepareAlloc, PrepareScratchSpace}, }, }; @@ -85,13 +85,19 @@ impl GLWESecretPrepared { } } +impl PrepareScratchSpace for GLWESecretPrepared { + fn prepare_scratch_space(&self, _module: &Module, _infos: &A) -> usize { + 0 + } +} + impl PrepareAlloc, B>> for GLWESecret where Module: SvpPrepare + SvpPPolAlloc, { - fn prepare_alloc(&self, module: &Module, scratch: &mut poulpy_hal::layouts::Scratch) -> GLWESecretPrepared, B> { + fn prepare_alloc(&self, module: &Module, _scratch: &mut Scratch) -> GLWESecretPrepared, B> { let mut sk_dft: GLWESecretPrepared, B> = GLWESecretPrepared::alloc(module, self); - sk_dft.prepare(module, self, scratch); + sk_dft.prepare(module, self, _scratch); sk_dft } } @@ -100,7 +106,7 @@ impl Prepare> for GLWESe where Module: SvpPrepare, { - fn prepare(&mut self, module: &Module, other: &GLWESecret, _scratch: &mut poulpy_hal::layouts::Scratch) { + fn prepare(&mut self, module: &Module, other: &GLWESecret, _scratch: &mut Scratch) { (0..self.rank().into()).for_each(|i| { module.svp_prepare(&mut self.data, i, &other.data, i); }); diff --git a/poulpy-core/src/layouts/prepared/glwe_to_lwe_ksk.rs b/poulpy-core/src/layouts/prepared/glwe_to_lwe_ksk.rs index 1993878..2d5ba3e 100644 --- a/poulpy-core/src/layouts/prepared/glwe_to_lwe_ksk.rs +++ b/poulpy-core/src/layouts/prepared/glwe_to_lwe_ksk.rs @@ -5,7 +5,7 @@ use poulpy_hal::{ use crate::layouts::{ Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, GLWEToLWEKey, LWEInfos, Rank, TorusPrecision, - prepared::{GGLWESwitchingKeyPrepared, Prepare, PrepareAlloc}, + prepared::{GGLWESwitchingKeyPrepared, Prepare, PrepareAlloc, PrepareScratchSpace}, }; #[derive(PartialEq, Eq)] @@ -113,6 +113,15 @@ impl GLWEToLWESwitchingKeyPrepared, B> { } } +impl PrepareScratchSpace for GLWEToLWESwitchingKeyPrepared +where + GGLWESwitchingKeyPrepared: PrepareScratchSpace, +{ + fn prepare_scratch_space(&self, module: &Module, infos: &A) -> usize { + self.0.prepare_scratch_space(module, infos) + } +} + impl PrepareAlloc, B>> for GLWEToLWEKey where Module: VmpPrepare + VmpPMatAlloc, diff --git a/poulpy-core/src/layouts/prepared/lwe_ksk.rs b/poulpy-core/src/layouts/prepared/lwe_ksk.rs index 9110dfb..c938fe4 100644 --- a/poulpy-core/src/layouts/prepared/lwe_ksk.rs +++ b/poulpy-core/src/layouts/prepared/lwe_ksk.rs @@ -5,7 +5,7 @@ use poulpy_hal::{ use crate::layouts::{ Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, LWEInfos, LWESwitchingKey, Rank, TorusPrecision, - prepared::{GGLWESwitchingKeyPrepared, Prepare, PrepareAlloc}, + prepared::{GGLWESwitchingKeyPrepared, Prepare, PrepareAlloc, PrepareScratchSpace}, }; #[derive(PartialEq, Eq)] @@ -122,6 +122,15 @@ impl LWESwitchingKeyPrepared, B> { } } +impl PrepareScratchSpace for LWESwitchingKeyPrepared +where + GGLWESwitchingKeyPrepared: PrepareScratchSpace, +{ + fn prepare_scratch_space(&self, module: &Module, infos: &A) -> usize { + self.0.prepare_scratch_space(module, infos) + } +} + impl PrepareAlloc, B>> for LWESwitchingKey where Module: VmpPrepare + VmpPMatAlloc, diff --git a/poulpy-core/src/layouts/prepared/lwe_to_glwe_ksk.rs b/poulpy-core/src/layouts/prepared/lwe_to_glwe_ksk.rs index 010ba6d..bd8184b 100644 --- a/poulpy-core/src/layouts/prepared/lwe_to_glwe_ksk.rs +++ b/poulpy-core/src/layouts/prepared/lwe_to_glwe_ksk.rs @@ -5,7 +5,7 @@ use poulpy_hal::{ use crate::layouts::{ Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, LWEInfos, LWEToGLWESwitchingKey, Rank, TorusPrecision, - prepared::{GGLWESwitchingKeyPrepared, Prepare, PrepareAlloc}, + prepared::{GGLWESwitchingKeyPrepared, Prepare, PrepareAlloc, PrepareScratchSpace}, }; /// A special [GLWESwitchingKey] required to for the conversion from [LWECiphertext] to [GLWECiphertext]. @@ -114,6 +114,15 @@ impl LWEToGLWESwitchingKeyPrepared, B> { } } +impl PrepareScratchSpace for LWEToGLWESwitchingKeyPrepared +where + GGLWESwitchingKeyPrepared: PrepareScratchSpace, +{ + fn prepare_scratch_space(&self, module: &Module, infos: &A) -> usize { + self.0.prepare_scratch_space(module, infos) + } +} + impl PrepareAlloc, B>> for LWEToGLWESwitchingKey where Module: VmpPrepare + VmpPMatAlloc, diff --git a/poulpy-core/src/layouts/prepared/mod.rs b/poulpy-core/src/layouts/prepared/mod.rs index 9da8f27..fdf3f09 100644 --- a/poulpy-core/src/layouts/prepared/mod.rs +++ b/poulpy-core/src/layouts/prepared/mod.rs @@ -21,6 +21,10 @@ pub use lwe_ksk::*; pub use lwe_to_glwe_ksk::*; use poulpy_hal::layouts::{Backend, Module, Scratch}; +pub trait PrepareScratchSpace { + fn prepare_scratch_space(&self, module: &Module, infos: &T) -> usize; +} + pub trait PrepareAlloc { fn prepare_alloc(&self, module: &Module, scratch: &mut Scratch) -> T; }