From faf200218bd1254f8c3b641d9a083a9e20f1ebcb Mon Sep 17 00:00:00 2001
From: Pedro Alves <pedro.alves@zama.ai>
Date: Thu, 12 Sep 2024 11:36:44 -0300
Subject: [PATCH] chore(gpu): add checks to ensure limits for compression

---
 backends/tfhe-cuda-backend/cuda/src/crypto/keyswitch.cuh    | 6 ++++++
 .../cuda/src/integer/compression/compression.cuh            | 6 ++++++
 tfhe/src/integer/gpu/mod.rs                                 | 4 ++--
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/backends/tfhe-cuda-backend/cuda/src/crypto/keyswitch.cuh b/backends/tfhe-cuda-backend/cuda/src/crypto/keyswitch.cuh
index 3d1618e413..abf3e71804 100644
--- a/backends/tfhe-cuda-backend/cuda/src/crypto/keyswitch.cuh
+++ b/backends/tfhe-cuda-backend/cuda/src/crypto/keyswitch.cuh
@@ -280,6 +280,12 @@ __host__ void host_packing_keyswitch_lwe_list_to_glwe(
     uint32_t lwe_dimension_in, uint32_t glwe_dimension,
     uint32_t polynomial_size, uint32_t base_log, uint32_t level_count,
     uint32_t num_lwes) {
+
+  if (num_lwes > polynomial_size)
+    PANIC("Cuda error: too many LWEs to pack. The number of LWEs should be "
+          "smaller than "
+          "polynomial_size.")
+
   cudaSetDevice(gpu_index);
   int glwe_accumulator_size = (glwe_dimension + 1) * polynomial_size;
 
diff --git a/backends/tfhe-cuda-backend/cuda/src/integer/compression/compression.cuh b/backends/tfhe-cuda-backend/cuda/src/integer/compression/compression.cuh
index f9072ac376..29ec8639f4 100644
--- a/backends/tfhe-cuda-backend/cuda/src/integer/compression/compression.cuh
+++ b/backends/tfhe-cuda-backend/cuda/src/integer/compression/compression.cuh
@@ -180,6 +180,12 @@ host_integer_decompress(cudaStream_t *streams, uint32_t *gpu_indexes,
                         uint32_t indexes_array_size, void **bsks,
                         int_decompression<Torus> *mem_ptr) {
 
+  auto polynomial_size = mem_ptr->encryption_params.polynomial_size;
+  if (indexes_array_size > polynomial_size)
+    PANIC("Cuda error: too many LWEs to decompress. The number of LWEs should "
+          "be smaller than "
+          "polynomial_size.")
+
   auto extracted_glwe = mem_ptr->tmp_extracted_glwe;
   auto compression_params = mem_ptr->compression_params;
   host_extract<Torus>(streams[0], gpu_indexes[0], extracted_glwe,
diff --git a/tfhe/src/integer/gpu/mod.rs b/tfhe/src/integer/gpu/mod.rs
index ba0bb62986..1a61768ff2 100644
--- a/tfhe/src/integer/gpu/mod.rs
+++ b/tfhe/src/integer/gpu/mod.rs
@@ -373,7 +373,7 @@ pub unsafe fn decompress_integer_radix_async<T: UnsignedInteger, B: Numeric>(
     pbs_level: DecompositionLevelCount,
     storage_log_modulus: u32,
     vec_indexes: &CudaVec<u32>,
-    num_blocks: u32,
+    num_lwes: u32,
 ) {
     assert_eq!(
         streams.gpu_indexes[0],
@@ -403,7 +403,7 @@ pub unsafe fn decompress_integer_radix_async<T: UnsignedInteger, B: Numeric>(
         lwe_dimension.0 as u32,
         pbs_level.0 as u32,
         pbs_base_log.0 as u32,
-        num_blocks,
+        num_lwes,
         message_modulus.0 as u32,
         carry_modulus.0 as u32,
         PBSType::Classical as u32,