M core/src/ggsw_ciphertext.rs → core/src/ggsw_ciphertext.rs

@@ -124,6 +124,16 @@ let res_dft: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
         res_znx + ci_dft + (ks | expand_rows | res_dft)
     }
 
+    pub fn keyswitch_inplace_scratch_space(
+        module: &Module<FFT64>,
+        out_size: usize,
+        ksk_size: usize,
+        tensor_key_size: usize,
+        rank: usize,
+    ) -> usize {
+        GGSWCiphertext::keyswitch_scratch_space(module, out_size, out_size, ksk_size, tensor_key_size, rank)
+    }
+
     pub fn automorphism_scratch_space(
         module: &Module<FFT64>,
         out_size: usize,
@@ -136,6 +146,23 @@ GGSWCiphertext::keyswitch_scratch_space(
             module,
             out_size,
             in_size,
+            auto_key_size,
+            tensor_key_size,
+            rank,
+        )
+    }
+
+    pub fn automorphism_inplace_scratch_space(
+        module: &Module<FFT64>,
+        out_size: usize,
+        auto_key_size: usize,
+        tensor_key_size: usize,
+        rank: usize,
+    ) -> usize {
+        GGSWCiphertext::automorphism_scratch_space(
+            module,
+            out_size,
+            out_size,
             auto_key_size,
             tensor_key_size,
             rank,
@@ -389,6 +416,22 @@ })
         })
     }
 
+    pub fn keyswitch_inplace<DataKsk, DataTsk>(
+        &mut self,
+        module: &Module<FFT64>,
+        ksk: &GLWESwitchingKey<DataKsk, FFT64>,
+        tsk: &TensorKey<DataTsk, FFT64>,
+        scratch: &mut Scratch,
+    ) where
+        MatZnxDft<DataKsk, FFT64>: MatZnxDftToRef<FFT64>,
+        MatZnxDft<DataTsk, FFT64>: MatZnxDftToRef<FFT64>,
+    {
+        unsafe {
+            let self_ptr: *mut GGSWCiphertext<DataSelf, FFT64> = self as *mut GGSWCiphertext<DataSelf, FFT64>;
+            self.keyswitch(module, &*self_ptr, ksk, tsk, scratch);
+        }
+    }
+
     pub fn automorphism<DataLhs, DataAk, DataTsk>(
         &mut self,
         module: &Module<FFT64>,
@@ -468,6 +511,22 @@
                 self.set_row(module, row_i, col_j, &res_dft);
             })
         })
+    }
+
+    pub fn automorphism_inplace<DataKsk, DataTsk>(
+        &mut self,
+        module: &Module<FFT64>,
+        auto_key: &AutomorphismKey<DataKsk, FFT64>,
+        tensor_key: &TensorKey<DataTsk, FFT64>,
+        scratch: &mut Scratch,
+    ) where
+        MatZnxDft<DataKsk, FFT64>: MatZnxDftToRef<FFT64>,
+        MatZnxDft<DataTsk, FFT64>: MatZnxDftToRef<FFT64>,
+    {
+        unsafe {
+            let self_ptr: *mut GGSWCiphertext<DataSelf, FFT64> = self as *mut GGSWCiphertext<DataSelf, FFT64>;
+            self.automorphism(module, &*self_ptr, auto_key, tensor_key, scratch);
+        }
     }
 
     pub fn external_product<DataLhs, DataRhs>(

M core/src/test_fft64/ggsw.rs → core/src/test_fft64/ggsw.rs

@@ -26,6 +26,38 @@ });
 }
 
 #[test]
+fn keyswitch() {
+    (1..4).for_each(|rank| {
+        println!("test keyswitch rank: {}", rank);
+        test_keyswitch(12, 15, 60, rank, 3.2);
+    });
+}
+
+#[test]
+fn keyswitch_inplace() {
+    (1..4).for_each(|rank| {
+        println!("test keyswitch_inplace rank: {}", rank);
+        test_keyswitch_inplace(12, 15, 60, rank, 3.2);
+    });
+}
+
+#[test]
+fn automorphism() {
+    (1..4).for_each(|rank| {
+        println!("test automorphism rank: {}", rank);
+        test_automorphism(-5, 12, 15, 60, rank, 3.2);
+    });
+}
+
+#[test]
+fn automorphism_inplace() {
+    (1..4).for_each(|rank| {
+        println!("test automorphism_inplace rank: {}", rank);
+        test_automorphism_inplace(-5, 12, 15, 60, rank, 3.2);
+    });
+}
+
+#[test]
 fn external_product() {
     (1..4).for_each(|rank| {
         println!("test external_product rank: {}", rank);
@@ -108,14 +140,6 @@ });
     });
 }
 
