wip packing

poulpy-core/src/glwe_packing.rs

@@ -2,18 +2,13 @@ use std::collections::HashMap;
 
 use poulpy_hal::{
     api::{
-        ScratchAvailable, VecZnxAddInplace, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace, VecZnxBigAutomorphismInplace,
+        ModuleLogN, ModuleN, ScratchAvailable, VecZnxAddInplace, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace, VecZnxBigAutomorphismInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallNegateInplace, VecZnxCopy, VecZnxDftApply, VecZnxDftBytesOf, VecZnxDftCopy, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNegateInplace, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotate, VecZnxRotateInplace, VecZnxRshInplace, VecZnxSub, VecZnxSubInplace, VecZnxSwitchRing, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes
-        VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallNegateInplace, VecZnxCopy, VecZnxDftApply,
-        VecZnxDftBytesOf, VecZnxDftCopy, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNegateInplace, VecZnxNormalize,
-        VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotate, VecZnxRotateInplace, VecZnxRshInplace, VecZnxSub,
-        VecZnxSubInplace, VecZnxSwitchRing, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
     },
-    layouts::{Backend, DataMut, DataRef, Module, Scratch},
+    layouts::{Backend, DataMut, DataRef, GaloisElement, Module, Scratch},
 };
 
 use crate::{
-    GLWEOperations,
+    layouts::{prepared::{AutomorphismKeyPrepared, AutomorphismKeyPreparedToRef}, GGLWEInfos, GLWEAlloc, GLWEInfos, GLWEToRef, LWEInfos, GLWE}, GLWEAutomorphism, GLWEOperations, ScratchTakeCore
-    layouts::{GGLWEInfos, GLWE, GLWEInfos, LWEInfos, prepared::AutomorphismKeyPrepared},
 };
 
 /// [GLWEPacker] enables only the fly GLWE packing
