refactor(shortint): use view -> decompress to simplify key decompression

- the decompress_into primitives consume the input entity, but we can use a view that won't consume the original owned key
zama-ai · Dec 2, 2024 · 8dd419f · 8dd419f
1 parent a0ad0c7
commit 8dd419f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 63 deletions.
diff --git a/tfhe/src/shortint/key_switching_key/mod.rs b/tfhe/src/shortint/key_switching_key/mod.rs
@@ -4,8 +4,8 @@
 
 use crate::conformance::ParameterSetConformant;
 use crate::core_crypto::prelude::{
-    decompress_seeded_lwe_keyswitch_key, keyswitch_lwe_ciphertext, ActivatedRandomGenerator,
-    KeyswitchKeyConformanceParams, LweKeyswitchKeyOwned, SeededLweKeyswitchKeyOwned,
+    keyswitch_lwe_ciphertext, KeyswitchKeyConformanceParams, LweKeyswitchKeyOwned,
+    SeededLweKeyswitchKeyOwned,
 };
 use crate::shortint::ciphertext::Degree;
 use crate::shortint::client_key::secret_encryption_key::SecretEncryptionKeyView;
@@ -783,22 +783,10 @@ pub struct CompressedKeySwitchingKeyMaterial {
 
 impl CompressedKeySwitchingKeyMaterial {
     pub fn decompress(&self) -> KeySwitchingKeyMaterial {
-        let key_switching_key = {
-            let cksk = &self.key_switching_key;
-            let mut decompressed_ksk = LweKeyswitchKeyOwned::new(
-                0u64,
-                cksk.decomposition_base_log(),
-                cksk.decomposition_level_count(),
-                cksk.input_key_lwe_dimension(),
-                cksk.output_key_lwe_dimension(),
-                cksk.ciphertext_modulus(),
-            );
-            decompress_seeded_lwe_keyswitch_key::<_, _, _, ActivatedRandomGenerator>(
-                &mut decompressed_ksk,
-                cksk,
-            );
-            decompressed_ksk
-        };
+        let key_switching_key = self
+            .key_switching_key
+            .as_view()
+            .par_decompress_into_lwe_keyswitch_key();
 
         KeySwitchingKeyMaterial {
             key_switching_key,

diff --git a/tfhe/src/shortint/server_key/compressed.rs b/tfhe/src/shortint/server_key/compressed.rs
@@ -162,35 +162,15 @@ impl CompressedServerKey {
 
         let (key_switching_key, bootstrapping_key) = rayon::join(
             || {
-                let mut decompressed_key_switching_key = LweKeyswitchKey::new(
-                    0,
-                    compressed_key_switching_key.decomposition_base_log(),
-                    compressed_key_switching_key.decomposition_level_count(),
-                    compressed_key_switching_key.input_key_lwe_dimension(),
-                    compressed_key_switching_key.output_key_lwe_dimension(),
-                    compressed_key_switching_key.ciphertext_modulus(),
-                );
-                par_decompress_seeded_lwe_keyswitch_key::<_, _, _, ActivatedRandomGenerator>(
-                    &mut decompressed_key_switching_key,
-                    compressed_key_switching_key,
-                );
-                decompressed_key_switching_key
+                compressed_key_switching_key
+                    .as_view()
+                    .par_decompress_into_lwe_keyswitch_key()
             },
             || match compressed_bootstrapping_key {
                 ShortintCompressedBootstrappingKey::Classic(compressed_bootstrapping_key) => {
-                    let mut decompressed_bootstrapping_key = LweBootstrapKey::new(
-                        0,
-                        compressed_bootstrapping_key.glwe_size(),
-                        compressed_bootstrapping_key.polynomial_size(),
-                        compressed_bootstrapping_key.decomposition_base_log(),
-                        compressed_bootstrapping_key.decomposition_level_count(),
-                        compressed_bootstrapping_key.input_lwe_dimension(),
-                        compressed_bootstrapping_key.ciphertext_modulus(),
-                    );
-                    par_decompress_seeded_lwe_bootstrap_key::<_, _, _, ActivatedRandomGenerator>(
-                        &mut decompressed_bootstrapping_key,
-                        compressed_bootstrapping_key,
-                    );
+                    let decompressed_bootstrapping_key = compressed_bootstrapping_key
+                        .as_view()
+                        .par_decompress_into_lwe_bootstrap_key();
 
                     let mut fourier_bsk = FourierLweBootstrapKeyOwned::new(
                         decompressed_bootstrapping_key.input_lwe_dimension(),
@@ -211,25 +191,9 @@ impl CompressedServerKey {
                     seeded_bsk: compressed_bootstrapping_key,
                     deterministic_execution,
                 } => {
-                    let mut decompressed_bootstrapping_key = LweMultiBitBootstrapKeyOwned::new(
-                        0,
-                        compressed_bootstrapping_key.glwe_size(),
-                        compressed_bootstrapping_key.polynomial_size(),
-                        compressed_bootstrapping_key.decomposition_base_log(),
-                        compressed_bootstrapping_key.decomposition_level_count(),
-                        compressed_bootstrapping_key.input_lwe_dimension(),
-                        compressed_bootstrapping_key.grouping_factor(),
-                        compressed_bootstrapping_key.ciphertext_modulus(),
-                    );
-                    par_decompress_seeded_lwe_multi_bit_bootstrap_key::<
-                        _,
-                        _,
-                        _,
-                        ActivatedRandomGenerator,
-                    >(
-                        &mut decompressed_bootstrapping_key,
-                        compressed_bootstrapping_key,
-                    );
+                    let decompressed_bootstrapping_key = compressed_bootstrapping_key
+                        .as_view()
+                        .par_decompress_into_lwe_multi_bit_bootstrap_key();
 
                     let mut fourier_bsk = FourierLweMultiBitBootstrapKeyOwned::new(
                         decompressed_bootstrapping_key.input_lwe_dimension(),