-#[test]
-fn keyswitch() {
-    (1..4).for_each(|rank| {
-        println!("test keyswitch rank: {}", rank);
-        test_keyswitch(12, 15, 60, rank, 3.2);
-    });
-}
-
 fn test_keyswitch(log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64) {
     let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
     let rows: usize = (k + basek - 1) / basek;
@@ -243,6 +267,125 @@ });
     });
 }
 
+fn test_keyswitch_inplace(log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64) {
+    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
+    let rows: usize = (k + basek - 1) / basek;
+
+    let mut ct: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k, rows, rank);
+    let mut tsk: TensorKey<Vec<u8>, FFT64> = TensorKey::new(&module, basek, k, rows, rank);
+    let mut ksk: GLWESwitchingKey<Vec<u8>, FFT64> = GLWESwitchingKey::new(&module, basek, k, rows, rank, rank);
+    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k);
+    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k);
+    let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
+
+    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(
+        GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct.size())
+            | GLWECiphertextFourier::decrypt_scratch_space(&module, ct.size())
+            | GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank, ksk.size())
+            | TensorKey::encrypt_sk_scratch_space(&module, rank, ksk.size())
+            | GGSWCiphertext::keyswitch_inplace_scratch_space(&module, ct.size(), ksk.size(), tsk.size(), rank),
+    );
+
+    let var_xs: f64 = 0.5;
+
+    let mut sk_in: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
+    sk_in.fill_ternary_prob(var_xs, &mut source_xs);
+
+    let mut sk_in_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
+    sk_in_dft.dft(&module, &sk_in);
+
+    let mut sk_out: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
+    sk_out.fill_ternary_prob(var_xs, &mut source_xs);
+
+    let mut sk_out_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
+    sk_out_dft.dft(&module, &sk_out);
+
+    ksk.encrypt_sk(
+        &module,
+        &sk_in,
+        &sk_out_dft,
+        &mut source_xa,
+        &mut source_xe,
+        sigma,
+        scratch.borrow(),
+    );
+    tsk.encrypt_sk(
+        &module,
+        &sk_out_dft,
+        &mut source_xa,
+        &mut source_xe,
+        sigma,
+        scratch.borrow(),
+    );
+
+    pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
+
+    ct.encrypt_sk(
+        &module,
+        &pt_scalar,
+        &sk_in_dft,
+        &mut source_xa,
+        &mut source_xe,
+        sigma,
+        scratch.borrow(),
+    );
+
+    ct.keyswitch_inplace(&module, &ksk, &tsk, scratch.borrow());
+
+    let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::new(&module, basek, k, rank);
+    let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct.size());
+    let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct.size());
+
+    (0..ct.rank() + 1).for_each(|col_j| {
+        (0..ct.rows()).for_each(|row_i| {
+            module.vec_znx_add_scalar_inplace(&mut pt_want, 0, row_i, &pt_scalar, 0);
+
+            // mul with sk[col_j-1]
+            if col_j > 0 {
+                module.vec_znx_dft(&mut pt_dft, 0, &pt_want, 0);
+                module.svp_apply_inplace(&mut pt_dft, 0, &sk_out_dft, col_j - 1);
+                module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
+                module.vec_znx_big_normalize(basek, &mut pt_want, 0, &pt_big, 0, scratch.borrow());
+            }
+
+            ct.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
+
+            ct_glwe_fourier.decrypt(&module, &mut pt_have, &sk_out_dft, scratch.borrow());
+
+            module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);
+
+            let noise_have: f64 = pt_have.data.std(0, basek).log2();
+            let noise_want: f64 = noise_ggsw_keyswitch(
+                module.n() as f64,
+                basek,
+                col_j,
+                var_xs,
+                0f64,
+                sigma * sigma,
+                0f64,
+                rank as f64,
+                k,
+                k,
+            );
+
+            println!("{} {}", noise_have, noise_want);
+
+            assert!(
+                (noise_have - noise_want).abs() <= 0.1,
+                "{} {}",
+                noise_have,
+                noise_want
+            );
+
+            pt_want.data.zero();
+        });
+    });
+}
+
 pub(crate) fn noise_ggsw_keyswitch(
     n: f64,
     basek: usize,
@@ -290,14 +433,6 @@ }
 
     noise = noise.sqrt();
     noise.log2().min(-1.0) // max noise is [-2^{-1}, 2^{-1}]
