feat(core): chunked lwe_bsk generation

Author: youben11

tfhe/src/core_crypto/algorithms/lwe_bootstrap_key_generation.rs

@@ -6,8 +6,8 @@ use crate::core_crypto::commons::math::random::{
 };
 use crate::core_crypto::commons::math::torus::UnsignedTorus;
 use crate::core_crypto::commons::parameters::{
-    CiphertextModulus, DecompositionBaseLog, DecompositionLevelCount, GlweDimension, LweDimension,
-    PolynomialSize,
+    ChunkSize, CiphertextModulus, DecompositionBaseLog, DecompositionLevelCount, GlweDimension,
+    LweDimension, PolynomialSize,
 };
 use crate::core_crypto::commons::test_tools::new_secret_random_generator;
 use crate::core_crypto::entities::*;
@@ -17,7 +17,7 @@ const NB_TESTS: usize = 10;
 #[cfg(tarpaulin)]
 const NB_TESTS: usize = 1;
 
-fn test_parallel_and_seeded_bsk_gen_equivalence<T: UnsignedTorus + Sync + Send>(
+fn test_parallel_and_seeded_and_chunked_bsk_gen_equivalence<T: UnsignedTorus + Sync + Send>(
     ciphertext_modulus: CiphertextModulus<T>,
 ) {
     for _ in 0..NB_TESTS {
@@ -144,29 +144,75 @@ fn test_parallel_and_seeded_bsk_gen_equivalence<T: UnsignedTorus + Sync + Send>(
         let par_decompressed_bsk = parallel_seeded_bsk.par_decompress_into_lwe_bootstrap_key();
 
         assert_eq!(ser_decompressed_bsk, par_decompressed_bsk);
+
+        let encryption_generator = EncryptionRandomGenerator::<DefaultRandomGenerator>::new(
+            mask_seed,
+            &mut DeterministicSeeder::<DefaultRandomGenerator>::new(deterministic_seeder_seed),
+        );
+
+        let chunk_generator = LweBootstrapKeyChunkGenerator::new(
+            encryption_generator,
+            ChunkSize(crate::core_crypto::commons::test_tools::random_usize_between(1..5)),
+            lwe_dim,
+            glwe_dim.to_glwe_size(),
+            poly_size,
+            base_log,
+            level,
+            ciphertext_modulus,
+            lwe_sk.clone(),
+            glwe_sk.clone(),
+            noise_distribution,
+            false,
+        );
+        let chunks = chunk_generator.collect::<Vec<_>>();
+        let assembled_bsk = allocate_and_assemble_lwe_bootstrap_key_from_chunks(chunks);
+        assert_eq!(assembled_bsk, sequential_bsk);
+
+        let encryption_generator = EncryptionRandomGenerator::<DefaultRandomGenerator>::new(
+            mask_seed,
+            &mut DeterministicSeeder::<DefaultRandomGenerator>::new(deterministic_seeder_seed),
+        );
+
+        let par_chunk_generator = LweBootstrapKeyChunkGenerator::new(
+            encryption_generator,
+            ChunkSize(crate::core_crypto::commons::test_tools::random_usize_between(1..5)),
+            lwe_dim,
+            glwe_dim.to_glwe_size(),
+            poly_size,
+            base_log,
+            level,
+            ciphertext_modulus,
+            lwe_sk,
+            glwe_sk,
+            noise_distribution,
+            true,
+        );
+        let chunks = par_chunk_generator.collect::<Vec<_>>();
+        let assembled_bsk = allocate_and_assemble_lwe_bootstrap_key_from_chunks(chunks);
+        assert_eq!(assembled_bsk, sequential_bsk);
     }
 }
 
 #[test]
-fn test_parallel_and_seeded_bsk_gen_equivalence_u32_native_mod() {
-    test_parallel_and_seeded_bsk_gen_equivalence::<u32>(CiphertextModulus::new_native());
+fn test_parallel_and_seeded_and_chunked_bsk_gen_equivalence_u32_native_mod() {
+    test_parallel_and_seeded_and_chunked_bsk_gen_equivalence::<u32>(CiphertextModulus::new_native());
 }
 
 #[test]
-fn test_parallel_and_seeded_bsk_gen_equivalence_u32_custom_mod() {
-    test_parallel_and_seeded_bsk_gen_equivalence::<u32>(
+fn test_parallel_and_seeded_and_chunked_bsk_gen_equivalence_u32_custom_mod() {
+    test_parallel_and_seeded_and_chunked_bsk_gen_equivalence::<u32>(
         CiphertextModulus::try_new_power_of_2(31).unwrap(),
     );
 }
 
 #[test]
-fn test_parallel_and_seeded_bsk_gen_equivalence_u64_native_mod() {
-    test_parallel_and_seeded_bsk_gen_equivalence::<u64>(CiphertextModulus::new_native());
+fn test_parallel_and_seeded_and_chunked_bsk_gen_equivalence_u64_native_mod() {
+    test_parallel_and_seeded_and_chunked_bsk_gen_equivalence::<u64>(CiphertextModulus::new_native());
 }
 
 #[test]
-fn test_parallel_and_seeded_bsk_gen_equivalence_u64_custom_mod() {
-    test_parallel_and_seeded_bsk_gen_equivalence::<u64>(
+fn test_parallel_and_seeded_and_chunked_bsk_gen_equivalence_u64_custom_mod() {
+    test_parallel_and_seeded_and_chunked_bsk_gen_equivalence::<u64>(
         CiphertextModulus::try_new_power_of_2(63).unwrap(),
     );
 }

tfhe/src/core_crypto/algorithms/test/lwe_bootstrap_key_generation.rs

@@ -182,3 +182,8 @@ pub enum RSigmaFactorVersions {
 pub enum NoiseEstimationMeasureBoundVersions {
     V0(NoiseEstimationMeasureBound),
 }
+
+#[derive(VersionsDispatch)]
+pub enum ChunkSizeVersions {
+    V0(ChunkSize),
+}

tfhe/src/core_crypto/backward_compatibility/commons/parameters.rs

@@ -0,0 +1,21 @@
+use tfhe_versionable::deprecation::Deprecable;
+use tfhe_versionable::VersionsDispatch;
+
+use crate::core_crypto::prelude::{Container, LweBootstrapKeyChunk, UnsignedInteger};
+
+impl<C: Container> Deprecable for LweBootstrapKeyChunk<C>
+where
+    C::Element: UnsignedInteger,
+{
+    const TYPE_NAME: &'static str = "LweBootstrapKeyChunk";
+    // TODO
+    const MIN_SUPPORTED_APP_VERSION: &'static str = "TFHE-rs v0.10";
+}
+
+#[derive(VersionsDispatch)]
+pub enum LweBootstrapKeyChunkVersions<C: Container>
+where
+    C::Element: UnsignedInteger,
+{
+    V0(LweBootstrapKeyChunk<C>),
+}

tfhe/src/core_crypto/backward_compatibility/entities/lwe_bootstrap_key_chunk.rs

@@ -9,6 +9,7 @@ pub mod glwe_ciphertext_list;
 pub mod glwe_secret_key;
 pub mod gsw_ciphertext;
 pub mod lwe_bootstrap_key;
+pub mod lwe_bootstrap_key_chunk;
 pub mod lwe_ciphertext;
 pub mod lwe_ciphertext_list;
 pub mod lwe_compact_ciphertext_list;

tfhe/src/core_crypto/backward_compatibility/entities/mod.rs

@@ -388,3 +388,8 @@ pub struct RSigmaFactor(pub f64);
 #[derive(Copy, Clone, Debug, PartialEq, Serialize, Deserialize, Versionize)]
 #[versionize(NoiseEstimationMeasureBoundVersions)]
 pub struct NoiseEstimationMeasureBound(pub f64);
+
+/// The size of a chunk in a chunked key.
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Versionize)]
+#[versionize(ChunkSizeVersions)]
+pub struct ChunkSize(pub usize);