@@ -43,9 +38,10 @@ impl Accumulator {
     /// * `base2k`: base 2 logarithm of the GLWE ciphertext in memory digit representation.
     /// * `k`: base 2 precision of the GLWE ciphertext precision over the Torus.
     /// * `rank`: rank of the GLWE ciphertext.
-    pub fn alloc<A, B: Backend>(module: &Module<B>, infos: &A) -> Self
+    pub fn alloc<A, M>(module: &M, infos: &A) -> Self
     where
         A: GLWEInfos,
+        M: GLWEAlloc
     {
         Self {
             data: GLWE::alloc_from_infos(module, infos),
@@ -66,9 +62,10 @@ impl GLWEPacker {
     ///   and N GLWE ciphertext can be packed. With `log_batch=2` all coefficients
     ///   which are multiples of X^{N/4} are packed. Meaning that N/4 ciphertexts
     ///   can be packed.
-    pub fn new<A, B: Backend>(module: Module<B>, infos: &A, log_batch: usize) -> Self
+    pub fn new<A, M>(module: &M, infos: &A, log_batch: usize) -> Self
     where
         A: GLWEInfos,
+        M: GLWEAlloc
     {
         let mut accumulators: Vec<Accumulator> = Vec::<Accumulator>::new();
         let log_n: usize = infos.n().log2();
@@ -90,13 +87,13 @@ impl GLWEPacker {
     }
 
     /// Number of scratch space bytes required to call [Self::add].
-    pub fn tmp_bytes<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_infos: &KEY) -> usize
+    pub fn tmp_bytes<R, K, M, BE: Backend>(module: &M, res_infos: &R, key_infos: &K) -> usize
     where
-        OUT: GLWEInfos,
+        R: GLWEInfos,
-        KEY: GGLWEInfos,
+        K: GGLWEInfos,
-        Module<B>: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
+        M: GLWEAlloc + GLWEAutomorphism<BE>,
     {
-        pack_core_tmp_bytes(module, out_infos, key_infos)
+        pack_core_tmp_bytes(module, res_infos, key_infos)
     }
 
     pub fn galois_elements<B: Backend>(module: &Module<B>) -> Vec<i64> {
@@ -112,37 +109,17 @@ impl GLWEPacker {
     /// * `a`: ciphertext to pack. Can optionally give None to pack a 0 ciphertext.
     /// * `auto_keys`: a [HashMap] containing the [AutomorphismKeyExec]s.
     /// * `scratch`: scratch space of size at least [Self::tmp_bytes].
-    pub fn add<DataA: DataRef, DataAK: DataRef, B: Backend>(
+    pub fn add<A, K, M, BE: Backend>(
         &mut self,
-        module: &Module<B>,
+        module: &M,
-        a: Option<&GLWE<DataA>>,
+        a: Option<&A>,
-        auto_keys: &HashMap<i64, AutomorphismKeyPrepared<DataAK, B>>,
+        auto_keys: &HashMap<i64, K>,
         scratch: &mut Scratch<B>,
     ) where
-        Module<B>: VecZnxDftBytesOf
+        A: GLWEToRef,
-            + VmpApplyDftToDftTmpBytes
+        K: AutomorphismKeyPreparedToRef<BE>,
-            + VecZnxBigNormalizeTmpBytes
+        M: GLWEAutomorphism<BE>,
-            + VmpApplyDftToDft<B>
+        Scratch<BE>: ScratchTakeCore<BE>,
-            + VmpApplyDftToDftAdd<B>
-            + VecZnxDftApply<B>
-            + VecZnxIdftApplyConsume<B>
-            + VecZnxBigAddSmallInplace<B>
-            + VecZnxBigNormalize<B>
-            + VecZnxCopy
-            + VecZnxRotateInplace<B>
-            + VecZnxSub
-            + VecZnxNegateInplace
-            + VecZnxRshInplace<B>
-            + VecZnxAddInplace
-            + VecZnxNormalizeInplace<B>
-            + VecZnxSubInplace
-            + VecZnxRotate
-            + VecZnxAutomorphismInplace<B>
-            + VecZnxBigSubSmallNegateInplace<B>
-            + VecZnxBigAutomorphismInplace<B>
-            + VecZnxNormalize<B>
-            + VecZnxNormalizeTmpBytes,
-        Scratch<B>: ScratchAvailable,
     {
         assert!(
             (self.counter as u32) < self.accumulators[0].data.n(),
@@ -177,47 +154,27 @@ impl GLWEPacker {
     }
 }
 
-fn pack_core_tmp_bytes<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_infos: &KEY) -> usize
+fn pack_core_tmp_bytes<R, K, M, BE: Backend>(module: &M, res_infos: &R, key_infos: &K) -> usize
 where
-    OUT: GLWEInfos,
+    R: GLWEInfos,
-    KEY: GGLWEInfos,
+    K: GGLWEInfos,
-    Module<B>: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
+    M: GLWEAlloc + GLWEAutomorphism<BE>,
 {
-    combine_tmp_bytes(module, out_infos, key_infos)
+    combine_tmp_bytes(module, res_infos, key_infos)
 }
 
-fn pack_core<D: DataRef, DataAK: DataRef, B: Backend>(
+fn pack_core<A, K, M, BE: Backend>(
-    module: &Module<B>,
+    module: &M,
-    a: Option<&GLWE<D>>,
+    a: Option<&A>,
     accumulators: &mut [Accumulator],
     i: usize,
-    auto_keys: &HashMap<i64, AutomorphismKeyPrepared<DataAK, B>>,
+    auto_keys: &HashMap<i64, K>,
-    scratch: &mut Scratch<B>,
+    scratch: &mut Scratch<BE>,
 ) where
-    Module<B>: VecZnxDftBytesOf
+    A: GLWEToRef + GLWEInfos,
-        + VmpApplyDftToDftTmpBytes
+    K: AutomorphismKeyPreparedToRef<BE>,
-        + VecZnxBigNormalizeTmpBytes
+    M: GLWEAutomorphism<BE> + ModuleLogN + VecZnxCopy,
-        + VmpApplyDftToDft<B>
+    Scratch<BE>: ScratchTakeCore<BE>,
-        + VmpApplyDftToDftAdd<B>
-        + VecZnxDftApply<B>
-        + VecZnxIdftApplyConsume<B>
-        + VecZnxBigAddSmallInplace<B>
-        + VecZnxBigNormalize<B>
-        + VecZnxCopy
-        + VecZnxRotateInplace<B>
-        + VecZnxSub
-        + VecZnxNegateInplace
-        + VecZnxRshInplace<B>
-        + VecZnxAddInplace
-        + VecZnxNormalizeInplace<B>
-        + VecZnxSubInplace
-        + VecZnxRotate
-        + VecZnxAutomorphismInplace<B>
-        + VecZnxBigSubSmallNegateInplace<B>
-        + VecZnxBigAutomorphismInplace<B>
-        + VecZnxNormalize<B>
-        + VecZnxNormalizeTmpBytes,
-    Scratch<B>: ScratchAvailable,
 {
     let log_n: usize = module.log_n();
 
@@ -268,49 +225,29 @@ fn pack_core<D: DataRef, DataAK: DataRef, B: Backend>(
     }
 }
 
-fn combine_tmp_bytes<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_infos: &KEY) -> usize
+fn combine_tmp_bytes<R, K, M, BE: Backend>(module: &M, res_infos: &R, key_infos: &K) -> usize
 where
-    OUT: GLWEInfos,
+    R: GLWEInfos,
-    KEY: GGLWEInfos,
+    K: GGLWEInfos,
-    Module<B>: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
+    M: GLWEAlloc + GLWEAutomorphism<BE>,
 {
-    GLWE::bytes_of_from_infos(module, out_infos)
+    GLWE::bytes_of_from_infos(module, res_infos)
-        + (GLWE::rsh_tmp_bytes(module.n()) | GLWE::automorphism_inplace_tmp_bytes(module, out_infos, key_infos))
+        + (GLWE::rsh_tmp_bytes(module.n()) | module.glwe_automorphism_tmp_bytes(res_infos, res_infos, key_infos))
 }
 
 /// [combine] merges two ciphertexts together.
-fn combine<D: DataRef, DataAK: DataRef, B: Backend>(
+fn combine<B, M, K, BE: Backend>(
-    module: &Module<B>,
+    module: &M,
     acc: &mut Accumulator,
-    b: Option<&GLWE<D>>,
+    b: Option<&B>,
     i: usize,
-    auto_keys: &HashMap<i64, AutomorphismKeyPrepared<DataAK, B>>,
+    auto_keys: &HashMap<i64, K>,
-    scratch: &mut Scratch<B>,
+    scratch: &mut Scratch<BE>,
 ) where
-    Module<B>: VecZnxDftBytesOf
+    B: GLWEToRef + GLWEInfos,
-        + VmpApplyDftToDftTmpBytes
+    K: AutomorphismKeyPreparedToRef<BE>,
-        + VecZnxBigNormalizeTmpBytes
+    M: GLWEAutomorphism<BE> + GaloisElement + VecZnxRotateInplace<BE>,
-        + VmpApplyDftToDft<B>
+    Scratch<BE>: ScratchTakeCore<BE>,
-        + VmpApplyDftToDftAdd<B>
-        + VecZnxDftApply<B>
-        + VecZnxIdftApplyConsume<B>
-        + VecZnxBigAddSmallInplace<B>
-        + VecZnxBigNormalize<B>
-        + VecZnxCopy
-        + VecZnxRotateInplace<B>
-        + VecZnxSub
-        + VecZnxNegateInplace
-        + VecZnxRshInplace<B>
-        + VecZnxAddInplace
-        + VecZnxNormalizeInplace<B>
-        + VecZnxSubInplace
-        + VecZnxRotate
-        + VecZnxAutomorphismInplace<B>
-        + VecZnxBigSubSmallNegateInplace<B>
-        + VecZnxBigAutomorphismInplace<B>
-        + VecZnxNormalize<B>
-        + VecZnxNormalizeTmpBytes,
-    Scratch<B>: ScratchAvailable,
 {
     let log_n: usize = acc.data.n().log2();
     let a: &mut GLWE<Vec<u8>> = &mut acc.data;
@@ -335,7 +272,7 @@ fn combine<D: DataRef, DataAK: DataRef, B: Backend>(
     // since 2*(I(X) * Q/2) = I(X) * Q = 0 mod Q.
     if acc.value {
         if let Some(b) = b {
-            let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(a);
+            let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(module, a);
 
             // a = a * X^-t
             a.rotate_inplace(module, -t, scratch_1);