Added PrepareScratchSpace trait

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<B: Backend> GGLWEAutomorphismKeyPrepared<Vec<u8>, B> {
     }
 }
 
+impl<DR: DataRef, B: Backend, A: GGLWEInfos> PrepareScratchSpace<B, A> for GGLWEAutomorphismKeyPrepared<DR, B>
+where
+    GGLWESwitchingKeyPrepared<DR, B>: PrepareScratchSpace<B, A>,
+{
+    fn prepare_scratch_space(&self, module: &Module<B>, infos: &A) -> usize {
+        self.key.prepare_scratch_space(module, infos)
+    }
+}
+
 impl<D: DataMut, DR: DataRef, B: Backend> Prepare<B, GGLWEAutomorphismKey<DR>> for GGLWEAutomorphismKeyPrepared<D, B>
 where
     Module<B>: VmpPrepare<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<B: Backend> GGLWECiphertextPrepared<Vec<u8>, B> {
     }
 }
 
+impl<DR: DataRef, B: Backend, A: GGLWEInfos> PrepareScratchSpace<B, A> for GGLWECiphertextPrepared<DR, B>
+where
+    Module<B>: VmpPrepareTmpBytes,
+{
+    fn prepare_scratch_space(&self, module: &Module<B>, infos: &A) -> usize {
+        module.vmp_prepare_tmp_bytes(
+            infos.dnum().into(),
+            infos.rank_in().into(),
+            (infos.rank() + 1).into(),
+            infos.size(),
+        )
+    }
+}
+
 impl<D: DataMut, DR: DataRef, B: Backend> Prepare<B, GGLWECiphertext<DR>> for GGLWECiphertextPrepared<D, B>
 where
     Module<B>: VmpPrepare<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<B: Backend> GGLWESwitchingKeyPrepared<Vec<u8>, B> {
     }
 }
 
+impl<DR: DataRef, B: Backend, A: GGLWEInfos> PrepareScratchSpace<B, A> for GGLWESwitchingKeyPrepared<DR, B>
+where
+    GGLWECiphertextPrepared<DR, B>: PrepareScratchSpace<B, A>,
+{
+    fn prepare_scratch_space(&self, module: &Module<B>, infos: &A) -> usize {
+        self.key.prepare_scratch_space(module, infos)
+    }
+}
+
 impl<D: DataMut, DR: DataRef, B: Backend> Prepare<B, GGLWESwitchingKey<DR>> for GGLWESwitchingKeyPrepared<D, B>
 where
     Module<B>: VmpPrepare<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<D: DataRef, B: Backend> GGLWETensorKeyPrepared<D, B> {
     }
 }
 
+impl<DR: DataRef, B: Backend, A: GGLWEInfos> PrepareScratchSpace<B, A> for GGLWETensorKeyPrepared<DR, B>
+where
+    GGLWESwitchingKeyPrepared<DR, B>: PrepareScratchSpace<B, A>,
+{
+    fn prepare_scratch_space(&self, module: &Module<B>, infos: &A) -> usize {
+        self.keys[0].prepare_scratch_space(module, infos)
+    }
+}
+
 impl<D: DataMut, DR: DataRef, B: Backend> Prepare<B, GGLWETensorKey<DR>> for GGLWETensorKeyPrepared<D, B>
 where
     Module<B>: VmpPrepare<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<D: DataRef, B: Backend> GGSWCiphertextPrepared<D, B> {
     }
 }
 
+impl<DR: DataRef, B: Backend, A: GGSWInfos> PrepareScratchSpace<B, A> for GGSWCiphertextPrepared<DR, B>
+where
+    Module<B>: VmpPrepareTmpBytes,
+{
+    fn prepare_scratch_space(&self, module: &Module<B>, infos: &A) -> usize {
+        module.vmp_prepare_tmp_bytes(
+            infos.dnum().into(),
+            (infos.rank() + 1).into(),
+            (infos.rank() + 1).into(),
+            infos.size(),
+        )
+    }
+}
+
 impl<D: DataMut, DR: DataRef, B: Backend> Prepare<B, GGSWCiphertext<DR>> for GGSWCiphertextPrepared<D, B>
 where
     Module<B>: VmpPrepare<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<DR: DataRef, B: Backend, A: GLWEInfos> PrepareScratchSpace<B, A> for GLWEPublicKeyPrepared<DR, B> {
+    fn prepare_scratch_space(&self, _module: &Module<B>, _infos: &A) -> usize {
+        0
+    }
+}
+
 impl<DM: DataMut, DR: DataRef, B: Backend> Prepare<B, GLWEPublicKey<DR>> for GLWEPublicKeyPrepared<DM, B>
 where
     Module<B>: VecZnxDftApply<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<D: Data, B: Backend> GLWESecretPrepared<D, B> {
     }
 }
 
