feat(integer): evaluate unpacking luts during casting to improve perf

- allows to avoid some LUT evaluations during expansions of the various CompactCiphertextList
zama-ai · Sep 26, 2024 · 2a7475d · 2a7475d
1 parent 6c804d8
commit 2a7475d
Show file tree

Hide file tree

Showing 14 changed files with 732 additions and 310 deletions.
diff --git a/tfhe/benches/integer/zk_pke.rs b/tfhe/benches/integer/zk_pke.rs
@@ -8,9 +8,7 @@ use std::io::Write;
 use std::path::Path;
 use tfhe::core_crypto::prelude::*;
 use tfhe::integer::key_switching_key::KeySwitchingKey;
-use tfhe::integer::parameters::{
-    IntegerCompactCiphertextListCastingMode, IntegerCompactCiphertextListUnpackingMode,
-};
+use tfhe::integer::parameters::IntegerCompactCiphertextListExpansionMode;
 use tfhe::integer::{ClientKey, CompactPrivateKey, CompactPublicKey, ServerKey};
 use tfhe::keycache::NamedParam;
 use tfhe::shortint::parameters::classic::tuniform::p_fail_2_minus_64::ks_pbs::PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
@@ -251,8 +249,7 @@ fn pke_zk_verify(c: &mut Criterion, results_file: &Path) {
                                 public_params,
                                 &pk,
                                 &metadata,
-                                IntegerCompactCiphertextListUnpackingMode::UnpackIfNecessary(&sks),
-                                IntegerCompactCiphertextListCastingMode::CastIfNecessary(
+                                IntegerCompactCiphertextListExpansionMode::CastAndUnpackIfNecessary(
                                     casting_key.as_view(),
                                 ),
                             )

diff --git a/tfhe/docs/guides/data_versioning.md b/tfhe/docs/guides/data_versioning.md
@@ -90,8 +90,8 @@ You will find below a list of breaking changes and how to upgrade them.
 ```rust
 use std::io::Cursor;
 use tfhe::integer::ciphertext::{
-    CompactCiphertextList, DataKind, IntegerCompactCiphertextListCastingMode,
-    IntegerCompactCiphertextListUnpackingMode, SignedRadixCiphertext,
+    CompactCiphertextList, DataKind, IntegerCompactCiphertextListExpansionMode,
+    SignedRadixCiphertext,
 };
 use tfhe::integer::{ClientKey, CompactPublicKey};
 use tfhe::shortint::parameters::classic::compact_pk::PARAM_MESSAGE_2_CARRY_2_COMPACT_PK_KS_PBS;
@@ -130,10 +130,7 @@ pub fn main() {
         .reinterpret_data(&[DataKind::Signed(num_blocks)])
         .unwrap();
     let expander = compact_ct
-        .expand(
-            IntegerCompactCiphertextListUnpackingMode::NoUnpacking,
-            IntegerCompactCiphertextListCastingMode::NoCasting,
-        )
+        .expand(IntegerCompactCiphertextListExpansionMode::NoCastingAndNoUnpacking)
         .unwrap();
     let expanded = expander.get::<SignedRadixCiphertext>(0).unwrap().unwrap();
     let decrypted: i8 = client_key.decrypt_signed_radix(&expanded);

diff --git a/tfhe/src/high_level_api/compact_list.rs b/tfhe/src/high_level_api/compact_list.rs
@@ -10,8 +10,7 @@ use crate::high_level_api::traits::Tagged;
 use crate::integer::ciphertext::{Compactable, DataKind, Expandable};
 use crate::integer::encryption::KnowsMessageModulus;
 use crate::integer::parameters::{
-    CompactCiphertextListConformanceParams, IntegerCompactCiphertextListCastingMode,
-    IntegerCompactCiphertextListUnpackingMode,
+    CompactCiphertextListConformanceParams, IntegerCompactCiphertextListExpansionMode,
 };
 use crate::named::Named;
 use crate::shortint::MessageModulus;
@@ -108,13 +107,7 @@ impl CompactCiphertextList {
         sks: &crate::ServerKey,
     ) -> crate::Result<CompactCiphertextListExpander> {
         self.inner
-            .expand(
-                IntegerCompactCiphertextListUnpackingMode::UnpackIfNecessary(sks.key.pbs_key()),
-                sks.cpk_casting_key().map_or(
-                    IntegerCompactCiphertextListCastingMode::NoCasting,
-                    IntegerCompactCiphertextListCastingMode::CastIfNecessary,
-                ),
-            )
+            .expand(sks.integer_compact_ciphertext_list_expansion_mode())
             .map(|inner| CompactCiphertextListExpander {
                 inner,
                 tag: self.tag.clone(),
@@ -126,44 +119,22 @@ impl CompactCiphertextList {
         if !self.inner.is_packed() && !self.inner.needs_casting() {
             // No ServerKey required, short-circuit to avoid the global state call
             return Ok(CompactCiphertextListExpander {
-                inner: self.inner.expand(
-                    IntegerCompactCiphertextListUnpackingMode::NoUnpacking,
-                    IntegerCompactCiphertextListCastingMode::NoCasting,
-                )?,
+                inner: self
+                    .inner
+                    .expand(IntegerCompactCiphertextListExpansionMode::NoCastingAndNoUnpacking)?,
                 tag: self.tag.clone(),
             });
         }
 
         global_state::try_with_internal_keys(|maybe_keys| match maybe_keys {
             None => Err(crate::high_level_api::errors::UninitializedServerKey.into()),
-            Some(InternalServerKey::Cpu(cpu_key)) => {
-                let unpacking_mode = if self.inner.is_packed() {
-                    IntegerCompactCiphertextListUnpackingMode::UnpackIfNecessary(cpu_key.pbs_key())
-                } else {
-                    IntegerCompactCiphertextListUnpackingMode::NoUnpacking
-                };
-
-                let casting_mode = if self.inner.needs_casting() {
-                    IntegerCompactCiphertextListCastingMode::CastIfNecessary(
-                        cpu_key.cpk_casting_key().ok_or_else(|| {
-                            crate::Error::new(
-                                "No casting key found in ServerKey, \
-                                required to expand this CompactCiphertextList"
-                                    .to_string(),
-                            )
-                        })?,
-                    )
-                } else {
-                    IntegerCompactCiphertextListCastingMode::NoCasting
-                };
-
-                self.inner
-                    .expand(unpacking_mode, casting_mode)
-                    .map(|inner| CompactCiphertextListExpander {
-                        inner,
-                        tag: self.tag.clone(),
-                    })
-            }
+            Some(InternalServerKey::Cpu(cpu_key)) => self
+                .inner
+                .expand(cpu_key.integer_compact_ciphertext_list_expansion_mode())
+                .map(|inner| CompactCiphertextListExpander {
+                    inner,
+                    tag: self.tag.clone(),
+                }),
             #[cfg(feature = "gpu")]
             Some(_) => Err(crate::Error::new("Expected a CPU server key".to_string())),
         })
@@ -255,51 +226,26 @@ mod zk {
                         public_params,
                         &pk.key.key,
                         metadata,
-                        IntegerCompactCiphertextListUnpackingMode::NoUnpacking,
-                        IntegerCompactCiphertextListCastingMode::NoCasting,
+                        IntegerCompactCiphertextListExpansionMode::NoCastingAndNoUnpacking,
                     )?,
                     tag: self.tag.clone(),
                 });
             }
 
             global_state::try_with_internal_keys(|maybe_keys| match maybe_keys {
                 None => Err(crate::high_level_api::errors::UninitializedServerKey.into()),
-                Some(InternalServerKey::Cpu(cpu_key)) => {
-                    let unpacking_mode = if self.inner.is_packed() {
-                        IntegerCompactCiphertextListUnpackingMode::UnpackIfNecessary(
-                            cpu_key.pbs_key(),
-                        )
-                    } else {
-                        IntegerCompactCiphertextListUnpackingMode::NoUnpacking
-                    };
-
-                    let casting_mode = if self.inner.needs_casting() {
-                        IntegerCompactCiphertextListCastingMode::CastIfNecessary(
-                            cpu_key.cpk_casting_key().ok_or_else(|| {
-                                crate::Error::new(
-                                    "No casting key found in ServerKey, \
-                                required to expand this CompactCiphertextList"
-                                        .to_string(),
-                                )
-                            })?,
-                        )
-                    } else {
-                        IntegerCompactCiphertextListCastingMode::NoCasting
-                    };
-
-                    self.inner
-                        .verify_and_expand(
-                            public_params,
-                            &pk.key.key,
-                            metadata,
-                            unpacking_mode,
-                            casting_mode,
-                        )
-                        .map(|expander| CompactCiphertextListExpander {
-                            inner: expander,
-                            tag: self.tag.clone(),
-                        })
-                }
+                Some(InternalServerKey::Cpu(cpu_key)) => self
+                    .inner
+                    .verify_and_expand(
+                        public_params,
+                        &pk.key.key,
+                        metadata,
+                        cpu_key.integer_compact_ciphertext_list_expansion_mode(),
+                    )
+                    .map(|expander| CompactCiphertextListExpander {
+                        inner: expander,
+                        tag: self.tag.clone(),
+                    }),
                 #[cfg(feature = "gpu")]
                 Some(_) => Err(crate::Error::new("Expected a CPU server key".to_string())),
             })
@@ -315,45 +261,23 @@ mod zk {
                 // No ServerKey required, short circuit to avoid the global state call
                 return Ok(CompactCiphertextListExpander {
                     inner: self.inner.expand_without_verification(
-                        IntegerCompactCiphertextListUnpackingMode::NoUnpacking,
-                        IntegerCompactCiphertextListCastingMode::NoCasting,
+                        IntegerCompactCiphertextListExpansionMode::NoCastingAndNoUnpacking,
                     )?,
                     tag: self.tag.clone(),
                 });
             }
 
             global_state::try_with_internal_keys(|maybe_keys| match maybe_keys {
                 None => Err(crate::high_level_api::errors::UninitializedServerKey.into()),
-                Some(InternalServerKey::Cpu(cpu_key)) => {
-                    let unpacking_mode = if self.inner.is_packed() {
-                        IntegerCompactCiphertextListUnpackingMode::UnpackIfNecessary(
-                            cpu_key.pbs_key(),
-                        )
-                    } else {
-                        IntegerCompactCiphertextListUnpackingMode::NoUnpacking
-                    };
-
-                    let casting_mode = if self.inner.needs_casting() {
-                        IntegerCompactCiphertextListCastingMode::CastIfNecessary(
-                            cpu_key.cpk_casting_key().ok_or_else(|| {
-                                crate::Error::new(
-                                    "No casting key found in ServerKey, \
-                                required to expand this CompactCiphertextList"
-                                        .to_string(),
-                                )
-                            })?,
-                        )
-                    } else {
-                        IntegerCompactCiphertextListCastingMode::NoCasting
-                    };
-
-                    self.inner
-                        .expand_without_verification(unpacking_mode, casting_mode)
-                        .map(|expander| CompactCiphertextListExpander {
-                            inner: expander,
-                            tag: self.tag.clone(),
-                        })
-                }
+                Some(InternalServerKey::Cpu(cpu_key)) => self
+                    .inner
+                    .expand_without_verification(
+                        cpu_key.integer_compact_ciphertext_list_expansion_mode(),
+                    )
+                    .map(|expander| CompactCiphertextListExpander {
+                        inner: expander,
+                        tag: self.tag.clone(),
+                    }),
                 #[cfg(feature = "gpu")]
                 Some(_) => Err(crate::Error::new("Expected a CPU server key".to_string())),
             })

diff --git a/tfhe/src/high_level_api/keys/server.rs b/tfhe/src/high_level_api/keys/server.rs
@@ -8,6 +8,7 @@ use crate::high_level_api::keys::{IntegerCompressedServerKey, IntegerServerKey};
 use crate::integer::compression_keys::{
     CompressedCompressionKey, CompressedDecompressionKey, CompressionKey, DecompressionKey,
 };
+use crate::integer::parameters::IntegerCompactCiphertextListExpansionMode;
 use crate::named::Named;
 use crate::prelude::Tagged;
 use crate::shortint::MessageModulus;
@@ -98,6 +99,19 @@ impl ServerKey {
     pub(in crate::high_level_api) fn message_modulus(&self) -> MessageModulus {
         self.key.message_modulus()
     }
+
+    pub(in crate::high_level_api) fn integer_compact_ciphertext_list_expansion_mode(
+        &self,
+    ) -> IntegerCompactCiphertextListExpansionMode {
+        self.cpk_casting_key().map_or_else(
+            || {
+                IntegerCompactCiphertextListExpansionMode::UnpackAndSanitizeIfNecessary(
+                    self.pbs_key(),
+                )
+            },
+            IntegerCompactCiphertextListExpansionMode::CastAndUnpackIfNecessary,
+        )
+    }
 }
 
 impl Tagged for ServerKey {