added test for automorphism on automorphism key

base2k/src/module.rs

@@ -74,15 +74,15 @@ impl<B: Backend> Module<B> {
     }
 
     // Returns gen^-1
-    pub fn galois_element_inv(&self, generator: i64) -> i64 {
+    pub fn galois_element_inv(&self, gal_el: i64) -> i64 {
-        if generator == 0 {
+        if gal_el == 0 {
             panic!("cannot invert 0")
         }
         ((mod_exp_u64(
-            generator.abs() as u64,
+            gal_el.abs() as u64,
             (self.cyclotomic_order() - 1) as usize,
         ) & (self.cyclotomic_order() - 1)) as i64)
-            * generator.signum()
+            * gal_el.signum()
     }
 }
 

base2k/src/vec_znx_ops.rs

@@ -574,6 +574,7 @@ impl<BACKEND: Backend> VecZnxOps for Module<BACKEND> {
         #[cfg(debug_assertions)]
         {
             assert_eq!(a.n(), self.n());
+            assert!(k & 1 != 0, "invalid galois element: must be odd but is {}", k);
         }
         unsafe {
             vec_znx::vec_znx_automorphism(

core/src/automorphism.rs

@@ -1,6 +1,7 @@
 use base2k::{
     Backend, FFT64, MatZnxDft, MatZnxDftOps, MatZnxDftToMut, MatZnxDftToRef, Module, ScalarZnx, ScalarZnxDftOps, ScalarZnxOps,
-    ScalarZnxToRef, Scratch, VecZnx, VecZnxDft, VecZnxDftOps, VecZnxDftToMut, VecZnxDftToRef, VecZnxOps, ZnxZero,
+    ScalarZnxToRef, Scratch, VecZnx, VecZnxBigAlloc, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut, VecZnxDftToRef,
+    VecZnxOps, ZnxZero,
 };
 use sampling::source::Source;
 
@@ -20,10 +21,10 @@ pub struct AutomorphismKey<Data, B: Backend> {
 }
 
 impl AutomorphismKey<Vec<u8>, FFT64> {
-    pub fn new(module: &Module<FFT64>, basek: usize, p: i64, k: usize, rows: usize, rank: usize) -> Self {
+    pub fn new(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
         AutomorphismKey {
             key: GLWESwitchingKey::new(module, basek, k, rows, rank, rank),
-            p: p,
+            p: 0,
         }
     }
 }
@@ -127,6 +128,20 @@ impl AutomorphismKey<Vec<u8>, FFT64> {
         GLWESwitchingKey::keyswitch_inplace_scratch_space(module, out_size, out_rank, ksk_size)
     }
 
+    pub fn automorphism_scratch_space(
+        module: &Module<FFT64>,
+        out_size: usize,
+        in_size: usize,
+        ksk_size: usize,
+        rank: usize,
+    ) -> usize {
+        let tmp_dft: usize = module.bytes_of_vec_znx_dft(rank + 1, in_size);
+        let tmp_idft: usize = module.bytes_of_vec_znx_big(rank + 1, out_size);
+        let idft: usize = module.vec_znx_idft_tmp_bytes();
+        let keyswitch: usize = GLWECiphertext::keyswitch_inplace_scratch_space(module, out_size, rank, ksk_size);
+        tmp_dft + tmp_idft + idft + keyswitch
+    }
+
     pub fn external_product_scratch_space(
         module: &Module<FFT64>,
         out_size: usize,
@@ -267,7 +282,7 @@ where
 
                 // Reverts the automorphis key from (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
                 (0..cols_out).for_each(|i| {
-                    module.vec_znx_automorphism_inplace(self.p(), &mut tmp_idft_small_data, i);
+                    module.vec_znx_automorphism_inplace(lhs.p(), &mut tmp_idft_small_data, i);
                 });
 
                 // Wraps into ciphertext
@@ -283,7 +298,7 @@ where
                 // Applies back the automorphism X^{k}: (-pi^{-1}_{k'}(s)a + pi_{k}(s), a) -> (-pi^{-1}_{k'+k}(s)a + s, a)
                 // and switches back to DFT domain
                 (0..self.rank_out() + 1).for_each(|i| {
-                    module.vec_znx_automorphism_inplace(rhs.p(), &mut tmp_idft, i);
+                    module.vec_znx_automorphism_inplace(lhs.p(), &mut tmp_idft, i);
                     module.vec_znx_dft(&mut tmp_dft, i, &tmp_idft, i);
                 });
 
@@ -299,6 +314,8 @@ where
                 self.set_row(module, row_i, col_j, &tmp_dft);
             });
         });
+
+        self.p = (lhs.p * rhs.p) % (module.cyclotomic_order() as i64);
     }
 
     pub fn keyswitch<DataLhs, DataRhs>(

core/src/test_fft64/automorphism_key.rs

@@ -0,0 +1,117 @@
+use base2k::{FFT64, Module, ScalarZnxOps, ScalarZnxToRef, ScratchOwned, Stats, VecZnxOps, ZnxView};
+use sampling::source::Source;
+
+use crate::{
+    automorphism::AutomorphismKey,
+    elem::{GetRow, Infos},
+    glwe_ciphertext_fourier::GLWECiphertextFourier,
+    glwe_plaintext::GLWEPlaintext,
+    keys::{SecretKey, SecretKeyFourier},
+    test_fft64::gglwe::noise_gglwe_product,
+};
+
+#[test]
+fn automorphism() {
+    (1..4).for_each(|rank| {
+        println!("test automorphism rank: {}", rank);
+        test_automorphism(-1, 5, 12, 12, 60, 3.2, rank);
+    });
+}
+
+fn test_automorphism(p0: i64, p1: i64, log_n: usize, basek: usize, k_ksk: usize, sigma: f64, rank: usize) {
+    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
+    let rows = (k_ksk + basek - 1) / basek;
+
+    let mut auto_key_in: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::new(&module, basek, k_ksk, rows, rank);
+    let mut auto_key_out: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::new(&module, basek, k_ksk, rows, rank);
+    let mut auto_key_apply: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::new(&module, basek, k_ksk, rows, rank);
+
+    let mut source_xs: Source = Source::new([0u8; 32]);
+    let mut source_xe: Source = Source::new([0u8; 32]);
+    let mut source_xa: Source = Source::new([0u8; 32]);
+
+    let mut scratch: ScratchOwned = ScratchOwned::new(
+        AutomorphismKey::encrypt_sk_scratch_space(&module, rank, auto_key_in.size())
+            | GLWECiphertextFourier::decrypt_scratch_space(&module, auto_key_out.size())
+            | AutomorphismKey::automorphism_scratch_space(
+                &module,
+                auto_key_out.size(),
+                auto_key_in.size(),
+                auto_key_apply.size(),
+                rank,
+            ),
+    );
+
+    let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
+    sk.fill_ternary_prob(0.5, &mut source_xs);
+
+    let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
+    sk_dft.dft(&module, &sk);
+
+    // gglwe_{s1}(s0) = s0 -> s1
+    auto_key_in.encrypt_sk(
+        &module,
+        p0,
+        &sk,
+        &mut source_xa,
+        &mut source_xe,
+        sigma,
+        scratch.borrow(),
+    );
+
+    // gglwe_{s2}(s1) -> s1 -> s2
+    auto_key_apply.encrypt_sk(
+        &module,
+        p1,
+        &sk,
+        &mut source_xa,
+        &mut source_xe,
+        sigma,
+        scratch.borrow(),
+    );
+
+    // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
+    auto_key_out.automorphism(&module, &auto_key_in, &auto_key_apply, scratch.borrow());
+
+    let mut ct_glwe_dft: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::new(&module, basek, k_ksk, rank);
+    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ksk);
+
+    let mut sk_auto: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
+    sk_auto.fill_zero(); // Necessary to avoid panic of unfilled sk
+    (0..rank).for_each(|i| {
+        module.scalar_znx_automorphism(module.galois_element_inv(p0 * p1), &mut sk_auto, i, &sk, i);
+    });
+
+    let mut sk_auto_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
+    sk_auto_dft.dft(&module, &sk_auto);
+
+    (0..auto_key_out.rank_in()).for_each(|col_i| {
+        (0..auto_key_out.rows()).for_each(|row_i| {
+            auto_key_out.get_row(&module, row_i, col_i, &mut ct_glwe_dft);
+
+            ct_glwe_dft.decrypt(&module, &mut pt, &sk_auto_dft, scratch.borrow());
+            module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &sk, col_i);
+
+            let noise_have: f64 = pt.data.std(0, basek).log2();
+            let noise_want: f64 = noise_gglwe_product(
+                module.n() as f64,
+                basek,
+                0.5,
+                0.5,
+                0f64,
+                sigma * sigma,
+                0f64,
+                rank as f64,
+                k_ksk,
+                k_ksk,
+            );
+
+            assert!(
+                (noise_have - noise_want).abs() <= 0.1,
+                "{} {}",
+                noise_have,
+                noise_want
+            );
+        });
+    });
+}