+impl<DR: DataRef, B: Backend, A: GLWEInfos> PrepareScratchSpace<B, A> for GLWESecretPrepared<DR, B> {
+    fn prepare_scratch_space(&self, _module: &Module<B>, _infos: &A) -> usize {
+        0
+    }
+}
+
 impl<D: DataRef, B: Backend> PrepareAlloc<B, GLWESecretPrepared<Vec<u8>, B>> for GLWESecret<D>
 where
     Module<B>: SvpPrepare<B> + SvpPPolAlloc<B>,
 {
-    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut poulpy_hal::layouts::Scratch<B>) -> GLWESecretPrepared<Vec<u8>, B> {
+    fn prepare_alloc(&self, module: &Module<B>, _scratch: &mut Scratch<B>) -> GLWESecretPrepared<Vec<u8>, B> {
         let mut sk_dft: GLWESecretPrepared<Vec<u8>, B> = GLWESecretPrepared::alloc(module, self);
-        sk_dft.prepare(module, self, scratch);
+        sk_dft.prepare(module, self, _scratch);
         sk_dft
     }
 }
@@ -100,7 +106,7 @@ impl<DM: DataMut, DR: DataRef, B: Backend> Prepare<B, GLWESecret<DR>> for GLWESe
 where
     Module<B>: SvpPrepare<B>,
 {
-    fn prepare(&mut self, module: &Module<B>, other: &GLWESecret<DR>, _scratch: &mut poulpy_hal::layouts::Scratch<B>) {
+    fn prepare(&mut self, module: &Module<B>, other: &GLWESecret<DR>, _scratch: &mut Scratch<B>) {
         (0..self.rank().into()).for_each(|i| {
             module.svp_prepare(&mut self.data, i, &other.data, i);
         });

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<B: Backend> GLWEToLWESwitchingKeyPrepared<Vec<u8>, B> {
     }
 }
 
+impl<DR: DataRef, B: Backend, A: GGLWEInfos> PrepareScratchSpace<B, A> for GLWEToLWESwitchingKeyPrepared<DR, B>
+where
+    GGLWESwitchingKeyPrepared<DR, B>: PrepareScratchSpace<B, A>,
+{
+    fn prepare_scratch_space(&self, module: &Module<B>, infos: &A) -> usize {
+        self.0.prepare_scratch_space(module, infos)
+    }
+}
+
 impl<D: DataRef, B: Backend> PrepareAlloc<B, GLWEToLWESwitchingKeyPrepared<Vec<u8>, B>> for GLWEToLWEKey<D>
 where
     Module<B>: VmpPrepare<B> + VmpPMatAlloc<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<B: Backend> LWESwitchingKeyPrepared<Vec<u8>, B> {
     }
 }
 
+impl<DR: DataRef, B: Backend, A: GGLWEInfos> PrepareScratchSpace<B, A> for LWESwitchingKeyPrepared<DR, B>
+where
+    GGLWESwitchingKeyPrepared<DR, B>: PrepareScratchSpace<B, A>,
+{
+    fn prepare_scratch_space(&self, module: &Module<B>, infos: &A) -> usize {
+        self.0.prepare_scratch_space(module, infos)
+    }
+}
+
 impl<D: DataRef, B: Backend> PrepareAlloc<B, LWESwitchingKeyPrepared<Vec<u8>, B>> for LWESwitchingKey<D>
 where
     Module<B>: VmpPrepare<B> + VmpPMatAlloc<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<B: Backend> LWEToGLWESwitchingKeyPrepared<Vec<u8>, B> {
     }
 }
 
+impl<DR: DataRef, B: Backend, A: GGLWEInfos> PrepareScratchSpace<B, A> for LWEToGLWESwitchingKeyPrepared<DR, B>
+where
+    GGLWESwitchingKeyPrepared<DR, B>: PrepareScratchSpace<B, A>,
+{
+    fn prepare_scratch_space(&self, module: &Module<B>, infos: &A) -> usize {
+        self.0.prepare_scratch_space(module, infos)
+    }
+}
+
 impl<D: DataRef, B: Backend> PrepareAlloc<B, LWEToGLWESwitchingKeyPrepared<Vec<u8>, B>> for LWEToGLWESwitchingKey<D>
 where
     Module<B>: VmpPrepare<B> + VmpPMatAlloc<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<B: Backend, T> {
+    fn prepare_scratch_space(&self, module: &Module<B>, infos: &T) -> usize;
+}
+
 pub trait PrepareAlloc<B: Backend, T> {
     fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> T;
 }