refactor(gpu): moving cast_to_signed to the backend

enzodimaria · enzodimaria · commit 51b7388dfde1 · 2025-12-02T09:26:34.000+01:00
diff --git a/backends/tfhe-cuda-backend/cuda/include/integer/cast.h b/backends/tfhe-cuda-backend/cuda/include/integer/cast.h
@@ -127,3 +127,38 @@ template <typename Torus> struct int_cast_to_unsigned_buffer {
     cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
   }
 };
+
+template <typename Torus> struct int_cast_to_signed_buffer {
+  int_radix_params params;
+  bool allocate_gpu_memory;
+  uint32_t num_input_blocks;
+  uint32_t target_num_blocks;
+
+  int_extend_radix_with_sign_msb_buffer<Torus> *extend_buffer;
+
+  int_cast_to_signed_buffer(CudaStreams streams, int_radix_params params,
+                            uint32_t num_input_blocks,
+                            uint32_t target_num_blocks, bool input_is_signed,
+                            bool allocate_gpu_memory, uint64_t &size_tracker) {
+    this->params = params;
+    this->allocate_gpu_memory = allocate_gpu_memory;
+    this->num_input_blocks = num_input_blocks;
+    this->target_num_blocks = target_num_blocks;
+    this->extend_buffer = nullptr;
+
+    if (input_is_signed && target_num_blocks > num_input_blocks) {
+      uint32_t num_additional_blocks = target_num_blocks - num_input_blocks;
+      this->extend_buffer = new int_extend_radix_with_sign_msb_buffer<Torus>(
+          streams, params, num_input_blocks, num_additional_blocks,
+          allocate_gpu_memory, size_tracker);
+    }
+  }
+
+  void release(CudaStreams streams) {
+    if (this->extend_buffer) {
+      this->extend_buffer->release(streams);
+      delete this->extend_buffer;
+    }
+    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
+  }
+};
diff --git a/backends/tfhe-cuda-backend/cuda/include/integer/integer.h b/backends/tfhe-cuda-backend/cuda/include/integer/integer.h
@@ -610,25 +610,6 @@ void cuda_integer_unsigned_scalar_div_radix_64(
 void cleanup_cuda_integer_unsigned_scalar_div_radix_64(CudaStreamsFFI streams,
                                                        int8_t **mem_ptr_void);
 
-uint64_t scratch_cuda_extend_radix_with_sign_msb_64(
-    CudaStreamsFFI streams, int8_t **mem_ptr, uint32_t glwe_dimension,
-    uint32_t polynomial_size, uint32_t lwe_dimension, uint32_t ks_level,
-    uint32_t ks_base_log, uint32_t pbs_level, uint32_t pbs_base_log,
-    uint32_t grouping_factor, uint32_t num_blocks,
-    uint32_t num_additional_blocks, uint32_t message_modulus,
-    uint32_t carry_modulus, PBS_TYPE pbs_type, bool allocate_gpu_memory,
-    PBS_MS_REDUCTION_T noise_reduction_type);
-
-void cuda_extend_radix_with_sign_msb_64(CudaStreamsFFI streams,
-                                        CudaRadixCiphertextFFI *output,
-                                        CudaRadixCiphertextFFI const *input,
-                                        int8_t *mem_ptr,
-                                        uint32_t num_additional_blocks,
-                                        void *const *bsks, void *const *ksks);
-
-void cleanup_cuda_extend_radix_with_sign_msb_64(CudaStreamsFFI streams,
-                                                int8_t **mem_ptr_void);
-
 uint64_t scratch_cuda_integer_signed_scalar_div_radix_64(
     CudaStreamsFFI streams, int8_t **mem_ptr, uint32_t glwe_dimension,
     uint32_t polynomial_size, uint32_t lwe_dimension, uint32_t ks_level,
@@ -1004,6 +985,24 @@ void cuda_unchecked_index_of_clear_64(
 
 void cleanup_cuda_unchecked_index_of_clear_64(CudaStreamsFFI streams,
                                               int8_t **mem_ptr_void);
+
+uint64_t scratch_cuda_cast_to_signed_64(
+    CudaStreamsFFI streams, int8_t **mem_ptr, uint32_t glwe_dimension,
+    uint32_t polynomial_size, uint32_t lwe_dimension, uint32_t ks_level,
+    uint32_t ks_base_log, uint32_t pbs_level, uint32_t pbs_base_log,
+    uint32_t grouping_factor, uint32_t num_input_blocks,
+    uint32_t target_num_blocks, uint32_t message_modulus,
+    uint32_t carry_modulus, PBS_TYPE pbs_type, bool input_is_signed,
+    bool allocate_gpu_memory, PBS_MS_REDUCTION_T noise_reduction_type);
+
+void cuda_cast_to_signed_64(CudaStreamsFFI streams,
+                            CudaRadixCiphertextFFI *output,
+                            CudaRadixCiphertextFFI const *input, int8_t *mem,
+                            bool input_is_signed, void *const *bsks,
+                            void *const *ksks);
+
+void cleanup_cuda_cast_to_signed_64(CudaStreamsFFI streams,
+                                    int8_t **mem_ptr_void);
 } // extern C
 
 #endif // CUDA_INTEGER_H
diff --git a/backends/tfhe-cuda-backend/cuda/src/integer/cast.cu b/backends/tfhe-cuda-backend/cuda/src/integer/cast.cu
@@ -27,53 +27,6 @@ void trim_radix_blocks_msb_64(CudaRadixCiphertextFFI *output,
   cuda_synchronize_stream(cuda_streams.stream(0), cuda_streams.gpu_index(0));
 }
 
-uint64_t scratch_cuda_extend_radix_with_sign_msb_64(
-    CudaStreamsFFI streams, int8_t **mem_ptr, uint32_t glwe_dimension,
-    uint32_t polynomial_size, uint32_t lwe_dimension, uint32_t ks_level,
-    uint32_t ks_base_log, uint32_t pbs_level, uint32_t pbs_base_log,
-    uint32_t grouping_factor, uint32_t num_blocks,
-    uint32_t num_additional_blocks, uint32_t message_modulus,
-    uint32_t carry_modulus, PBS_TYPE pbs_type, bool allocate_gpu_memory,
-    PBS_MS_REDUCTION_T noise_reduction_type) {
-
-  int_radix_params params(pbs_type, glwe_dimension, polynomial_size,
-                          glwe_dimension * polynomial_size, lwe_dimension,
-                          ks_level, ks_base_log, pbs_level, pbs_base_log,
-                          grouping_factor, message_modulus, carry_modulus,
-                          noise_reduction_type);
-
-  return scratch_extend_radix_with_sign_msb<uint64_t>(
-      CudaStreams(streams),
-      (int_extend_radix_with_sign_msb_buffer<uint64_t> **)mem_ptr, params,
-      num_blocks, num_additional_blocks, allocate_gpu_memory);
-}
-
-void cuda_extend_radix_with_sign_msb_64(CudaStreamsFFI streams,
-                                        CudaRadixCiphertextFFI *output,
-                                        CudaRadixCiphertextFFI const *input,
-                                        int8_t *mem_ptr,
-                                        uint32_t num_additional_blocks,
-                                        void *const *bsks, void *const *ksks) {
-  PUSH_RANGE("cast")
-  host_extend_radix_with_sign_msb<uint64_t>(
-      CudaStreams(streams), output, input,
-      (int_extend_radix_with_sign_msb_buffer<uint64_t> *)mem_ptr,
-      num_additional_blocks, bsks, (uint64_t **)ksks);
-  POP_RANGE()
-}
-
-void cleanup_cuda_extend_radix_with_sign_msb_64(CudaStreamsFFI streams,
-                                                int8_t **mem_ptr_void) {
-  PUSH_RANGE("clean cast")
-  int_extend_radix_with_sign_msb_buffer<uint64_t> *mem_ptr =
-      (int_extend_radix_with_sign_msb_buffer<uint64_t> *)(*mem_ptr_void);
-
-  mem_ptr->release(CudaStreams(streams));
-  POP_RANGE()
-  delete mem_ptr;
-  *mem_ptr_void = nullptr;
-}
-
 uint64_t scratch_cuda_cast_to_unsigned_64(
     CudaStreamsFFI streams, int8_t **mem_ptr, uint32_t glwe_dimension,
     uint32_t polynomial_size, uint32_t big_lwe_dimension,
@@ -116,3 +69,46 @@ void cleanup_cuda_cast_to_unsigned_64(CudaStreamsFFI streams,
   delete mem_ptr;
   *mem_ptr_void = nullptr;
 }
+
+uint64_t scratch_cuda_cast_to_signed_64(
+    CudaStreamsFFI streams, int8_t **mem_ptr, uint32_t glwe_dimension,
+    uint32_t polynomial_size, uint32_t lwe_dimension, uint32_t ks_level,
+    uint32_t ks_base_log, uint32_t pbs_level, uint32_t pbs_base_log,
+    uint32_t grouping_factor, uint32_t num_input_blocks,
+    uint32_t target_num_blocks, uint32_t message_modulus,
+    uint32_t carry_modulus, PBS_TYPE pbs_type, bool input_is_signed,
+    bool allocate_gpu_memory, PBS_MS_REDUCTION_T noise_reduction_type) {
+
+  int_radix_params params(pbs_type, glwe_dimension, polynomial_size,
+                          glwe_dimension * polynomial_size, lwe_dimension,
+                          ks_level, ks_base_log, pbs_level, pbs_base_log,
+                          grouping_factor, message_modulus, carry_modulus,
+                          noise_reduction_type);
+
+  return scratch_cuda_cast_to_signed<uint64_t>(
+      CudaStreams(streams), (int_cast_to_signed_buffer<uint64_t> **)mem_ptr,
+      params, num_input_blocks, target_num_blocks, input_is_signed,
+      allocate_gpu_memory);
+}
+
+void cuda_cast_to_signed_64(CudaStreamsFFI streams,
+                            CudaRadixCiphertextFFI *output,
+                            CudaRadixCiphertextFFI const *input, int8_t *mem,
+                            bool input_is_signed, void *const *bsks,
+                            void *const *ksks) {
+
+  host_cast_to_signed<uint64_t>(CudaStreams(streams), output, input,
+                                (int_cast_to_signed_buffer<uint64_t> *)mem,
+                                input_is_signed, bsks, (uint64_t **)ksks);
+}
+
+void cleanup_cuda_cast_to_signed_64(CudaStreamsFFI streams,
+                                    int8_t **mem_ptr_void) {
+  int_cast_to_signed_buffer<uint64_t> *mem_ptr =
+      (int_cast_to_signed_buffer<uint64_t> *)(*mem_ptr_void);
+
+  mem_ptr->release(CudaStreams(streams));
+
+  delete mem_ptr;
+  *mem_ptr_void = nullptr;
+}
diff --git a/backends/tfhe-cuda-backend/cuda/src/integer/cast.cuh b/backends/tfhe-cuda-backend/cuda/src/integer/cast.cuh
@@ -160,4 +160,49 @@ host_cast_to_unsigned(CudaStreams streams, CudaRadixCiphertextFFI *output,
   }
 }
 
+template <typename Torus>
+uint64_t
+scratch_cuda_cast_to_signed(CudaStreams streams,
+                            int_cast_to_signed_buffer<Torus> **mem_ptr,
+                            int_radix_params params, uint32_t num_input_blocks,
+                            uint32_t target_num_blocks, bool input_is_signed,
+                            bool allocate_gpu_memory) {
+
+  uint64_t size_tracker = 0;
+  *mem_ptr = new int_cast_to_signed_buffer<Torus>(
+      streams, params, num_input_blocks, target_num_blocks, input_is_signed,
+      allocate_gpu_memory, size_tracker);
+
+  return size_tracker;
+}
+
+template <typename Torus>
+__host__ void
+host_cast_to_signed(CudaStreams streams, CudaRadixCiphertextFFI *output,
+                    CudaRadixCiphertextFFI const *input,
+                    int_cast_to_signed_buffer<Torus> *mem_ptr,
+                    bool input_is_signed, void *const *bsks, Torus **ksks) {
+
+  uint32_t current_num_blocks = input->num_radix_blocks;
+  uint32_t target_num_blocks = mem_ptr->target_num_blocks;
+
+  if (input_is_signed) {
+    if (target_num_blocks > current_num_blocks) {
+      uint32_t num_blocks_to_add = target_num_blocks - current_num_blocks;
+      host_extend_radix_with_sign_msb<Torus>(streams, output, input,
+                                             mem_ptr->extend_buffer,
+                                             num_blocks_to_add, bsks, ksks);
+    } else {
+      host_trim_radix_blocks_msb<Torus>(output, input, streams);
+    }
+  } else {
+    if (target_num_blocks > current_num_blocks) {
+      host_extend_radix_with_trivial_zero_blocks_msb<Torus>(output, input,
+                                                            streams);
+    } else {
+      host_trim_radix_blocks_msb<Torus>(output, input, streams);
+    }
+  }
+}
+
 #endif
diff --git a/backends/tfhe-cuda-backend/src/bindings.rs b/backends/tfhe-cuda-backend/src/bindings.rs
@@ -1353,44 +1353,6 @@ unsafe extern "C" {
         mem_ptr_void: *mut *mut i8,
     );
 }
-unsafe extern "C" {
-    pub fn scratch_cuda_extend_radix_with_sign_msb_64(
-        streams: CudaStreamsFFI,
-        mem_ptr: *mut *mut i8,
-        glwe_dimension: u32,
-        polynomial_size: u32,
-        lwe_dimension: u32,
-        ks_level: u32,
-        ks_base_log: u32,
-        pbs_level: u32,
-        pbs_base_log: u32,
-        grouping_factor: u32,
-        num_blocks: u32,
-        num_additional_blocks: u32,
-        message_modulus: u32,
-        carry_modulus: u32,
-        pbs_type: PBS_TYPE,
-        allocate_gpu_memory: bool,
-        noise_reduction_type: PBS_MS_REDUCTION_T,
-    ) -> u64;
-}
-unsafe extern "C" {
-    pub fn cuda_extend_radix_with_sign_msb_64(
-        streams: CudaStreamsFFI,
-        output: *mut CudaRadixCiphertextFFI,
-        input: *const CudaRadixCiphertextFFI,
-        mem_ptr: *mut i8,
-        num_additional_blocks: u32,
-        bsks: *const *mut ffi::c_void,
-        ksks: *const *mut ffi::c_void,
-    );
-}
-unsafe extern "C" {
-    pub fn cleanup_cuda_extend_radix_with_sign_msb_64(
-        streams: CudaStreamsFFI,
-        mem_ptr_void: *mut *mut i8,
-    );
-}
 unsafe extern "C" {
     pub fn scratch_cuda_integer_signed_scalar_div_radix_64(
         streams: CudaStreamsFFI,
@@ -2186,6 +2148,42 @@ unsafe extern "C" {
         mem_ptr_void: *mut *mut i8,
     );
 }
+unsafe extern "C" {
+    pub fn scratch_cuda_cast_to_signed_64(
+        streams: CudaStreamsFFI,
+        mem_ptr: *mut *mut i8,
+        glwe_dimension: u32,
+        polynomial_size: u32,
+        lwe_dimension: u32,
+        ks_level: u32,
+        ks_base_log: u32,
+        pbs_level: u32,
+        pbs_base_log: u32,
+        grouping_factor: u32,
+        num_input_blocks: u32,
+        target_num_blocks: u32,
+        message_modulus: u32,
+        carry_modulus: u32,
+        pbs_type: PBS_TYPE,
+        input_is_signed: bool,
+        allocate_gpu_memory: bool,
+        noise_reduction_type: PBS_MS_REDUCTION_T,
+    ) -> u64;
+}
+unsafe extern "C" {
+    pub fn cuda_cast_to_signed_64(
+        streams: CudaStreamsFFI,
+        output: *mut CudaRadixCiphertextFFI,
+        input: *const CudaRadixCiphertextFFI,
+        mem: *mut i8,
+        input_is_signed: bool,
+        bsks: *const *mut ffi::c_void,
+        ksks: *const *mut ffi::c_void,
+    );
+}
+unsafe extern "C" {
+    pub fn cleanup_cuda_cast_to_signed_64(streams: CudaStreamsFFI, mem_ptr_void: *mut *mut i8);
+}
 unsafe extern "C" {
     pub fn scratch_cuda_integer_compress_radix_ciphertext_64(
         streams: CudaStreamsFFI,
diff --git a/tfhe/src/integer/gpu/mod.rs b/tfhe/src/integer/gpu/mod.rs
diff --git a/tfhe/src/integer/gpu/server_key/radix/mod.rs b/tfhe/src/integer/gpu/server_key/radix/mod.rs