core/src/test_fft64/gglwe.rs

@@ -405,6 +405,19 @@ fn test_external_product(log_n: usize, basek: usize, k: usize, sigma: f64, rank_
     // gglwe_(m) (x) RGSW_(X^k) = gglwe_(m * X^k)
     ct_gglwe_out.external_product(&module, &ct_gglwe_in, &ct_rgsw, scratch.borrow());
 
+    scratch = ScratchOwned::new(
+        GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank_out, ct_gglwe_in.size())
+            | GLWECiphertextFourier::decrypt_scratch_space(&module, ct_gglwe_out.size())
+            | GLWESwitchingKey::external_product_scratch_space(
+                &module,
+                ct_gglwe_out.size(),
+                ct_gglwe_in.size(),
+                ct_rgsw.size(),
+                rank_out,
+            )
+            | GGSWCiphertext::encrypt_sk_scratch_space(&module, rank_out, ct_rgsw.size()),
+    );
+
     let mut ct_glwe_dft: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::new(&module, basek, k, rank_out);
     let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k);
 

core/src/test_fft64/glwe.rs

@@ -437,7 +437,7 @@ fn test_automorphism(
     let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
     let rows: usize = (k_ct_in + basek - 1) / basek;
 
-    let mut autokey: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::new(&module, basek, p, k_autokey, rows, rank);
+    let mut autokey: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::new(&module, basek, k_autokey, rows, rank);
     let mut ct_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::new(&module, basek, k_ct_in, rank);
     let mut ct_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::new(&module, basek, k_ct_out, rank);
     let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ct_in);
@@ -523,19 +523,11 @@ fn automorphism_inplace() {
     });
 }
 
-fn test_automorphism_inplace(
+fn test_automorphism_inplace(log_n: usize, basek: usize, p: i64, k_autokey: usize, k_ct: usize, rank: usize, sigma: f64) {
-    log_n: usize,
-    basek: usize,
-    p: i64,
-    k_autokey: usize,
-    k_ct: usize,
-    rank: usize,
-    sigma: f64,
-) {
     let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
     let rows: usize = (k_ct + basek - 1) / basek;
 
-    let mut autokey: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::new(&module, basek, p, k_autokey, rows, rank);
+    let mut autokey: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::new(&module, basek, k_autokey, rows, rank);
     let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::new(&module, basek, k_ct, rank);
     let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ct);
     let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ct);

core/src/test_fft64/mod.rs

@@ -1,3 +1,4 @@
+mod automorphism_key;
 mod gglwe;
 mod ggsw;
 mod glwe;