Add cross-basek normalization (#90)

* added cross_basek_normalization

* updated method signatures to take layouts

* fixed cross-base normalization

fix #91
fix #93

poulpy-core/src/keyswitching/lwe_ct.rs

@@ -1,26 +1,31 @@
 use poulpy_hal::{
     api::{
         ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
-        VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume, VmpApplyDftToDft, VmpApplyDftToDftAdd,
-        VmpApplyDftToDftTmpBytes,
+        VecZnxCopy, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes,
+        VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
     },
     layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxView, ZnxViewMut, ZnxZero},
 };
 
 use crate::{
     TakeGLWECt,
-    layouts::{GLWECiphertext, Infos, LWECiphertext, prepared::LWESwitchingKeyPrepared},
+    layouts::{
+        GGLWELayoutInfos, GLWECiphertext, GLWECiphertextLayout, LWECiphertext, LWEInfos, Rank, TorusPrecision,
+        prepared::LWESwitchingKeyPrepared,
+    },
 };
 
 impl LWECiphertext<Vec<u8>> {
-    pub fn keyswitch_scratch_space<B: Backend>(
+    pub fn keyswitch_scratch_space<B: Backend, OUT, IN, KEY>(
         module: &Module<B>,
-        basek: usize,
-        k_lwe_out: usize,
-        k_lwe_in: usize,
-        k_ksk: usize,
+        out_infos: &OUT,
+        in_infos: &IN,
+        key_infos: &KEY,
     ) -> usize
     where
+        OUT: LWEInfos,
+        IN: LWEInfos,
+        KEY: GGLWELayoutInfos,
         Module<B>: VecZnxDftAllocBytes
             + VmpApplyDftToDftTmpBytes
             + VecZnxBigNormalizeTmpBytes
@@ -30,10 +35,30 @@ impl LWECiphertext<Vec<u8>> {
             + VecZnxDftApply<B>
             + VecZnxIdftApplyConsume<B>
             + VecZnxBigAddSmallInplace<B>
-            + VecZnxBigNormalize<B>,
+            + VecZnxBigNormalize<B>
+            + VecZnxNormalizeTmpBytes,
     {
-        GLWECiphertext::bytes_of(module.n(), basek, k_lwe_out.max(k_lwe_in), 1)
-            + GLWECiphertext::keyswitch_inplace_scratch_space(module, basek, k_lwe_out, k_ksk, 1, 1)
+        let max_k: TorusPrecision = in_infos.k().max(out_infos.k());
+
+        let glwe_in_infos: GLWECiphertextLayout = GLWECiphertextLayout {
+            n: module.n().into(),
+            base2k: in_infos.base2k(),
+            k: max_k,
+            rank: Rank(1),
+        };
+
+        let glwe_out_infos: GLWECiphertextLayout = GLWECiphertextLayout {
+            n: module.n().into(),
+            base2k: out_infos.base2k(),
+            k: max_k,
+            rank: Rank(1),
+        };
+
+        let glwe_in: usize = GLWECiphertext::alloc_bytes(&glwe_in_infos);
+        let glwe_out: usize = GLWECiphertext::alloc_bytes(&glwe_out_infos);
+        let ks: usize = GLWECiphertext::keyswitch_scratch_space(module, &glwe_out_infos, &glwe_in_infos, key_infos);
+
+        glwe_in + glwe_out + ks
     }
 }
 
@@ -55,32 +80,47 @@ impl<DLwe: DataMut> LWECiphertext<DLwe> {
             + VecZnxDftApply<B>
             + VecZnxIdftApplyConsume<B>
             + VecZnxBigAddSmallInplace<B>
-            + VecZnxBigNormalize<B>,
+            + VecZnxBigNormalize<B>
+            + VecZnxNormalize<B>
+            + VecZnxNormalizeTmpBytes
+            + VecZnxCopy,
         Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
     {
         #[cfg(debug_assertions)]
         {
-            assert!(self.n() <= module.n());
-            assert!(a.n() <= module.n());
-            assert_eq!(self.basek(), a.basek());
+            assert!(self.n() <= module.n() as u32);
+            assert!(a.n() <= module.n() as u32);
+            assert!(scratch.available() >= LWECiphertext::keyswitch_scratch_space(module, self, a, ksk));
         }
 
-        let max_k: usize = self.k().max(a.k());
-        let basek: usize = self.basek();
+        let max_k: TorusPrecision = self.k().max(a.k());
 
-        let (mut glwe, scratch_1) = scratch.take_glwe_ct(ksk.n(), basek, max_k, 1);
-        glwe.data.zero();
+        let a_size: usize = a.k().div_ceil(ksk.base2k()) as usize;
 
-        let n_lwe: usize = a.n();
+        let (mut glwe_in, scratch_1) = scratch.take_glwe_ct(&GLWECiphertextLayout {
+            n: ksk.n(),
+            base2k: a.base2k(),
+            k: max_k,
+            rank: Rank(1),
+        });
+        glwe_in.data.zero();
 
-        (0..a.size()).for_each(|i| {
-            let data_lwe: &[i64] = a.data.at(0, i);
-            glwe.data.at_mut(0, i)[0] = data_lwe[0];
-            glwe.data.at_mut(1, i)[..n_lwe].copy_from_slice(&data_lwe[1..]);
+        let (mut glwe_out, scratch_1) = scratch_1.take_glwe_ct(&GLWECiphertextLayout {
+            n: ksk.n(),
+            base2k: self.base2k(),
+            k: max_k,
+            rank: Rank(1),
         });
 
-        glwe.keyswitch_inplace(module, &ksk.0, scratch_1);
+        let n_lwe: usize = a.n().into();
 
-        self.sample_extract(&glwe);
+        for i in 0..a_size {
+            let data_lwe: &[i64] = a.data.at(0, i);
+            glwe_in.data.at_mut(0, i)[0] = data_lwe[0];
+            glwe_in.data.at_mut(1, i)[..n_lwe].copy_from_slice(&data_lwe[1..]);
+        }
+
+        glwe_out.keyswitch(module, &glwe_in, &ksk.0, scratch_1);
+        self.sample_extract(&glwe_out);
     }
 }