chore(gpu): refactor prepare_count_of_consecutive_bits

Author: enzodimaria

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

@@ -37,6 +37,10 @@ enum SIGNED_OPERATION { ADDITION = 1, SUBTRACTION = -1 };
 
 enum outputFlag { FLAG_NONE = 0, FLAG_OVERFLOW = 1, FLAG_CARRY = 2 };
 
+enum Direction { TRAILING = 0, LEADING = 1 };
+
+enum BitValue { ZERO = 0, ONE = 1 };
+
 extern "C" {
 
 typedef struct {
@@ -538,5 +542,26 @@ void cleanup_cuda_integer_is_at_least_one_comparisons_block_true(
     void *const *streams, uint32_t const *gpu_indexes, uint32_t gpu_count,
     int8_t **mem_ptr_void);
 
+uint64_t scratch_cuda_prepare_count_of_consecutive_bits_buffer_kb_64(
+    void *const *streams, uint32_t const *gpu_indexes, const uint32_t gpu_count,
+    int8_t **mem_ptr, const uint32_t num_radix_blocks, const Direction dir,
+    const BitValue bit_value, const bool allocate_gpu_memory,
+    const uint32_t glwe_dimension, const uint32_t polynomial_size,
+    const uint32_t lwe_dimension, const uint32_t ks_level,
+    const uint32_t ks_base_log, const uint32_t pbs_level,
+    const uint32_t pbs_base_log, const uint32_t grouping_factor,
+    const uint32_t message_modulus, const uint32_t carry_modulus,
+    const PBS_TYPE pbs_type, const bool allocate_ms_array);
+
+void host_cuda_prepare_count_of_consecutive_bits_buffer_kb_64(
+    void *const *streams, uint32_t const *gpu_indexes, const uint32_t gpu_count,
+    CudaRadixCiphertextFFI *output, CudaRadixCiphertextFFI const *input,
+    int8_t **mem_ptr, void *const *bsks, void *const *ksks,
+    CudaModulusSwitchNoiseReductionKeyFFI const *ms_noise_reduction_key);
+
+void cleanup_cuda_prepare_count_of_consecutive_bits_buffer_kb_64(
+    void *const *streams, uint32_t const *gpu_indexes, uint32_t gpu_count,
+    int8_t **mem_ptr_void);
+
 } // extern C
 #endif // CUDA_INTEGER_H

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

@@ -1,6 +1,8 @@
 #ifndef CUDA_INTEGER_UTILITIES_H
 #define CUDA_INTEGER_UTILITIES_H
 
+#include <cinttypes>
+
 #include "integer.h"
 #include "integer/radix_ciphertext.cuh"
 #include "integer/radix_ciphertext.h"
@@ -17,6 +19,10 @@ class NoiseLevel {
   static const uint64_t UNKNOWN = std::numeric_limits<uint64_t>::max();
 };
 
+constexpr BitValue opposite(const BitValue bit) noexcept {
+  return (bit == BitValue::ZERO) ? BitValue::ONE : BitValue::ZERO;
+}
+
 template <typename Torus>
 __global__ void radix_blocks_rotate_right(Torus *dst, Torus *src,
                                           uint32_t value, uint32_t blocks_count,
@@ -4400,6 +4406,94 @@ template <typename Torus> struct unsigned_int_div_rem_memory {
   }
 };
 
+template <typename Torus> struct int_prepare_count_of_consecutive_bits_buffer {
+  int_radix_params params;
+  int_radix_lut<Torus> *uni_lut;
+  int_radix_lut<Torus> *bi_lut;
+  CudaRadixCiphertextFFI *copy_ct;
+  Direction dir;
+  BitValue bit_val;
+  bool gpu_memory_allocated;
+
+  int_prepare_count_of_consecutive_bits_buffer(
+      cudaStream_t const *streams, uint32_t const *gpu_indexes,
+      uint32_t gpu_count, Direction dir, BitValue bit_val,
+      int_radix_params params, uint32_t num_radix_blocks,
+      bool allocate_gpu_memory, uint64_t *size_tracker) {
+    this->params = params;
+    this->dir = dir;
+    this->bit_val = bit_val;
+    this->gpu_memory_allocated = allocate_gpu_memory;
+    this->copy_ct = new CudaRadixCiphertextFFI;
+
+    const uint32_t bits = (uint32_t)std::log2(params.message_modulus);
+
+    create_zero_radix_ciphertext_async<Torus>(
+        streams[0], gpu_indexes[0], copy_ct, num_radix_blocks,
+        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
+
+    this->uni_lut = new int_radix_lut<Torus>(streams, gpu_indexes, gpu_count,
+                                             params, 1, num_radix_blocks,
+                                             allocate_gpu_memory, size_tracker);
+
+    auto lut_fn = [params, dir, bit_val, bits](Torus x) -> Torus {
+      x %= params.message_modulus;
+      Torus count = 0;
+
+      const int start = (dir == TRAILING ? 0 : (int)bits - 1);
+      const int end = (dir == TRAILING ? (int)bits : -1);
+      const int step = (dir == TRAILING ? 1 : -1);
+
+      for (uint32_t i = start; i != end; i += step) {
+        if (((x >> i) & (Torus)1) == (Torus)opposite(bit_val))
+          break;
+        ++count;
+      }
+      return count;
+    };
+
+    generate_device_accumulator<Torus>(
+        streams[0], gpu_indexes[0], uni_lut->get_lut(0, 0),
+        uni_lut->get_degree(0), uni_lut->get_max_degree(0),
+        params.glwe_dimension, params.polynomial_size, params.message_modulus,
+        params.carry_modulus, lut_fn, gpu_memory_allocated);
+
+    uni_lut->broadcast_lut(streams, gpu_indexes, 0);
+
+    this->bi_lut = new int_radix_lut<Torus>(streams, gpu_indexes, gpu_count,
+                                            params, 1, num_radix_blocks,
+                                            allocate_gpu_memory, size_tracker);
+
+    auto sum_lut = [bits](Torus block_num_bit_count,
+                          Torus more_significant_block_bit_count) -> Torus {
+      return (more_significant_block_bit_count == bits) ? block_num_bit_count
+                                                        : 0;
+    };
+
+    generate_device_accumulator_bivariate<Torus>(
+        streams[0], gpu_indexes[0], bi_lut->get_lut(0, 0),
+        bi_lut->get_degree(0), bi_lut->get_max_degree(0), params.glwe_dimension,
+        params.polynomial_size, params.message_modulus, params.carry_modulus,
+        sum_lut, gpu_memory_allocated);
+
+    bi_lut->broadcast_lut(streams, gpu_indexes, 0);
+  }
+
+  void release(cudaStream_t const *streams, uint32_t const *gpu_indexes,
+               uint32_t gpu_count) {
+
+    release_radix_ciphertext_async(streams[0], gpu_indexes[0], copy_ct,
+                                   gpu_memory_allocated);
+    delete copy_ct;
+
+    uni_lut->release(streams, gpu_indexes, gpu_count);
+    delete uni_lut;
+
+    bi_lut->release(streams, gpu_indexes, gpu_count);
+    delete bi_lut;
+  }
+};
+
 template <typename Torus> struct int_bitop_buffer {
 
   int_radix_params params;

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

@@ -0,0 +1,46 @@
+#include "ilog2.cuh"
+
+uint64_t scratch_cuda_prepare_count_of_consecutive_bits_buffer_kb_64(
+    void *const *streams, uint32_t const *gpu_indexes, const uint32_t gpu_count,
+    int8_t **mem_ptr, const uint32_t num_radix_blocks, const Direction dir,
+    const BitValue bit_value, const bool allocate_gpu_memory,
+    const uint32_t glwe_dimension, const uint32_t polynomial_size,
+    const uint32_t lwe_dimension, const uint32_t ks_level,
+    const uint32_t ks_base_log, const uint32_t pbs_level,
+    const uint32_t pbs_base_log, const uint32_t grouping_factor,
+    const uint32_t message_modulus, const uint32_t carry_modulus,
+    const PBS_TYPE pbs_type, const bool allocate_ms_array) {
+
+  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,
+                          allocate_ms_array);
+
+  return scratch_cuda_prepare_count_of_consecutive_bits_buffer_kb<uint64_t>(
+      (cudaStream_t *)streams, gpu_indexes, gpu_count,
+      (int_prepare_count_of_consecutive_bits_buffer<uint64_t> **)mem_ptr,
+      num_radix_blocks, params, dir, bit_value, allocate_gpu_memory);
+}
+
+void host_cuda_prepare_count_of_consecutive_bits_buffer_kb_64(
+    void *const *streams, uint32_t const *gpu_indexes, const uint32_t gpu_count,
+    CudaRadixCiphertextFFI *output, CudaRadixCiphertextFFI const *input,
+    int8_t **mem_ptr, void *const *bsks, void *const *ksks,
+    CudaModulusSwitchNoiseReductionKeyFFI const *ms_noise_reduction_key) {
+
+  host_cuda_prepare_count_of_consecutive_bits_buffer_kb<uint64_t>(
+      (cudaStream_t *)streams, gpu_indexes, gpu_count, output, input,
+      (int_prepare_count_of_consecutive_bits_buffer<uint64_t> **)mem_ptr, bsks,
+      (uint64_t **)ksks, ms_noise_reduction_key);
+}
+
+void cleanup_cuda_prepare_count_of_consecutive_bits_buffer_kb_64(
+    void *const *streams, uint32_t const *gpu_indexes, uint32_t gpu_count,
+    int8_t **mem_ptr) {
+
+  auto *buf =
+      (int_prepare_count_of_consecutive_bits_buffer<uint64_t> *)(*mem_ptr);
+
+  buf->release((cudaStream_t *)(streams), gpu_indexes, gpu_count);
+}

backends/tfhe-cuda-backend/cuda/src/integer/ilog2.cuh

@@ -0,0 +1,49 @@
+#ifndef ILOG2_CUH
+#define ILOG2_CUH
+
+#include "device.h"
+#include "integer.cuh"
+#include "integer/integer_utilities.h"
+#include <inttypes.h>
+
+template <typename Torus>
+__host__ uint64_t scratch_cuda_prepare_count_of_consecutive_bits_buffer_kb(
+    cudaStream_t const *streams, uint32_t const *gpu_indexes,
+    uint32_t gpu_count,
+    int_prepare_count_of_consecutive_bits_buffer<Torus> **mem_ptr,
+    uint32_t num_radix_blocks, int_radix_params params, Direction dir,
+    BitValue bit_value, bool allocate_gpu_memory) {
+
+  uint64_t size_tracker = 0;
+  *mem_ptr = new int_prepare_count_of_consecutive_bits_buffer<Torus>(
+      streams, gpu_indexes, gpu_count, dir, bit_value, params, num_radix_blocks,
+      allocate_gpu_memory, &size_tracker);
+
+  return size_tracker;
+}
+
+template <typename Torus>
+__host__ void host_cuda_prepare_count_of_consecutive_bits_buffer_kb(
+    cudaStream_t const *streams, uint32_t const *gpu_indexes,
+    uint32_t gpu_count, CudaRadixCiphertextFFI *output,
+    CudaRadixCiphertextFFI const *input,
+    int_prepare_count_of_consecutive_bits_buffer<Torus> **mem_ptr,
+    void *const *bsks, Torus *const *ksks,
+    CudaModulusSwitchNoiseReductionKeyFFI const *ms_noise_reduction_key) {
+
+  integer_radix_apply_univariate_lookup_table_kb<Torus>(
+      streams, gpu_indexes, gpu_count, (*mem_ptr)->copy_ct, input, bsks, ksks,
+      ms_noise_reduction_key, (*mem_ptr)->uni_lut, output->num_radix_blocks);
+
+  if ((*mem_ptr)->dir == Direction::LEADING) {
+    host_radix_blocks_reverse_inplace<Torus>((cudaStream_t *)streams,
+                                             gpu_indexes, (*mem_ptr)->copy_ct);
+  }
+
+  host_compute_prefix_sum_hillis_steele<Torus>(
+      (cudaStream_t *)(streams), gpu_indexes, gpu_count, output,
+      (*mem_ptr)->copy_ct, (*mem_ptr)->bi_lut, bsks, ksks,
+      ms_noise_reduction_key, output->num_radix_blocks);
+}
+
+#endif

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

@@ -210,6 +210,12 @@ pub const COMPARISON_TYPE_LE: COMPARISON_TYPE = 5;
 pub const COMPARISON_TYPE_MAX: COMPARISON_TYPE = 6;
 pub const COMPARISON_TYPE_MIN: COMPARISON_TYPE = 7;
 pub type COMPARISON_TYPE = ffi::c_uint;
+pub const Direction_TRAILING: Direction = 0;
+pub const Direction_LEADING: Direction = 1;
+pub type Direction = ffi::c_uint;
+pub const BitValue_ZERO: BitValue = 0;
+pub const BitValue_ONE: BitValue = 1;
+pub type BitValue = ffi::c_uint;
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]
 pub struct CudaRadixCiphertextFFI {
@@ -1316,6 +1322,51 @@ unsafe extern "C" {
         mem_ptr_void: *mut *mut i8,
     );
 }
+unsafe extern "C" {
+    pub fn scratch_cuda_prepare_count_of_consecutive_bits_buffer_kb_64(
+        streams: *const *mut ffi::c_void,
+        gpu_indexes: *const u32,
+        gpu_count: u32,
+        mem_ptr: *mut *mut i8,
+        num_radix_blocks: u32,
+        dir: Direction,
+        bit_value: BitValue,
+        allocate_gpu_memory: bool,
+        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,
+        message_modulus: u32,
+        carry_modulus: u32,
+        pbs_type: PBS_TYPE,
+        allocate_ms_array: bool,
+    ) -> u64;
+}
+unsafe extern "C" {
+    pub fn host_cuda_prepare_count_of_consecutive_bits_buffer_kb_64(
+        streams: *const *mut ffi::c_void,
+        gpu_indexes: *const u32,
+        gpu_count: u32,
+        output: *mut CudaRadixCiphertextFFI,
+        input: *const CudaRadixCiphertextFFI,
+        mem_ptr: *mut *mut i8,
+        bsks: *const *mut ffi::c_void,
+        ksks: *const *mut ffi::c_void,
+        ms_noise_reduction_key: *const CudaModulusSwitchNoiseReductionKeyFFI,
+    );
+}
+unsafe extern "C" {
+    pub fn cleanup_cuda_prepare_count_of_consecutive_bits_buffer_kb_64(
+        streams: *const *mut ffi::c_void,
+        gpu_indexes: *const u32,
+        gpu_count: u32,
+        mem_ptr_void: *mut *mut i8,
+    );
+}
 pub const KS_TYPE_BIG_TO_SMALL: KS_TYPE = 0;
 pub const KS_TYPE_SMALL_TO_BIG: KS_TYPE = 1;
 pub type KS_TYPE = ffi::c_uint;