poulpy · arnaucube's git mirror

wip

Jean-Philippe Bossuat jeanphilippe.bossuat@gmail.com

Sun, 25 May 2025 11:15:59 +0200

commit

cb284a4c4c50736c2e1d62912689f94c12b708a3

parent

43c7f21033cdb3dbf363f0346f66cf9934f316a5

6 files changed, 83 insertions(+), 24 deletions(-)

jump to

backend/src/vec_znx.rs

core/src/glwe_ciphertext.rs

core/src/glwe_ops.rs

core/src/tensor_key.rs

core/src/test_fft64/ggsw.rs

core/src/test_fft64/tensor_key.rs

M backend/src/vec_znx.rs → backend/src/vec_znx.rs

@@ -71,7 +71,7 @@ impl<D: AsRef<[u8]>> ZnxView for VecZnx<D> {
     type Scalar = i64;
 }
 
-impl<D: AsRef<[u8]>> VecZnx<D> {
+impl VecZnx<Vec<u8>> {
     pub fn rsh_scratch_space(n: usize) -> usize {
         n * std::mem::size_of::<i64>()
     }

M core/src/glwe_ciphertext.rs → core/src/glwe_ciphertext.rs

@@ -334,7 +334,7 @@ ) where
         VecZnx<DataLhs>: VecZnxToRef,
         MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
     {
-        Self::keyswitch_private(self, true, rhs.p(), module, lhs, &rhs.key, scratch);
+        Self::keyswitch_private::<_, _, 1>(self, rhs.p(), module, lhs, &rhs.key, scratch);
     }
 
     pub fn automorphism_add_inplace<DataRhs>(
@@ -347,7 +347,61 @@ MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
     {
         unsafe {
             let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
-            Self::keyswitch_private(self, true, rhs.p(), module, &*self_ptr, &rhs.key, scratch);
+            Self::keyswitch_private::<_, _, 1>(self, rhs.p(), module, &*self_ptr, &rhs.key, scratch);
+        }
+    }
+
+    pub fn automorphism_sub_ab<DataLhs, DataRhs>(
+        &mut self,
+        module: &Module<FFT64>,
+        lhs: &GLWECiphertext<DataLhs>,
+        rhs: &AutomorphismKey<DataRhs, FFT64>,
+        scratch: &mut Scratch,
+    ) where
+        VecZnx<DataLhs>: VecZnxToRef,
+        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
+    {
+        Self::keyswitch_private::<_, _, 2>(self, rhs.p(), module, lhs, &rhs.key, scratch);
+    }
+
+    pub fn automorphism_sub_ab_inplace<DataRhs>(
+        &mut self,
+        module: &Module<FFT64>,
+        rhs: &AutomorphismKey<DataRhs, FFT64>,
+        scratch: &mut Scratch,
+    ) where
+        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
+    {
+        unsafe {
+            let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
+            Self::keyswitch_private::<_, _, 2>(self, rhs.p(), module, &*self_ptr, &rhs.key, scratch);
+        }
+    }
+
+    pub fn automorphism_sub_ba<DataLhs, DataRhs>(
+        &mut self,
+        module: &Module<FFT64>,
+        lhs: &GLWECiphertext<DataLhs>,
+        rhs: &AutomorphismKey<DataRhs, FFT64>,
+        scratch: &mut Scratch,
+    ) where
+        VecZnx<DataLhs>: VecZnxToRef,
+        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
+    {
+        Self::keyswitch_private::<_, _, 3>(self, rhs.p(), module, lhs, &rhs.key, scratch);
+    }
+
+    pub fn automorphism_sub_ba_inplace<DataRhs>(
+        &mut self,
+        module: &Module<FFT64>,
+        rhs: &AutomorphismKey<DataRhs, FFT64>,
+        scratch: &mut Scratch,
+    ) where
+        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
+    {
+        unsafe {
+            let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
+            Self::keyswitch_private::<_, _, 3>(self, rhs.p(), module, &*self_ptr, &rhs.key, scratch);
         }
     }
 
@@ -420,12 +474,11 @@ ) where
         VecZnx<DataLhs>: VecZnxToRef,
         MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
     {
-        Self::keyswitch_private(self, false, 0, module, lhs, rhs, scratch);
+        Self::keyswitch_private::<_, _, 0>(self, 0, module, lhs, rhs, scratch);
     }
 
-    pub(crate) fn keyswitch_private<DataLhs, DataRhs>(
+    pub(crate) fn keyswitch_private<DataLhs, DataRhs, const OP: u8>(
         &mut self,
-        add_self: bool,
         apply_auto: i64,
         module: &Module<FFT64>,
         lhs: &GLWECiphertext<DataLhs>,
@@ -481,8 +534,11 @@ if apply_auto != 0 {
                 module.vec_znx_big_automorphism_inplace(apply_auto, &mut res_big, i);
             }
 
-            if add_self {
-                module.vec_znx_big_add_small_inplace(&mut res_big, i, lhs, i);
+            match OP{
+                1=> module.vec_znx_big_add_small_inplace(&mut res_big, i, lhs, i),
+                2=> module.vec_znx_big_sub_small_a_inplace(&mut res_big, i, lhs, i),
+                3=> module.vec_znx_big_sub_small_b_inplace(&mut res_big, i, lhs, i),
+                _=>{},
             }
             module.vec_znx_big_normalize(basek, self, i, &res_big, i, scratch1);
         });

M core/src/glwe_ops.rs → core/src/glwe_ops.rs

@@ -175,10 +175,7 @@ self.set_basek(a.basek());
         self.set_k(a.k());
     }
 
-    pub fn rotate_inplace<A>(&mut self, module: &Module<FFT64>, k: i64)
-    where
-        A: VecZnxToRef + Infos,
-    {
+    pub fn rotate_inplace(&mut self, module: &Module<FFT64>, k: i64){
         #[cfg(debug_assertions)]
         {
             assert_eq!(self.n(), module.n());
@@ -242,3 +239,9 @@ module.vec_znx_normalize_inplace(self.basek(), self, i, scratch);
         });
     }
 }
