feat(gpu): create noise and pfail tests for pbs + ks + ms

guillermo-oyarzun · guillermo-oyarzun · commit 4d34ee47d91c · 2025-12-02T10:09:28.000+01:00
diff --git a/backends/tfhe-cuda-backend/cuda/include/ciphertext.h b/backends/tfhe-cuda-backend/cuda/include/ciphertext.h
@@ -35,17 +35,9 @@ void cuda_centered_modulus_switch_64(void *stream, uint32_t gpu_index,
                                      uint32_t lwe_dimension,
                                      uint32_t log_modulus);
 
-void cuda_improve_noise_modulus_switch_64(
-    void *stream, uint32_t gpu_index, void *lwe_array_out,
-    void const *lwe_array_in, void const *lwe_array_indexes,
-    void const *encrypted_zeros, uint32_t lwe_size, uint32_t num_lwes,
-    uint32_t num_zeros, double input_variance, double r_sigma, double bound,
-    uint32_t log_modulus);
-
 void cuda_glwe_sample_extract_128(
     void *stream, uint32_t gpu_index, void *lwe_array_out,
     void const *glwe_array_in, uint32_t const *nth_array, uint32_t num_nths,
     uint32_t lwe_per_glwe, uint32_t glwe_dimension, uint32_t polynomial_size);
 }
-
 #endif
diff --git a/backends/tfhe-cuda-backend/cuda/include/integer/integer.h b/backends/tfhe-cuda-backend/cuda/include/integer/integer.h
@@ -921,6 +921,10 @@ void cuda_unchecked_first_index_in_clears_64(
     uint32_t num_unique, uint32_t num_blocks, uint32_t num_blocks_index,
     int8_t *mem, void *const *bsks, void *const *ksks);
 
+void cuda_small_scalar_multiplication_integer_64_inplace(
+    CudaStreamsFFI streams, CudaRadixCiphertextFFI *lwe_array, uint64_t scalar,
+    const uint32_t message_modulus, const uint32_t carry_modulus);
+
 void cleanup_cuda_unchecked_first_index_in_clears_64(CudaStreamsFFI streams,
                                                      int8_t **mem_ptr_void);
 
diff --git a/backends/tfhe-cuda-backend/cuda/include/integer/integer_utilities.h b/backends/tfhe-cuda-backend/cuda/include/integer/integer_utilities.h
@@ -43,6 +43,8 @@ class NoiseLevel {
               "parameters");                                                   \
     } else if ((msg_mod) == 0 && (carry_mod) == 0) {                           \
       break;                                                                   \
+    } else if ((msg_mod) == 4 && (carry_mod) == 32) {                          \
+      constexpr int max_noise_level = 9;                                       \
     } else {                                                                   \
       PANIC("Invalid message modulus or carry modulus")                        \
     }                                                                          \
diff --git a/backends/tfhe-cuda-backend/cuda/src/crypto/torus.cuh b/backends/tfhe-cuda-backend/cuda/src/crypto/torus.cuh
@@ -144,6 +144,20 @@ __device__ __forceinline__ T modulus_switch(T input, uint32_t log_modulus) {
   return output;
 }
 
