refactor(gpu): moving cast_to_signed to the backend

Author: enzodimaria

backends/tfhe-cuda-backend/cuda/include/integer/cast.h

@@ -127,3 +127,37 @@ 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 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 (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;
+    }
+  }
+};

backends/tfhe-cuda-backend/cuda/include/integer/integer.h

@@ -1004,6 +1004,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 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

backends/tfhe-cuda-backend/cuda/src/integer/cast.cu

@@ -116,3 +116,45 @@ 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 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, 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;
+}

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 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, 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

backends/tfhe-cuda-backend/src/bindings.rs

@@ -2186,6 +2186,41 @@ 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,
+        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,

tfhe/src/integer/gpu/mod.rs

@@ -5910,82 +5910,6 @@ pub(crate) unsafe fn cuda_backend_unchecked_partial_sum_ciphertexts_assign<
     update_noise_degree(result, &cuda_ffi_result);
 }
 
-#[allow(clippy::too_many_arguments)]
-/// # Safety
-///
-/// - The data must not be moved or dropped while being used by the CUDA kernel.
-/// - This function assumes exclusive access to the passed data; violating this may lead to
-///   undefined behavior.
-pub(crate) unsafe fn cuda_backend_extend_radix_with_sign_msb<T: UnsignedInteger, B: Numeric>(
-    streams: &CudaStreams,
-    output: &mut CudaRadixCiphertext,
-    ct: &CudaRadixCiphertext,
-    bootstrapping_key: &CudaVec<B>,
-    keyswitch_key: &CudaVec<T>,
-    lwe_dimension: LweDimension,
-    glwe_dimension: GlweDimension,
-    polynomial_size: PolynomialSize,
-    ks_level: DecompositionLevelCount,
-    ks_base_log: DecompositionBaseLog,
-    pbs_level: DecompositionLevelCount,
-    pbs_base_log: DecompositionBaseLog,
-    num_additional_blocks: u32,
-    pbs_type: PBSType,
-    grouping_factor: LweBskGroupingFactor,
-    ms_noise_reduction_configuration: Option<&CudaModulusSwitchNoiseReductionConfiguration>,
-) {
-    let message_modulus = ct.info.blocks.first().unwrap().message_modulus;
-    let carry_modulus = ct.info.blocks.first().unwrap().carry_modulus;
-
-    let noise_reduction_type = resolve_ms_noise_reduction_config(ms_noise_reduction_configuration);
-
-    let mut mem_ptr: *mut i8 = std::ptr::null_mut();
-
-    let mut input_degrees = ct.info.blocks.iter().map(|b| b.degree.0).collect();
-    let mut input_noise_levels = ct.info.blocks.iter().map(|b| b.noise_level.0).collect();
-    let cuda_ffi_radix_input =
-        prepare_cuda_radix_ffi(ct, &mut input_degrees, &mut input_noise_levels);
-
-    let mut output_degrees = output.info.blocks.iter().map(|b| b.degree.0).collect();
-    let mut output_noise_levels = output.info.blocks.iter().map(|b| b.noise_level.0).collect();
-    let mut cuda_ffi_radix_output =
-        prepare_cuda_radix_ffi(output, &mut output_degrees, &mut output_noise_levels);
-
-    scratch_cuda_extend_radix_with_sign_msb_64(
-        streams.ffi(),
-        std::ptr::addr_of_mut!(mem_ptr),
-        glwe_dimension.0 as u32,
-        polynomial_size.0 as u32,
-        lwe_dimension.0 as u32,
-        ks_level.0 as u32,
-        ks_base_log.0 as u32,
-        pbs_level.0 as u32,
-        pbs_base_log.0 as u32,
-        grouping_factor.0 as u32,
-        1u32,
-        num_additional_blocks,
-        message_modulus.0 as u32,
-        carry_modulus.0 as u32,
-        pbs_type as u32,
-        true,
-        noise_reduction_type as u32,
-    );
-
-    cuda_extend_radix_with_sign_msb_64(
-        streams.ffi(),
-        &raw mut cuda_ffi_radix_output,
-        &raw const cuda_ffi_radix_input,
-        mem_ptr,
-        num_additional_blocks,
-        bootstrapping_key.ptr.as_ptr(),
-        keyswitch_key.ptr.as_ptr(),
-    );
-
-    cleanup_cuda_extend_radix_with_sign_msb_64(streams.ffi(), std::ptr::addr_of_mut!(mem_ptr));
-
-    update_noise_degree(output, &cuda_ffi_radix_output);
-}
-
 #[allow(clippy::too_many_arguments)]
 /// # Safety
 ///
