Add cross-basek normalization (#90)

* added cross_basek_normalization * updated method signatures to take layouts * fixed cross-base normalization fix #91 fix #93
2026-02-10 05:06:44 +01:00 · 2025-09-30 14:40:10 +02:00
parent 4da790ea6a
commit 37e13b965c
216 changed files with 12481 additions and 7745 deletions
--- a/poulpy-core/src/automorphism/glwe_ct.rs
+++ b/poulpy-core/src/automorphism/glwe_ct.rs
@@ -1,44 +1,38 @@
 use poulpy_hal::{
    api::{
-        ScratchAvailable, TakeVecZnxDft, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace, VecZnxBigAutomorphismInplace,
-        VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallAInplace, VecZnxBigSubSmallBInplace,
-        VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume, VmpApplyDftToDft, VmpApplyDftToDftAdd,
-        VmpApplyDftToDftTmpBytes,
+        ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace,
+        VecZnxBigAutomorphismInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallInplace,
+        VecZnxBigSubSmallNegateInplace, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume, VecZnxNormalize,
+        VecZnxNormalizeTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
    },
    layouts::{Backend, DataMut, DataRef, Module, Scratch, VecZnxBig},
 };

-use crate::layouts::{GLWECiphertext, Infos, prepared::GGLWEAutomorphismKeyPrepared};
+use crate::layouts::{GGLWELayoutInfos, GLWECiphertext, GLWEInfos, LWEInfos, prepared::GGLWEAutomorphismKeyPrepared};

 impl GLWECiphertext<Vec<u8>> {
-    #[allow(clippy::too_many_arguments)]
-    pub fn automorphism_scratch_space<B: Backend>(
+    pub fn automorphism_scratch_space<B: Backend, OUT, IN, KEY>(
        module: &Module<B>,
-        basek: usize,
-        k_out: usize,
-        k_in: usize,
-        k_ksk: usize,
-        digits: usize,
-        rank: usize,
+        out_infos: &OUT,
+        in_infos: &IN,
+        key_infos: &KEY,
    ) -> usize
    where
-        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes,
+        OUT: GLWEInfos,
+        IN: GLWEInfos,
+        KEY: GGLWELayoutInfos,
+        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
    {
-        Self::keyswitch_scratch_space(module, basek, k_out, k_in, k_ksk, digits, rank, rank)
+        Self::keyswitch_scratch_space(module, out_infos, in_infos, key_infos)
    }

-    pub fn automorphism_inplace_scratch_space<B: Backend>(
-        module: &Module<B>,
-        basek: usize,
-        k_out: usize,
-        k_ksk: usize,
-        digits: usize,
-        rank: usize,
-    ) -> usize
+    pub fn automorphism_inplace_scratch_space<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_infos: &KEY) -> usize
    where
-        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes,
+        OUT: GLWEInfos,
+        KEY: GGLWELayoutInfos,
+        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
    {
-        Self::keyswitch_inplace_scratch_space(module, basek, k_out, k_ksk, digits, rank)
+        Self::keyswitch_inplace_scratch_space(module, out_infos, key_infos)
    }
 }

@@ -59,11 +53,13 @@ impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
            + VecZnxIdftApplyConsume<B>
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
-            + VecZnxAutomorphismInplace<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
+            + VecZnxAutomorphismInplace<B>
+            + VecZnxNormalize<B>
+            + VecZnxNormalizeTmpBytes,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
    {
        self.keyswitch(module, lhs, &rhs.key, scratch);
-        (0..self.rank() + 1).for_each(|i| {
+        (0..(self.rank() + 1).into()).for_each(|i| {
            module.vec_znx_automorphism_inplace(rhs.p(), &mut self.data, i, scratch);
        })
    }
@@ -83,11 +79,13 @@ impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
            + VecZnxIdftApplyConsume<B>
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
-            + VecZnxAutomorphismInplace<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
+            + VecZnxAutomorphismInplace<B>
+            + VecZnxNormalize<B>
+            + VecZnxNormalizeTmpBytes,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
    {
        self.keyswitch_inplace(module, &rhs.key, scratch);
-        (0..self.rank() + 1).for_each(|i| {
+        (0..(self.rank() + 1).into()).for_each(|i| {
            module.vec_znx_automorphism_inplace(rhs.p(), &mut self.data, i, scratch);
        })
    }
@@ -108,19 +106,29 @@ impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
            + VecZnxIdftApplyConsume<B>
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
-            + VecZnxBigAutomorphismInplace<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
+            + VecZnxBigAutomorphismInplace<B>
+            + VecZnxNormalizeTmpBytes
+            + VecZnxNormalize<B>,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
    {
        #[cfg(debug_assertions)]
        {
            self.assert_keyswitch(module, lhs, &rhs.key, scratch);
        }
-        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
+        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n().into(), (self.rank() + 1).into(), rhs.size()); // TODO: optimise size
        let mut res_big: VecZnxBig<_, B> = lhs.keyswitch_internal(module, res_dft, &rhs.key, scratch_1);
-        (0..self.cols()).for_each(|i| {
+        (0..(self.rank() + 1).into()).for_each(|i| {
            module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i, scratch_1);
            module.vec_znx_big_add_small_inplace(&mut res_big, i, &lhs.data, i);
-            module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch_1);
+            module.vec_znx_big_normalize(
+                self.base2k().into(),
+                &mut self.data,
+                i,
+                rhs.base2k().into(),
+                &res_big,
+                i,
+                scratch_1,
+            );
        })
    }

@@ -139,19 +147,29 @@ impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
            + VecZnxIdftApplyConsume<B>
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
-            + VecZnxBigAutomorphismInplace<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
+            + VecZnxBigAutomorphismInplace<B>
+            + VecZnxNormalizeTmpBytes
+            + VecZnxNormalize<B>,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
    {
        #[cfg(debug_assertions)]
        {
            self.assert_keyswitch_inplace(module, &rhs.key, scratch);
        }
-        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
+        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n().into(), (self.rank() + 1).into(), rhs.size()); // TODO: optimise size
        let mut res_big: VecZnxBig<_, B> = self.keyswitch_internal(module, res_dft, &rhs.key, scratch_1);
-        (0..self.cols()).for_each(|i| {
+        (0..(self.rank() + 1).into()).for_each(|i| {
            module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i, scratch_1);
            module.vec_znx_big_add_small_inplace(&mut res_big, i, &self.data, i);
-            module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch_1);
+            module.vec_znx_big_normalize(
+                self.base2k().into(),
+                &mut self.data,
+                i,
+                rhs.base2k().into(),
+                &res_big,
+                i,
+                scratch_1,
+            );
        })
    }

@@ -172,23 +190,33 @@ impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
            + VecZnxBigAutomorphismInplace<B>
-            + VecZnxBigSubSmallAInplace<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
+            + VecZnxBigSubSmallInplace<B>
+            + VecZnxNormalizeTmpBytes
+            + VecZnxNormalize<B>,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
    {
        #[cfg(debug_assertions)]
        {
            self.assert_keyswitch(module, lhs, &rhs.key, scratch);
        }
-        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
+        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n().into(), (self.rank() + 1).into(), rhs.size()); // TODO: optimise size
        let mut res_big: VecZnxBig<_, B> = lhs.keyswitch_internal(module, res_dft, &rhs.key, scratch_1);
-        (0..self.cols()).for_each(|i| {
+        (0..(self.rank() + 1).into()).for_each(|i| {
            module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i, scratch_1);
-            module.vec_znx_big_sub_small_a_inplace(&mut res_big, i, &lhs.data, i);
-            module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch_1);
+            module.vec_znx_big_sub_small_inplace(&mut res_big, i, &lhs.data, i);
+            module.vec_znx_big_normalize(
+                self.base2k().into(),
+                &mut self.data,
+                i,
+                rhs.base2k().into(),
+                &res_big,
+                i,
+                scratch_1,
+            );
        })
    }

-    pub fn automorphism_sub_ab_inplace<DataRhs: DataRef, B: Backend>(
+    pub fn automorphism_sub_inplace<DataRhs: DataRef, B: Backend>(
        &mut self,
        module: &Module<B>,
        rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
@@ -204,23 +232,33 @@ impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
            + VecZnxBigAutomorphismInplace<B>
-            + VecZnxBigSubSmallAInplace<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
+            + VecZnxBigSubSmallInplace<B>
+            + VecZnxNormalizeTmpBytes
+            + VecZnxNormalize<B>,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
    {
        #[cfg(debug_assertions)]
        {
            self.assert_keyswitch_inplace(module, &rhs.key, scratch);
        }
-        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
+        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n().into(), (self.rank() + 1).into(), rhs.size()); // TODO: optimise size
        let mut res_big: VecZnxBig<_, B> = self.keyswitch_internal(module, res_dft, &rhs.key, scratch_1);
-        (0..self.cols()).for_each(|i| {
+        (0..(self.rank() + 1).into()).for_each(|i| {
            module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i, scratch_1);
-            module.vec_znx_big_sub_small_a_inplace(&mut res_big, i, &self.data, i);
-            module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch_1);
+            module.vec_znx_big_sub_small_inplace(&mut res_big, i, &self.data, i);
+            module.vec_znx_big_normalize(
+                self.base2k().into(),
+                &mut self.data,
+                i,
+                rhs.base2k().into(),
+                &res_big,
+                i,
+                scratch_1,
+            );
        })
    }

-    pub fn automorphism_sub_ba<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
+    pub fn automorphism_sub_negate<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
        &mut self,
        module: &Module<B>,
        lhs: &GLWECiphertext<DataLhs>,
@@ -237,23 +275,33 @@ impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
            + VecZnxBigAutomorphismInplace<B>
-            + VecZnxBigSubSmallBInplace<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
+            + VecZnxBigSubSmallNegateInplace<B>
+            + VecZnxNormalizeTmpBytes
+            + VecZnxNormalize<B>,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
    {
        #[cfg(debug_assertions)]
        {
            self.assert_keyswitch(module, lhs, &rhs.key, scratch);
        }
-        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
+        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n().into(), (self.rank() + 1).into(), rhs.size()); // TODO: optimise size
        let mut res_big: VecZnxBig<_, B> = lhs.keyswitch_internal(module, res_dft, &rhs.key, scratch_1);
-        (0..self.cols()).for_each(|i| {
+        (0..(self.rank() + 1).into()).for_each(|i| {
            module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i, scratch_1);
-            module.vec_znx_big_sub_small_b_inplace(&mut res_big, i, &lhs.data, i);
-            module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch_1);
+            module.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, &lhs.data, i);
+            module.vec_znx_big_normalize(
+                self.base2k().into(),
+                &mut self.data,
+                i,
+                rhs.base2k().into(),
+                &res_big,
+                i,
+                scratch_1,
+            );
        })
    }

-    pub fn automorphism_sub_ba_inplace<DataRhs: DataRef, B: Backend>(
+    pub fn automorphism_sub_negate_inplace<DataRhs: DataRef, B: Backend>(
        &mut self,
        module: &Module<B>,
        rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
@@ -269,19 +317,29 @@ impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
            + VecZnxBigAddSmallInplace<B>
            + VecZnxBigNormalize<B>
            + VecZnxBigAutomorphismInplace<B>
-            + VecZnxBigSubSmallBInplace<B>,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
+            + VecZnxBigSubSmallNegateInplace<B>
+            + VecZnxNormalizeTmpBytes
+            + VecZnxNormalize<B>,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
    {
        #[cfg(debug_assertions)]
        {
            self.assert_keyswitch_inplace(module, &rhs.key, scratch);
        }
-        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
+        let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n().into(), (self.rank() + 1).into(), rhs.size()); // TODO: optimise size
        let mut res_big: VecZnxBig<_, B> = self.keyswitch_internal(module, res_dft, &rhs.key, scratch_1);
-        (0..self.cols()).for_each(|i| {
+        (0..(self.rank() + 1).into()).for_each(|i| {
            module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i, scratch_1);
-            module.vec_znx_big_sub_small_b_inplace(&mut res_big, i, &self.data, i);
-            module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch_1);
+            module.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, &self.data, i);
+            module.vec_znx_big_normalize(
+                self.base2k().into(),
+                &mut self.data,
+                i,
+                rhs.base2k().into(),
+                &res_big,
+                i,
+                scratch_1,
+            );
        })
    }
 }