chore(tests): add support for AP in tests and benches

Author: nsarlin-zama

tests/backward_compatibility/high_level_api.rs

@@ -2,7 +2,7 @@ use super::shortint::load_params;
 use crate::{load_and_unversionize, TestedModule};
 use std::path::Path;
 use tfhe::prelude::{CiphertextList, FheDecrypt, FheEncrypt};
-use tfhe::shortint::PBSParameters;
+use tfhe::shortint::{AtomicPatternParameters, PBSParameters};
 #[cfg(feature = "zk-pok")]
 use tfhe::zk::CompactPkeCrs;
 use tfhe::{
@@ -22,10 +22,10 @@ use tfhe_backward_compat_data::{
 };
 use tfhe_versionable::Unversionize;
 
-fn load_hl_params(test_params: &TestParameterSet) -> PBSParameters {
+fn load_hl_params(test_params: &TestParameterSet) -> AtomicPatternParameters {
     let pbs_params = load_params(test_params);
 
-    PBSParameters::PBS(pbs_params)
+    PBSParameters::PBS(pbs_params).into()
 }
 
 /// Test HL ciphertext: loads the ciphertext and compare the decrypted value to the one in the

tfhe/benches/utilities.rs

@@ -38,6 +38,7 @@ pub mod shortint_utils {
     use super::*;
     use itertools::iproduct;
     use std::vec::IntoIter;
+    use tfhe::shortint::atomic_pattern::AtomicPatternParameters;
     use tfhe::shortint::parameters::compact_public_key_only::CompactPublicKeyEncryptionParameters;
     #[cfg(not(feature = "gpu"))]
     use tfhe::shortint::parameters::current_params::V1_1_PARAM_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M128;
@@ -56,8 +57,10 @@ pub mod shortint_utils {
     /// of parameters and a num_block to achieve a certain bit_size ciphertext
     /// in radix decomposition
     pub struct ParamsAndNumBlocksIter {
-        params_and_bit_sizes:
-            itertools::Product<IntoIter<tfhe::shortint::PBSParameters>, IntoIter<usize>>,
+        params_and_bit_sizes: itertools::Product<
+            IntoIter<tfhe::shortint::atomic_pattern::AtomicPatternParameters>,
+            IntoIter<usize>,
+        >,
     }
 
     impl Default for ParamsAndNumBlocksIter {
@@ -90,7 +93,11 @@ pub mod shortint_utils {
     }
 
     impl Iterator for ParamsAndNumBlocksIter {
-        type Item = (tfhe::shortint::PBSParameters, usize, usize);
+        type Item = (
+            tfhe::shortint::atomic_pattern::AtomicPatternParameters,
+            usize,
+            usize,
+        );
 
         fn next(&mut self) -> Option<Self::Item> {
             let (param, bit_size) = self.params_and_bit_sizes.next()?;
@@ -126,6 +133,31 @@ pub mod shortint_utils {
         }
     }
 
+    impl From<AtomicPatternParameters> for CryptoParametersRecord<u64> {
+        fn from(params: AtomicPatternParameters) -> Self {
+            CryptoParametersRecord {
+                lwe_dimension: Some(params.lwe_dimension()),
+                glwe_dimension: Some(params.glwe_dimension()),
+                polynomial_size: Some(params.polynomial_size()),
+                lwe_noise_distribution: Some(params.lwe_noise_distribution()),
+                glwe_noise_distribution: Some(params.glwe_noise_distribution()),
+                pbs_base_log: Some(params.pbs_base_log()),
+                pbs_level: Some(params.pbs_level()),
+                ks_base_log: Some(params.ks_base_log()),
+                ks_level: Some(params.ks_level()),
+                message_modulus: Some(params.message_modulus().0),
+                carry_modulus: Some(params.carry_modulus().0),
+                ciphertext_modulus: Some(
+                    params
+                        .ciphertext_modulus()
+                        .try_to()
+                        .expect("failed to convert ciphertext modulus"),
+                ),
+                ..Default::default()
+            }
+        }
+    }
+
     impl From<ShortintKeySwitchingParameters> for CryptoParametersRecord<u64> {
         fn from(params: ShortintKeySwitchingParameters) -> Self {
             CryptoParametersRecord {

tfhe/src/core_crypto/gpu/entities/lwe_bootstrap_key.rs

@@ -65,7 +65,7 @@ pub struct CudaLweBootstrapKey {
 impl CudaLweBootstrapKey {
     pub fn from_lwe_bootstrap_key<InputBskCont: Container>(
         bsk: &LweBootstrapKey<InputBskCont>,
-        ms_noise_reduction_key: Option<ModulusSwitchNoiseReductionKey>,
+        ms_noise_reduction_key: Option<ModulusSwitchNoiseReductionKey<u64>>,
         streams: &CudaStreams,
     ) -> Self
     where

tfhe/src/high_level_api/keys/client.rs

@@ -70,7 +70,7 @@ impl ClientKey {
         }
     }
 
-    pub fn computation_parameters(&self) -> crate::shortint::PBSParameters {
+    pub fn computation_parameters(&self) -> crate::shortint::AtomicPatternParameters {
         self.key.block_parameters()
     }
 

tfhe/src/high_level_api/keys/inner.rs

@@ -174,7 +174,7 @@ impl IntegerClientKey {
         }
     }
 
-    pub(crate) fn block_parameters(&self) -> crate::shortint::parameters::PBSParameters {
+    pub(crate) fn block_parameters(&self) -> crate::shortint::parameters::AtomicPatternParameters {
         self.key.parameters()
     }
 }

tfhe/src/integer/client_key/crt.rs

@@ -76,7 +76,7 @@ impl CrtClientKey {
     }
 
     /// Returns the parameters used by the client key.
-    pub fn parameters(&self) -> crate::shortint::PBSParameters {
+    pub fn parameters(&self) -> crate::shortint::AtomicPatternParameters {
         self.key.parameters()
     }
 

tfhe/src/integer/client_key/mod.rs

@@ -142,8 +142,8 @@ impl ClientKey {
         Self { key }
     }
 
-    pub fn parameters(&self) -> crate::shortint::PBSParameters {
-        self.key.parameters.pbs_parameters().unwrap()
+    pub fn parameters(&self) -> crate::shortint::AtomicPatternParameters {
+        self.key.parameters.ap_parameters().unwrap()
     }
 
     #[cfg(test)]

tfhe/src/integer/client_key/radix.rs

@@ -11,7 +11,9 @@ use crate::integer::compression_keys::{
     CompressedCompressionKey, CompressedDecompressionKey, CompressionPrivateKeys,
 };
 use crate::integer::BooleanBlock;
-use crate::shortint::{Ciphertext as ShortintCiphertext, PBSParameters as ShortintParameters};
+use crate::shortint::{
+    AtomicPatternParameters as ShortintParameters, Ciphertext as ShortintCiphertext,
+};
 use serde::{Deserialize, Serialize};
 use tfhe_versionable::Versionize;
 

tfhe/src/integer/gpu/server_key/mod.rs

@@ -13,7 +13,9 @@ use crate::integer::ClientKey;
 use crate::shortint::ciphertext::{MaxDegree, MaxNoiseLevel};
 use crate::shortint::engine::ShortintEngine;
 use crate::shortint::server_key::ModulusSwitchNoiseReductionKey;
-use crate::shortint::{CarryModulus, CiphertextModulus, MessageModulus, PBSOrder};
+use crate::shortint::{
+    AtomicPatternParameters, CarryModulus, CiphertextModulus, MessageModulus, PBSOrder,
+};
 mod radix;
 
 pub enum CudaBootstrappingKey {
@@ -83,7 +85,10 @@ impl CudaServerKey {
         let mut engine = ShortintEngine::new();
 
         // Generate a regular keyset and convert to the GPU
-        let pbs_params_base = &cks.parameters();
+        let AtomicPatternParameters::Classical(pbs_params_base) = &cks.parameters() else {
+            panic!("Only the classical atomic pattern is supported on GPU")
+        };
+
         let d_bootstrapping_key = match pbs_params_base {
             crate::shortint::PBSParameters::PBS(pbs_params) => {
                 let h_bootstrap_key: LweBootstrapKeyOwned<u64> =

tfhe/src/integer/gpu/server_key/radix/tests_long_run/test_erc20.rs

@@ -18,7 +18,7 @@ create_gpu_parameterized_test!(no_cmux_erc20 {
 
 fn safe_erc20<P>(param: P)
 where
-    P: Into<PBSParameters>,
+    P: Into<AtomicPatternParameters>,
 {
     let overflowing_add_executor =
         GpuMultiDeviceFunctionExecutor::new(&CudaServerKey::unsigned_overflowing_add);
@@ -37,7 +37,7 @@ where
 
 fn whitepaper_erc20<P>(param: P)
 where
-    P: Into<PBSParameters>,
+    P: Into<AtomicPatternParameters>,
 {
     let ge_executor = GpuMultiDeviceFunctionExecutor::new(&CudaServerKey::ge);
     let add_executor = GpuMultiDeviceFunctionExecutor::new(&CudaServerKey::add);
@@ -54,7 +54,7 @@ where
 
 fn no_cmux_erc20<P>(param: P)
 where
-    P: Into<PBSParameters>,
+    P: Into<AtomicPatternParameters>,
 {
     let ge_executor = GpuMultiDeviceFunctionExecutor::new(&CudaServerKey::ge);
     let mul_executor = GpuMultiDeviceFunctionExecutor::new(&CudaServerKey::mul);