refactor(tfhe): reproduce sequential BSK generation

2026-01-08 22:28:01 -05:00 · 2022-11-25 10:25:02 +01:00
parent d0d0b542ac
commit 2aa84d2b3c
5 changed files with 267 additions and 11 deletions
--- a/tfhe/src/core_crypto/algorithms/ggsw_encryption.rs
+++ b/tfhe/src/core_crypto/algorithms/ggsw_encryption.rs
@@ -20,8 +20,6 @@ pub fn encrypt_ggsw_ciphertext<Scalar, KeyCont, OutputCont, Gen>(
    KeyCont: Container<Element = Scalar>,
    OutputCont: ContainerMut<Element = Scalar>,
 {
-    // ck_dim_eq!(self.key_size() => output.glwe_size().to_glwe_dimension());
-
    assert!(
        output.polynomial_size() == glwe_secret_key.polynomial_size(),
        "Mismatch between polynomial sizes of output cipertexts and input secret key. \
@@ -128,11 +126,12 @@ fn encrypt_ggsw_level_matrix_row<Scalar, KeyCont, InputCont, OutputCont, Gen>(
 mod test {
    use crate::core_crypto::algorithms::encrypt_ggsw_ciphertext;
    use crate::core_crypto::commons::crypto::encoding::PlaintextList;
+    use crate::core_crypto::commons::crypto::ggsw::StandardGgswCiphertext;
    use crate::core_crypto::commons::crypto::secret::generators::{
        DeterministicSeeder, EncryptionRandomGenerator,
    };
    use crate::core_crypto::commons::crypto::secret::GlweSecretKey;
-    use crate::core_crypto::commons::math::random::Seeder;
+    use crate::core_crypto::commons::math::random::Seed;
    use crate::core_crypto::commons::math::tensor::*;
    use crate::core_crypto::commons::math::torus::UnsignedTorus;
    use crate::core_crypto::commons::test_tools;
@@ -142,8 +141,6 @@ mod test {
    };
    use concrete_csprng::generators::SoftwareRandomGenerator;

-    use crate::core_crypto::commons::crypto::ggsw::StandardGgswCiphertext;
-
    fn test_refactored_ggsw<T: UnsignedTorus>() {
        // random settings
        let nb_ct = test_tools::random_ciphertext_count(10);
@@ -163,12 +160,7 @@ mod test {

        for plaintext in plaintext_vector.plaintext_iter() {
            let main_seed = test_tools::random_seed();
-
-            // Use a deterministic seeder to get the seeds that will be used during the tests
-            let mut deterministic_seeder =
-                DeterministicSeeder::<SoftwareRandomGenerator>::new(main_seed);
-            let _ = deterministic_seeder.seed();
-            let mask_seed = deterministic_seeder.seed();
+            let mask_seed = Seed(crate::core_crypto::commons::test_tools::any_usize() as u128);

            let mut generator = EncryptionRandomGenerator::<SoftwareRandomGenerator>::new(
                mask_seed,
--- a/tfhe/src/core_crypto/algorithms/lwe_bootstrap_key_generation.rs
+++ b/tfhe/src/core_crypto/algorithms/lwe_bootstrap_key_generation.rs
@@ -0,0 +1,175 @@
+use crate::core_crypto::algorithms::*;
+use crate::core_crypto::commons::crypto::secret::generators::EncryptionRandomGenerator;
+use crate::core_crypto::commons::math::random::ByteRandomGenerator;
+use crate::core_crypto::commons::math::torus::UnsignedTorus;
+use crate::core_crypto::commons::traits::*;
+use crate::core_crypto::entities::*;
+use crate::core_crypto::specification::dispersion::DispersionParameter;
+
+pub fn generate_lwe_bootstrap_key<Scalar, InputKeyCont, OutputKeyCont, OutputCont, Gen>(
+    input_lwe_secret_key: &LweSecretKeyBase<InputKeyCont>,
+    output_glwe_secret_key: &GlweSecretKeyBase<OutputKeyCont>,
+    output: &mut LweBootstrapKeyBase<OutputCont>,
+    noise_parameters: impl DispersionParameter,
+    generator: &mut EncryptionRandomGenerator<Gen>,
+) where
+    Scalar: UnsignedTorus,
+    InputKeyCont: Container<Element = Scalar>,
+    OutputKeyCont: Container<Element = Scalar>,
+    OutputCont: ContainerMut<Element = Scalar>,
+    Gen: ByteRandomGenerator,
+{
+    assert!(
+        output.input_lwe_dimension() == input_lwe_secret_key.lwe_dimension(),
+        "Mismatched LweDimension between input LWE secret key and LWE bootstrap key. \
+        Input LWE secret key LweDimension: {:?}, LWE bootstrap key input LweDimension {:?}.",
+        input_lwe_secret_key.lwe_dimension(),
+        output.input_lwe_dimension()
+    );
+
+    assert!(
+        output.glwe_size() == output_glwe_secret_key.glwe_dimension().to_glwe_size(),
+        "Mismatched GlweSize between output GLWE secret key and LWE bootstrap key. \
+        Output GLWE secret key GlweSize: {:?}, LWE bootstrap key GlweSize {:?}.",
+        output_glwe_secret_key.glwe_dimension().to_glwe_size(),
+        output.glwe_size()
+    );
+
+    assert!(
+        output.polynomial_size() == output_glwe_secret_key.polynomial_size(),
+        "Mismatched PolynomialSize between output GLWE secret key and LWE bootstrap key. \
+        Output GLWE secret key PolynomialSize: {:?}, LWE bootstrap key PolynomialSize {:?}.",
+        output_glwe_secret_key.polynomial_size(),
+        output.polynomial_size()
+    );
+
+    let gen_iter = generator
+        .fork_bsk_to_ggsw::<Scalar>(
+            output.input_lwe_dimension(),
+            output.decomposition_level_count(),
+            output.glwe_size(),
+            output.polynomial_size(),
+        )
+        .unwrap();
+
+    for ((mut ggsw, &input_key_element), mut generator) in output
+        .iter_mut()
+        .zip(input_lwe_secret_key.as_ref())
+        .zip(gen_iter)
+    {
+        encrypt_ggsw_ciphertext(
+            output_glwe_secret_key,
+            &mut ggsw,
+            Plaintext(input_key_element),
+            noise_parameters,
+            &mut generator,
+        );
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::core_crypto::algorithms::generate_lwe_bootstrap_key;
+    use crate::core_crypto::commons::crypto::bootstrap::StandardBootstrapKey;
+    use crate::core_crypto::commons::crypto::secret::generators::{
+        DeterministicSeeder, EncryptionRandomGenerator,
+    };
+    use crate::core_crypto::commons::crypto::secret::{GlweSecretKey, LweSecretKey};
+    use crate::core_crypto::commons::math::random::Seed;
+    use crate::core_crypto::commons::math::tensor::*;
+    use crate::core_crypto::commons::math::torus::UnsignedTorus;
+    use crate::core_crypto::commons::test_tools::new_secret_random_generator;
+    use crate::core_crypto::entities::{GlweSecretKeyBase, LweBootstrapKey, LweSecretKeyBase};
+    use crate::core_crypto::prelude::{
+        DecompositionBaseLog, DecompositionLevelCount, GlweDimension, LweDimension, PolynomialSize,
+        StandardDev,
+    };
+    use concrete_csprng::generators::SoftwareRandomGenerator;
+
+    fn test_refactored_bsk_equivalence<T: UnsignedTorus + Send + Sync>() {
+        for _ in 0..10 {
+            let lwe_dim =
+                LweDimension(crate::core_crypto::commons::test_tools::random_usize_between(5..10));
+            let glwe_dim =
+                GlweDimension(crate::core_crypto::commons::test_tools::random_usize_between(5..10));
+            let poly_size = PolynomialSize(
+                crate::core_crypto::commons::test_tools::random_usize_between(5..10),
+            );
+            let level = DecompositionLevelCount(
+                crate::core_crypto::commons::test_tools::random_usize_between(2..5),
+            );
+            let base_log = DecompositionBaseLog(
+                crate::core_crypto::commons::test_tools::random_usize_between(2..5),
+            );
+            let mask_seed = Seed(crate::core_crypto::commons::test_tools::any_usize() as u128);
+            let deterministic_seeder_seed =
+                Seed(crate::core_crypto::commons::test_tools::any_usize() as u128);
+
+            let mut secret_generator = new_secret_random_generator();
+            let lwe_sk = LweSecretKey::generate_binary(lwe_dim, &mut secret_generator);
+            let glwe_sk =
+                GlweSecretKey::generate_binary(glwe_dim, poly_size, &mut secret_generator);
+
+            let mut bsk = StandardBootstrapKey::allocate(
+                T::ZERO,
+                glwe_dim.to_glwe_size(),
+                poly_size,
+                level,
+                base_log,
+                lwe_dim,
+            );
+
+            let mut encryption_generator =
+                EncryptionRandomGenerator::<SoftwareRandomGenerator>::new(
+                    mask_seed,
+                    &mut DeterministicSeeder::<SoftwareRandomGenerator>::new(
+                        deterministic_seeder_seed,
+                    ),
+                );
+
+            bsk.fill_with_new_key(
+                &lwe_sk,
+                &glwe_sk,
+                StandardDev::from_standard_dev(10.),
+                &mut encryption_generator,
+            );
+
+            let mut refactored_bsk = LweBootstrapKey::new(
+                T::ZERO,
+                glwe_dim.to_glwe_size(),
+                poly_size,
+                base_log,
+                level,
+                lwe_dim,
+            );
+
+            let mut encryption_generator =
+                EncryptionRandomGenerator::<SoftwareRandomGenerator>::new(
+                    mask_seed,
+                    &mut DeterministicSeeder::<SoftwareRandomGenerator>::new(
+                        deterministic_seeder_seed,
+                    ),
+                );
+
+            generate_lwe_bootstrap_key(
+                &LweSecretKeyBase::from_container(lwe_sk.as_tensor().as_slice()),
+                &GlweSecretKeyBase::from_container(glwe_sk.as_tensor().as_slice(), poly_size),
+                &mut refactored_bsk,
+                StandardDev::from_standard_dev(10.),
+                &mut encryption_generator,
+            );
+
+            assert_eq!(bsk.as_tensor().as_slice(), refactored_bsk.as_ref());
+        }
+    }
+
+    #[test]
+    fn test_refactored_bsk_u32() {
+        test_refactored_bsk_equivalence::<u32>()
+    }
+
+    #[test]
+    fn test_refactored_bsk_u64() {
+        test_refactored_bsk_equivalence::<u64>()
+    }
+}
--- a/tfhe/src/core_crypto/algorithms/mod.rs
+++ b/tfhe/src/core_crypto/algorithms/mod.rs
@@ -1,6 +1,7 @@
 pub mod ggsw_encryption;
 pub mod glwe_encryption;
 pub mod glwe_secret_key_generation;
+pub mod lwe_bootstrap_key_generation;
 pub mod lwe_ciphertext_keyswitch;
 pub mod lwe_encryption;
 pub mod lwe_keyswitch_key_generation;
@@ -14,6 +15,7 @@ pub mod slice_algorithms;
 pub use ggsw_encryption::*;
 pub use glwe_encryption::*;
 pub use glwe_secret_key_generation::*;
+pub use lwe_bootstrap_key_generation::*;
 pub use lwe_ciphertext_keyswitch::*;
 pub use lwe_encryption::*;
 pub use lwe_keyswitch_key_generation::*;
--- a/tfhe/src/core_crypto/entities/lwe_bootstrap_key.rs
+++ b/tfhe/src/core_crypto/entities/lwe_bootstrap_key.rs
@@ -0,0 +1,85 @@
+use crate::core_crypto::commons::traits::*;
+use crate::core_crypto::entities::*;
+use crate::core_crypto::specification::parameters::*;
+
+// An LweBootstrapKey is literally a GgswCiphertextList, so we wrap a GgswCiphetextList and use
+// Deref to have access to all the primitives of the GgswCiphertextList easily
+pub struct LweBootstrapKeyBase<C: Container> {
+    ggsw_list: GgswCiphertextListBase<C>,
+}
+
+impl<C: Container> std::ops::Deref for LweBootstrapKeyBase<C> {
+    type Target = GgswCiphertextListBase<C>;
+
+    fn deref(&self) -> &GgswCiphertextListBase<C> {
+        self.as_ggsw_ciphertext_list()
+    }
+}
+
+impl<C: ContainerMut> std::ops::DerefMut for LweBootstrapKeyBase<C> {
+    fn deref_mut(&mut self) -> &mut GgswCiphertextListBase<C> {
+        self.as_mut_ggsw_ciphertext_list()
+    }
+}
+
+impl<C: Container> LweBootstrapKeyBase<C> {
+    pub fn from_container(
+        container: C,
+        glwe_size: GlweSize,
+        polynomial_size: PolynomialSize,
+        decomp_base_log: DecompositionBaseLog,
+        decomp_level_count: DecompositionLevelCount,
+    ) -> LweBootstrapKeyBase<C> {
+        LweBootstrapKeyBase {
+            ggsw_list: GgswCiphertextListBase::from_container(
+                container,
+                glwe_size,
+                polynomial_size,
+                decomp_base_log,
+                decomp_level_count,
+            ),
+        }
+    }
+
+    pub fn input_lwe_dimension(&self) -> LweDimension {
+        LweDimension(self.ggsw_ciphertext_count().0)
+    }
+
+    pub fn output_lwe_dimension(&self) -> LweDimension {
+        LweDimension(self.glwe_size().to_glwe_dimension().0 * self.polynomial_size().0)
+    }
+
+    pub fn as_ggsw_ciphertext_list(&self) -> &GgswCiphertextListBase<C> {
+        &self.ggsw_list
+    }
+}
+
+impl<C: ContainerMut> LweBootstrapKeyBase<C> {
+    pub fn as_mut_ggsw_ciphertext_list(&mut self) -> &mut GgswCiphertextListBase<C> {
+        &mut self.ggsw_list
+    }
+}
+
+pub type LweBootstrapKey<Scalar> = LweBootstrapKeyBase<Vec<Scalar>>;
+
+impl<Scalar: Copy> LweBootstrapKey<Scalar> {
+    pub fn new(
+        fill_with: Scalar,
+        glwe_size: GlweSize,
+        polynomial_size: PolynomialSize,
+        decomp_base_log: DecompositionBaseLog,
+        decomp_level_count: DecompositionLevelCount,
+        input_lwe_dimension: LweDimension,
+    ) -> LweBootstrapKey<Scalar> {
+        LweBootstrapKey {
+            ggsw_list: GgswCiphertextListBase::new(
+                fill_with,
+                glwe_size,
+                polynomial_size,
+                decomp_base_log,
+                decomp_level_count,
+                GgswCiphertextCount(input_lwe_dimension.0),
+            ),
+        }
+    }
+}
--- a/tfhe/src/core_crypto/entities/mod.rs
+++ b/tfhe/src/core_crypto/entities/mod.rs
@@ -4,6 +4,7 @@ pub mod ggsw_ciphertext_list;
 pub mod glwe_ciphertext;
 pub mod glwe_ciphertext_list;
 pub mod glwe_secret_key;
+pub mod lwe_bootstrap_key;
 pub mod lwe_ciphertext;
 pub mod lwe_ciphertext_list;
 pub mod lwe_keyswitch_key;
@@ -19,6 +20,7 @@ pub use ggsw_ciphertext_list::*;
 pub use glwe_ciphertext::*;
 pub use glwe_ciphertext_list::*;
 pub use glwe_secret_key::*;
+pub use lwe_bootstrap_key::*;
 pub use lwe_ciphertext::*;
 pub use lwe_ciphertext_list::*;
 pub use lwe_keyswitch_key::*;