-}
-
-#[test]
-fn automorphism() {
-    (1..4).for_each(|rank| {
-        println!("test automorphism rank: {}", rank);
-        test_automorphism(-5, 12, 15, 60, rank, 3.2);
-    });
 }
 
 fn test_automorphism(p: i64, log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64) {
@@ -390,6 +525,119 @@ module.vec_znx_big_normalize(basek, &mut pt_want, 0, &pt_big, 0, scratch.borrow());
             }
 
             ct_out.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
+
+            ct_glwe_fourier.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());
+
+            module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);
+
+            let noise_have: f64 = pt_have.data.std(0, basek).log2();
+            let noise_want: f64 = noise_ggsw_keyswitch(
+                module.n() as f64,
+                basek,
+                col_j,
+                var_xs,
+                0f64,
+                sigma * sigma,
+                0f64,
+                rank as f64,
+                k,
+                k,
+            );
+
+            assert!(
+                (noise_have - noise_want).abs() <= 0.1,
+                "{} {}",
+                noise_have,
+                noise_want
+            );
+
+            pt_want.data.zero();
+        });
+    });
+}
+
+fn test_automorphism_inplace(p: i64, log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64) {
+    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
+    let rows: usize = (k + basek - 1) / basek;
+
+    let mut ct: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k, rows, rank);
+    let mut tensor_key: TensorKey<Vec<u8>, FFT64> = TensorKey::new(&module, basek, k, rows, rank);
+    let mut auto_key: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::new(&module, basek, k, rows, rank);
+    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k);
+    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k);
+    let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
+
+    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(
+        GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct.size())
+            | GLWECiphertextFourier::decrypt_scratch_space(&module, ct.size())
+            | AutomorphismKey::encrypt_sk_scratch_space(&module, rank, auto_key.size())
+            | TensorKey::encrypt_sk_scratch_space(&module, rank, tensor_key.size())
+            | GGSWCiphertext::automorphism_inplace_scratch_space(&module, ct.size(), auto_key.size(), tensor_key.size(), rank),
+    );
+
+    let var_xs: f64 = 0.5;
+
+    let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
+    sk.fill_ternary_prob(var_xs, &mut source_xs);
+
+    let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
+    sk_dft.dft(&module, &sk);
+
+    auto_key.encrypt_sk(
+        &module,
+        p,
+        &sk,
+        &mut source_xa,
+        &mut source_xe,
+        sigma,
+        scratch.borrow(),
+    );
+    tensor_key.encrypt_sk(
+        &module,
+        &sk_dft,
+        &mut source_xa,
+        &mut source_xe,
+        sigma,
+        scratch.borrow(),
+    );
+
+    pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
+
+    ct.encrypt_sk(
+        &module,
+        &pt_scalar,
+        &sk_dft,
+        &mut source_xa,
+        &mut source_xe,
+        sigma,
+        scratch.borrow(),
+    );
+
+    ct.automorphism_inplace(&module, &auto_key, &tensor_key, scratch.borrow());
+
+    module.scalar_znx_automorphism_inplace(p, &mut pt_scalar, 0);
+
+    let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::new(&module, basek, k, rank);
+    let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct.size());
+    let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct.size());
+
+    (0..ct.rank() + 1).for_each(|col_j| {
+        (0..ct.rows()).for_each(|row_i| {
+            module.vec_znx_add_scalar_inplace(&mut pt_want, 0, row_i, &pt_scalar, 0);
+
+            // mul with sk[col_j-1]
+            if col_j > 0 {
+                module.vec_znx_dft(&mut pt_dft, 0, &pt_want, 0);
+                module.svp_apply_inplace(&mut pt_dft, 0, &sk_dft, col_j - 1);
+                module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
+                module.vec_znx_big_normalize(basek, &mut pt_want, 0, &pt_big, 0, scratch.borrow());
+            }
+
+            ct.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
 
             ct_glwe_fourier.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());