+template <typename Torus, class params>
+__device__ uint32_t calculates_monomial_degree(const Torus *lwe_array_group,
+                                               uint32_t ggsw_idx,
+                                               uint32_t grouping_factor) {
+  Torus x = 0;
+  for (int i = 0; i < grouping_factor; i++) {
+    uint32_t mask_position = grouping_factor - (i + 1);
+    int selection_bit = (ggsw_idx >> mask_position) & 1;
+    x += selection_bit * lwe_array_group[i];
+  }
+
+  return modulus_switch(x, params::log2_degree + 1);
+}
+
 template <typename Torus>
 __global__ void modulus_switch_inplace(Torus *array, uint32_t size,
                                        uint32_t log_modulus) {
diff --git a/backends/tfhe-cuda-backend/cuda/src/integer/scalar_mul.cu b/backends/tfhe-cuda-backend/cuda/src/integer/scalar_mul.cu
@@ -40,3 +40,12 @@ void cleanup_cuda_scalar_mul(CudaStreamsFFI streams, int8_t **mem_ptr_void) {
   delete mem_ptr;
   *mem_ptr_void = nullptr;
 }
+
+void cuda_small_scalar_multiplication_integer_64_inplace(
+    CudaStreamsFFI streams, CudaRadixCiphertextFFI *lwe_array, uint64_t scalar,
+    const uint32_t message_modulus, const uint32_t carry_modulus) {
+
+  host_integer_small_scalar_mul_radix<uint64_t>(CudaStreams(streams), lwe_array,
+                                                lwe_array, scalar,
+                                                message_modulus, carry_modulus);
+}
diff --git a/backends/tfhe-cuda-backend/cuda/src/pbs/programmable_bootstrap_multibit.cuh b/backends/tfhe-cuda-backend/cuda/src/pbs/programmable_bootstrap_multibit.cuh
@@ -18,20 +18,6 @@
 #include "types/complex/operations.cuh"
 #include <vector>
 
-template <typename Torus, class params>
-__device__ uint32_t calculates_monomial_degree(const Torus *lwe_array_group,
-                                               uint32_t ggsw_idx,
-                                               uint32_t grouping_factor) {
-  Torus x = 0;
-  for (int i = 0; i < grouping_factor; i++) {
-    uint32_t mask_position = grouping_factor - (i + 1);
-    int selection_bit = (ggsw_idx >> mask_position) & 1;
-    x += selection_bit * lwe_array_group[i];
-  }
-
-  return modulus_switch(x, params::log2_degree + 1);
-}
-
 __device__ __forceinline__ int
 get_start_ith_ggsw_offset(uint32_t polynomial_size, int glwe_dimension,
                           uint32_t level_count) {
diff --git a/backends/tfhe-cuda-backend/src/bindings.rs b/backends/tfhe-cuda-backend/src/bindings.rs
@@ -64,23 +64,6 @@ unsafe extern "C" {
         log_modulus: u32,
     );
 }
-unsafe extern "C" {
-    pub fn cuda_improve_noise_modulus_switch_64(
-        stream: *mut ffi::c_void,
-        gpu_index: u32,
-        lwe_array_out: *mut ffi::c_void,
-        lwe_array_in: *const ffi::c_void,
-        lwe_array_indexes: *const ffi::c_void,
-        encrypted_zeros: *const ffi::c_void,
-        lwe_size: u32,
-        num_lwes: u32,
-        num_zeros: u32,
-        input_variance: f64,
-        r_sigma: f64,
-        bound: f64,
-        log_modulus: u32,
-    );
-}
 unsafe extern "C" {
     pub fn cuda_glwe_sample_extract_128(
         stream: *mut ffi::c_void,
@@ -2005,6 +1988,15 @@ unsafe extern "C" {
         ksks: *const *mut ffi::c_void,
     );
 }
+unsafe extern "C" {
+    pub fn cuda_small_scalar_multiplication_integer_64_inplace(
+        streams: CudaStreamsFFI,
+        lwe_array: *mut CudaRadixCiphertextFFI,
+        scalar: u64,
+        message_modulus: u32,
+        carry_modulus: u32,
+    );
+}
 unsafe extern "C" {
     pub fn cleanup_cuda_unchecked_first_index_in_clears_64(
         streams: CudaStreamsFFI,
diff --git a/tfhe/src/core_crypto/commons/noise_formulas/noise_simulation/lwe_programmable_bootstrap.rs b/tfhe/src/core_crypto/commons/noise_formulas/noise_simulation/lwe_programmable_bootstrap.rs
@@ -22,6 +22,9 @@ use crate::core_crypto::fft_impl::fft128::crypto::bootstrap::Fourier128LweBootst
 use crate::core_crypto::fft_impl::fft64::c64;
 use crate::core_crypto::fft_impl::fft64::crypto::bootstrap::FourierLweBootstrapKey;
 
+#[cfg(feature = "gpu")]
+use crate::integer::gpu::server_key::CudaBootstrappingKey;
+
 #[derive(Clone, Copy)]
 pub struct NoiseSimulationLweFourierBsk {
     input_lwe_dimension: LweDimension,
@@ -81,6 +84,48 @@ impl NoiseSimulationLweFourierBsk {
             && decomp_level_count == bsk_decomp_level_count
     }
 
+    #[cfg(feature = "gpu")]
+    pub fn matches_actual_bsk_gpu(&self, lwe_bsk: &CudaBootstrappingKey<u64>) -> bool {
+        let Self {
+            input_lwe_dimension,
+            output_glwe_size: glwe_size,
+            output_polynomial_size: polynomial_size,
+            decomp_base_log,
+            decomp_level_count,
+            noise_distribution: _,
+            modulus: _,
+        } = *self;
+
+        match lwe_bsk {
+            CudaBootstrappingKey::Classic(cuda_bsk) => {
+                let bsk_input_lwe_dimension = cuda_bsk.input_lwe_dimension();
+                let bsk_glwe_size = cuda_bsk.glwe_dimension().to_glwe_size();
+                let bsk_polynomial_size = cuda_bsk.polynomial_size();
+                let bsk_decomp_base_log = cuda_bsk.decomp_base_log();
+                let bsk_decomp_level_count = cuda_bsk.decomp_level_count();
+
+                input_lwe_dimension == bsk_input_lwe_dimension
+                    && glwe_size == bsk_glwe_size
+                    && polynomial_size == bsk_polynomial_size
+                    && decomp_base_log == bsk_decomp_base_log
+                    && decomp_level_count == bsk_decomp_level_count
+            }
+            CudaBootstrappingKey::MultiBit(cuda_mb_bsk) => {
+                let bsk_input_lwe_dimension = cuda_mb_bsk.input_lwe_dimension();
+                let bsk_glwe_size = cuda_mb_bsk.glwe_dimension().to_glwe_size();
+                let bsk_polynomial_size = cuda_mb_bsk.polynomial_size();
+                let bsk_decomp_base_log = cuda_mb_bsk.decomp_base_log();
+                let bsk_decomp_level_count = cuda_mb_bsk.decomp_level_count();
+
+                input_lwe_dimension == bsk_input_lwe_dimension
+                    && glwe_size == bsk_glwe_size
+                    && polynomial_size == bsk_polynomial_size
+                    && decomp_base_log == bsk_decomp_base_log
+                    && decomp_level_count == bsk_decomp_level_count
+            }
+        }
+    }
+
     pub fn input_lwe_dimension(&self) -> LweDimension {
         self.input_lwe_dimension
     }
diff --git a/tfhe/src/core_crypto/gpu/algorithms/lwe_programmable_bootstrapping.rs b/tfhe/src/core_crypto/gpu/algorithms/lwe_programmable_bootstrapping.rs
@@ -60,15 +60,15 @@ pub unsafe fn cuda_programmable_bootstrap_lwe_ciphertext_async<Scalar>(
         accumulator.polynomial_size(),
         bsk.polynomial_size(),
     );
-
+    #[cfg(not(test))]
     assert_eq!(
         input.ciphertext_modulus(),
         output.ciphertext_modulus(),
         "Mismatched CiphertextModulus between input ({:?}) and output ({:?})",
         input.ciphertext_modulus(),
         output.ciphertext_modulus(),
     );
-
+    #[cfg(not(test))]
     assert_eq!(
         input.ciphertext_modulus(),
         accumulator.ciphertext_modulus(),
diff --git a/tfhe/src/core_crypto/gpu/algorithms/lwe_programmable_bootstrapping/cuda_fft64_pbs.rs b/tfhe/src/core_crypto/gpu/algorithms/lwe_programmable_bootstrapping/cuda_fft64_pbs.rs
@@ -0,0 +1,56 @@
+//! CUDA implementations of the LWE programmable bootstrap for noise measurement traits.
+
+use crate::core_crypto::commons::noise_formulas::noise_simulation::traits::LweClassicFftBootstrap;
+use crate::core_crypto::commons::numeric::CastFrom;
+use crate::core_crypto::gpu::algorithms::lwe_programmable_bootstrapping::cuda_programmable_bootstrap_lwe_ciphertext;
+use crate::core_crypto::gpu::entities::glwe_ciphertext_list::CudaGlweCiphertextList;
+use crate::core_crypto::gpu::entities::lwe_bootstrap_key::CudaLweBootstrapKey;
+use crate::core_crypto::gpu::entities::lwe_ciphertext_list::CudaLweCiphertextList;
+use crate::core_crypto::gpu::vec::CudaVec;
+use crate::core_crypto::gpu::CudaSideResources;
+use crate::core_crypto::prelude::{CastInto, UnsignedTorus};
+
+impl<Scalar>
+    LweClassicFftBootstrap<
+        CudaLweCiphertextList<Scalar>,
+        CudaLweCiphertextList<Scalar>,
+        CudaGlweCiphertextList<Scalar>,
+    > for CudaLweBootstrapKey
+where
+    Scalar: UnsignedTorus + CastInto<usize> + CastFrom<usize>,
+{
+    type SideResources = CudaSideResources;
+
+    fn lwe_classic_fft_pbs(
+        &self,
+        input: &CudaLweCiphertextList<Scalar>,
+        output: &mut CudaLweCiphertextList<Scalar>,
+        accumulator: &CudaGlweCiphertextList<Scalar>,
+        side_resources: &mut Self::SideResources,
+    ) {
+        // Create simple index vectors for single operation
+        let count = input.lwe_ciphertext_count().0;
+        let indexes: Vec<Scalar> = (0..count).map(|i| Scalar::cast_from(i)).collect();
+
+        let mut lut_indexes = unsafe { CudaVec::new_async(count, &side_resources.streams, 0) };
+        let mut output_indexes = unsafe { CudaVec::new_async(count, &side_resources.streams, 0) };
+        let mut input_indexes = unsafe { CudaVec::new_async(count, &side_resources.streams, 0) };
+
+        unsafe {
+            lut_indexes.copy_from_cpu_async(&indexes, &side_resources.streams, 0);
+            output_indexes.copy_from_cpu_async(&indexes, &side_resources.streams, 0);
+            input_indexes.copy_from_cpu_async(&indexes, &side_resources.streams, 0);
+        }
+
+        cuda_programmable_bootstrap_lwe_ciphertext(
+            input,
+            output,
+            accumulator,
+            &input_indexes,
+            &lut_indexes,
+            &output_indexes,
+            self,
+            &side_resources.streams,
+        );
+    }
+}
diff --git a/tfhe/src/core_crypto/gpu/mod.rs b/tfhe/src/core_crypto/gpu/mod.rs
@@ -4,7 +4,25 @@ pub mod slice;
 pub mod vec;
 
 use crate::core_crypto::gpu::lwe_bootstrap_key::CudaModulusSwitchNoiseReductionConfiguration;
-use crate::core_crypto::gpu::vec::{CudaVec, GpuIndex};
+use crate::core_crypto::gpu::vec::CudaVec;
+use crate::integer::gpu::ciphertext::info::CudaBlockInfo;
+use crate::GpuIndex;
+
+/// Side resources for CUDA operations in noise simulation
+#[derive(Clone)]
+pub struct CudaSideResources {
+    pub streams: CudaStreams,
+    pub block_info: CudaBlockInfo,
+}
+
+impl CudaSideResources {
+    pub fn new(streams: &CudaStreams, block_info: CudaBlockInfo) -> Self {
+        Self {
+            streams: streams.clone(),
+            block_info,
+        }
+    }
+}
 use crate::core_crypto::prelude::{
     CiphertextModulus, DecompositionBaseLog, DecompositionLevelCount, GlweCiphertextCount,
     GlweDimension, LweBskGroupingFactor, LweCiphertextCount, LweDimension, PolynomialSize,
@@ -831,6 +849,19 @@ pub unsafe fn cuda_modulus_switch_ciphertext_async<T: UnsignedInteger>(
     );
 }
 
+pub fn cuda_modulus_switch_ciphertext<Scalar>(
+    output_lwe_ciphertext: &mut CudaVec<Scalar>,
+    log_modulus: u32,
+    streams: &CudaStreams,
+) where
+    Scalar: UnsignedInteger,
+{
+    unsafe {
+        cuda_modulus_switch_ciphertext_async(streams, output_lwe_ciphertext, log_modulus);
+    }
+    streams.synchronize();
+}
+
 /// Addition of a vector of LWE ciphertexts
 ///
 /// # Safety
diff --git a/tfhe/src/integer/gpu/ciphertext/mod.rs b/tfhe/src/integer/gpu/ciphertext/mod.rs
@@ -10,12 +10,12 @@ use crate::core_crypto::gpu::lwe_ciphertext_list::CudaLweCiphertextList;
 use crate::core_crypto::gpu::vec::CudaVec;
 use crate::core_crypto::gpu::CudaStreams;
 use crate::core_crypto::prelude::{LweCiphertextList, LweCiphertextOwned};
-use crate::integer::gpu::ciphertext::info::{CudaBlockInfo, CudaRadixCiphertextInfo};
 use crate::integer::parameters::LweDimension;
 use crate::integer::{IntegerCiphertext, RadixCiphertext, SignedRadixCiphertext};
 use crate::shortint::{Ciphertext, EncryptionKeyChoice};
 use crate::GpuIndex;
 
+use crate::integer::gpu::ciphertext::info::{CudaBlockInfo, CudaRadixCiphertextInfo};
 pub use compressed_noise_squashed_ciphertext_list::*;
 
 pub trait CudaIntegerRadixCiphertext: Sized {
diff --git a/tfhe/src/integer/gpu/mod.rs b/tfhe/src/integer/gpu/mod.rs
@@ -10184,3 +10184,42 @@ pub(crate) unsafe fn cuda_backend_unchecked_index_of_clear<
     update_noise_degree(index_ct, &ffi_index);
     update_noise_degree(&mut match_ct.0.ciphertext, &ffi_match);
 }
+
+#[allow(clippy::too_many_arguments)]
+/// # Safety
+///
+/// - [CudaStreams::synchronize] __must__ be called after this function as soon as synchronization
+///   is required
+pub unsafe fn unchecked_small_scalar_mul_integer_async(
+    streams: &CudaStreams,
+    lwe_array: &mut CudaRadixCiphertext,
+    small_scalar: u64,
+    message_modulus: MessageModulus,
+    carry_modulus: CarryModulus,
+) {
+    assert_eq!(
+        streams.gpu_indexes[0],
+        lwe_array.d_blocks.0.d_vec.gpu_index(0),
+        "GPU error: all data should reside on the same GPU."
+    );
+    let mut lwe_array_degrees = lwe_array.info.blocks.iter().map(|b| b.degree.0).collect();
+    let mut lwe_array_noise_levels = lwe_array
+        .info
+        .blocks
+        .iter()
+        .map(|b| b.noise_level.0)
+        .collect();
+    let mut cuda_ffi_lwe_array = prepare_cuda_radix_ffi(
+        lwe_array,
+        &mut lwe_array_degrees,
+        &mut lwe_array_noise_levels,
+    );
+
+    cuda_small_scalar_multiplication_integer_64_inplace(
+        streams.ffi(),
+        &mut cuda_ffi_lwe_array,
+        small_scalar,
+        message_modulus.0 as u32,
+        carry_modulus.0 as u32,
+    );
+}
diff --git a/tfhe/src/integer/gpu/server_key/radix/mod.rs b/tfhe/src/integer/gpu/server_key/radix/mod.rs
@@ -63,6 +63,8 @@ mod aes256;
 #[cfg(test)]
 mod tests_long_run;
 #[cfg(test)]
+mod tests_noise_distribution;
+#[cfg(test)]
 mod tests_signed;
 #[cfg(test)]
 mod tests_unsigned;
diff --git a/tfhe/src/integer/gpu/server_key/radix/tests_noise_distribution/br_dp_ks_ms.rs b/tfhe/src/integer/gpu/server_key/radix/tests_noise_distribution/br_dp_ks_ms.rs
diff --git a/tfhe/src/integer/gpu/server_key/radix/tests_noise_distribution/dp_ks_ms.rs b/tfhe/src/integer/gpu/server_key/radix/tests_noise_distribution/dp_ks_ms.rs
diff --git a/tfhe/src/integer/gpu/server_key/radix/tests_noise_distribution/mod.rs b/tfhe/src/integer/gpu/server_key/radix/tests_noise_distribution/mod.rs
diff --git a/tfhe/src/integer/gpu/server_key/radix/tests_noise_distribution/utils/mod.rs b/tfhe/src/integer/gpu/server_key/radix/tests_noise_distribution/utils/mod.rs
diff --git a/tfhe/src/integer/gpu/server_key/radix/tests_noise_distribution/utils/noise_simulation.rs b/tfhe/src/integer/gpu/server_key/radix/tests_noise_distribution/utils/noise_simulation.rs
