poulpy · arnaucube's git mirror

@@ -5,7 +5,7 @@ };
 
 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)]
@@ -107,6 +107,15 @@ where
         Module<B>: VmpPMatAllocBytes,
     {
         GGLWESwitchingKeyPrepared::alloc_bytes_with(module, base2k, k, rank, rank, dnum, dsize)
+    }
+}
+
+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)
     }
 }

M poulpy-core/src/layouts/prepared/gglwe_ct.rs → 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)]
@@ -257,6 +257,20 @@ dsize.0,
         );
 
         module.vmp_pmat_alloc_bytes(dnum.into(), rank_in.into(), (rank_out + 1).into(), size)
+    }
+}
+
+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(),
+        )
     }
 }

M poulpy-core/src/layouts/prepared/gglwe_ksk.rs → poulpy-core/src/layouts/prepared/gglwe_ksk.rs

@@ -5,7 +5,7 @@ };
 
 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)]
@@ -112,6 +112,15 @@ where
         Module<B>: VmpPMatAllocBytes,
     {
         GGLWECiphertextPrepared::alloc_bytes_with(module, base2k, k, rank_in, rank_out, dnum, dsize)
+    }
+}
+
+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)
     }
 }

M poulpy-core/src/layouts/prepared/gglwe_tsk.rs → poulpy-core/src/layouts/prepared/gglwe_tsk.rs

@@ -5,7 +5,7 @@ };
 
 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)]
@@ -148,6 +148,15 @@ std::mem::swap(&mut i, &mut j);
         };
         let rank: usize = self.rank_out().into();
         &self.keys[i * rank + j - (i * (i + 1) / 2)]
+    }
+}
+
+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)
     }
 }

M poulpy-core/src/layouts/prepared/ggsw_ct.rs → 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)]
@@ -255,6 +255,20 @@
 impl<D: DataRef, B: Backend> GGSWCiphertextPrepared<D, B> {
     pub fn data(&self) -> &VmpPMat<D, B> {
         &self.data
+    }
+}
+
+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(),
+        )
     }
 }

M poulpy-core/src/layouts/prepared/glwe_pk.rs → 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},
     },
 };
 
@@ -177,6 +177,12 @@ fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> GLWEPublicKeyPrepared<Vec<u8>, B> {
         let mut pk_prepared: GLWEPublicKeyPrepared<Vec<u8>, B> = GLWEPublicKeyPrepared::alloc(module, self);
         pk_prepared.prepare(module, self, scratch);
         pk_prepared
+    }
+}
+
+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
     }
 }

M poulpy-core/src/layouts/prepared/glwe_sk.rs → 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 @@ Rank(self.data.cols() as u32)
     }
 }
 
+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 GLWESecretPrepared<DM, B>
 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);
         });

M poulpy-core/src/layouts/prepared/glwe_to_lwe_ksk.rs → poulpy-core/src/layouts/prepared/glwe_to_lwe_ksk.rs

@@ -5,7 +5,7 @@ };
 
 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)]
@@ -110,6 +110,15 @@ where
         Module<B>: VmpPMatAllocBytes,
     {
         GGLWESwitchingKeyPrepared::alloc_bytes_with(module, base2k, k, rank_in, Rank(1), dnum, Dsize(1))
+    }
+}
+
+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)
     }
 }

M poulpy-core/src/layouts/prepared/lwe_ksk.rs → poulpy-core/src/layouts/prepared/lwe_ksk.rs

@@ -5,7 +5,7 @@ };
 
 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)]
@@ -119,6 +119,15 @@ where
         Module<B>: VmpPMatAllocBytes,
     {
         GGLWESwitchingKeyPrepared::alloc_bytes_with(module, base2k, k, Rank(1), Rank(1), dnum, Dsize(1))
+    }
+}
+
+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)
     }
 }

M poulpy-core/src/layouts/prepared/lwe_to_glwe_ksk.rs → poulpy-core/src/layouts/prepared/lwe_to_glwe_ksk.rs

@@ -5,7 +5,7 @@ };
 
 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].
@@ -111,6 +111,15 @@ where
         Module<B>: VmpPMatAllocBytes,
     {
         GGLWESwitchingKeyPrepared::alloc_bytes_with(module, base2k, k, Rank(1), rank_out, dnum, Dsize(1))
+    }
+}
+
+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)
     }
 }

M poulpy-core/src/layouts/prepared/mod.rs → 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;
 }