+
+impl GLWECiphertext<Vec<u8>>{
+    pub fn rsh_scratch_space(module: &Module<FFT64>) -> usize{
+        VecZnx::rsh_scratch_space(module.n())
+    }
+}

M core/src/tensor_key.rs → core/src/tensor_key.rs

@@ -56,7 +56,7 @@ }
 }
 
 impl TensorKey<Vec<u8>, FFT64> {
-    pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, rank: usize, size: usize) -> usize {
+    pub fn generate_from_sk_scratch_space(module: &Module<FFT64>, rank: usize, size: usize) -> usize {
         module.bytes_of_scalar_znx_dft(1) + GLWESwitchingKey::encrypt_sk_scratch_space(module, rank, size)
     }
 }
@@ -65,7 +65,7 @@ impl<DataSelf> TensorKey<DataSelf, FFT64>
 where
     MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64>,
 {
-    pub fn encrypt_sk<DataSk>(
+    pub fn generate_from_sk<DataSk>(
         &mut self,
         module: &Module<FFT64>,
         sk_dft: &SecretKeyFourier<DataSk, FFT64>,

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

@@ -160,7 +160,7 @@ let mut scratch: ScratchOwned = ScratchOwned::new(
         GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct_in.size())
             | GLWECiphertextFourier::decrypt_scratch_space(&module, ct_out.size())
             | GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank, ksk.size())
-            | TensorKey::encrypt_sk_scratch_space(&module, rank, ksk.size())
+            | TensorKey::generate_from_sk_scratch_space(&module, rank, ksk.size())
             | GGSWCiphertext::keyswitch_scratch_space(
                 &module,
                 ct_out.size(),
@@ -194,7 +194,7 @@ &mut source_xe,
         sigma,
         scratch.borrow(),
     );
-    tsk.encrypt_sk(
+    tsk.generate_from_sk(
         &module,
         &sk_out_dft,
         &mut source_xa,
@@ -286,7 +286,7 @@ 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())
+            | TensorKey::generate_from_sk_scratch_space(&module, rank, ksk.size())
             | GGSWCiphertext::keyswitch_inplace_scratch_space(&module, ct.size(), ksk.size(), tsk.size(), rank),
     );
 
@@ -313,7 +313,7 @@ &mut source_xe,
         sigma,
         scratch.borrow(),
     );
-    tsk.encrypt_sk(
+    tsk.generate_from_sk(
         &module,
         &sk_out_dft,
         &mut source_xa,
@@ -455,7 +455,7 @@ let mut scratch: ScratchOwned = ScratchOwned::new(
         GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct_in.size())
             | GLWECiphertextFourier::decrypt_scratch_space(&module, ct_out.size())
             | AutomorphismKey::generate_from_sk_scratch_space(&module, rank, auto_key.size())
-            | TensorKey::encrypt_sk_scratch_space(&module, rank, tensor_key.size())
+            | TensorKey::generate_from_sk_scratch_space(&module, rank, tensor_key.size())
             | GGSWCiphertext::automorphism_scratch_space(
                 &module,
                 ct_out.size(),
@@ -483,7 +483,7 @@ &mut source_xe,
         sigma,
         scratch.borrow(),
     );
-    tensor_key.encrypt_sk(
+    tensor_key.generate_from_sk(
         &module,
         &sk_dft,
         &mut source_xa,
@@ -575,7 +575,7 @@ let mut scratch: ScratchOwned = ScratchOwned::new(
         GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct.size())
             | GLWECiphertextFourier::decrypt_scratch_space(&module, ct.size())
             | AutomorphismKey::generate_from_sk_scratch_space(&module, rank, auto_key.size())
-            | TensorKey::encrypt_sk_scratch_space(&module, rank, tensor_key.size())
+            | TensorKey::generate_from_sk_scratch_space(&module, rank, tensor_key.size())
             | GGSWCiphertext::automorphism_inplace_scratch_space(&module, ct.size(), auto_key.size(), tensor_key.size(), rank),
     );
 
@@ -596,7 +596,7 @@ &mut source_xe,
         sigma,
         scratch.borrow(),
     );
-    tensor_key.encrypt_sk(
+    tensor_key.generate_from_sk(
         &module,
         &sk_dft,
         &mut source_xa,

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

@@ -30,7 +30,7 @@ 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(TensorKey::encrypt_sk_scratch_space(
+    let mut scratch: ScratchOwned = ScratchOwned::new(TensorKey::generate_from_sk_scratch_space(
         &module,
         rank,
         tensor_key.size(),
@@ -40,7 +40,7 @@ let mut sk: SecretKey<Vec<u8>> = SecretKey::alloc(&module, rank);
     sk.fill_ternary_prob(0.5, &mut source_xs);
 
     let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::alloc(&module, rank);
-    sk_dft.dft(&module, &sk);
+    sk_dft.dft(generate_from_sksk);
 
     tensor_key.encrypt_sk(
         &module,