@@ -10184,3 +10108,83 @@ 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
+///
+/// - The data must not be moved or dropped while being used by the CUDA kernel.
+/// - This function assumes exclusive access to the passed data; violating this may lead to
+///   undefined behavior.
+pub(crate) unsafe fn cuda_backend_cast_to_signed<T: UnsignedInteger, B: Numeric>(
+    streams: &CudaStreams,
+    output: &mut CudaRadixCiphertext,
+    input: &CudaRadixCiphertext,
+    input_is_signed: bool,
+    bootstrapping_key: &CudaVec<B>,
+    keyswitch_key: &CudaVec<T>,
+    message_modulus: MessageModulus,
+    carry_modulus: CarryModulus,
+    glwe_dimension: GlweDimension,
+    polynomial_size: PolynomialSize,
+    big_lwe_dimension: LweDimension,
+    ks_level: DecompositionLevelCount,
+    ks_base_log: DecompositionBaseLog,
+    pbs_level: DecompositionLevelCount,
+    pbs_base_log: DecompositionBaseLog,
+    pbs_type: PBSType,
+    grouping_factor: LweBskGroupingFactor,
+    ms_noise_reduction_configuration: Option<&CudaModulusSwitchNoiseReductionConfiguration>,
+) {
+    assert_eq!(streams.gpu_indexes[0], bootstrapping_key.gpu_index(0));
+    assert_eq!(streams.gpu_indexes[0], keyswitch_key.gpu_index(0));
+
+    let num_input_blocks = input.d_blocks.lwe_ciphertext_count().0 as u32;
+    let target_num_blocks = output.d_blocks.lwe_ciphertext_count().0 as u32;
+
+    let noise_reduction_type = resolve_ms_noise_reduction_config(ms_noise_reduction_configuration);
+
+    let mut input_degrees = input.info.blocks.iter().map(|b| b.degree.0).collect();
+    let mut input_noise_levels = input.info.blocks.iter().map(|b| b.noise_level.0).collect();
+    let cuda_ffi_input = prepare_cuda_radix_ffi(input, &mut input_degrees, &mut input_noise_levels);
+
+    let mut output_degrees = output.info.blocks.iter().map(|b| b.degree.0).collect();
+    let mut output_noise_levels = output.info.blocks.iter().map(|b| b.noise_level.0).collect();
+    let mut cuda_ffi_output =
+        prepare_cuda_radix_ffi(output, &mut output_degrees, &mut output_noise_levels);
+
+    let mut mem_ptr: *mut i8 = std::ptr::null_mut();
+
+    scratch_cuda_cast_to_signed_64(
+        streams.ffi(),
+        std::ptr::addr_of_mut!(mem_ptr),
+        glwe_dimension.0 as u32,
+        polynomial_size.0 as u32,
+        big_lwe_dimension.0 as u32,
+        ks_level.0 as u32,
+        ks_base_log.0 as u32,
+        pbs_level.0 as u32,
+        pbs_base_log.0 as u32,
+        grouping_factor.0 as u32,
+        num_input_blocks,
+        target_num_blocks,
+        message_modulus.0 as u32,
+        carry_modulus.0 as u32,
+        pbs_type as u32,
+        true,
+        noise_reduction_type as u32,
+    );
+
+    cuda_cast_to_signed_64(
+        streams.ffi(),
+        &raw mut cuda_ffi_output,
+        &raw const cuda_ffi_input,
+        mem_ptr,
+        input_is_signed,
+        bootstrapping_key.ptr.as_ptr(),
+        keyswitch_key.ptr.as_ptr(),
+    );
+
+    cleanup_cuda_cast_to_signed_64(streams.ffi(), std::ptr::addr_of_mut!(mem_ptr));
+
+    update_noise_degree(output, &cuda_ffi_output);
+}