#ifndef CUDA_INTEGER_UTILITIES_H
#define CUDA_INTEGER_UTILITIES_H

#include "integer.h"
#include "integer/radix_ciphertext.cuh"
#include "integer/radix_ciphertext.h"
#include "keyswitch/keyswitch.h"
#include "pbs/programmable_bootstrap.cuh"
#include "utils/helper_multi_gpu.cuh"
#include <cmath>
#include <functional>
#include <queue>

#include <stdio.h>

class NoiseLevel {
public:
  // Constants equivalent to the Rust code
  static const uint64_t NOMINAL = 1;
  static const uint64_t ZERO = 0;
  static const uint64_t UNKNOWN = std::numeric_limits<uint64_t>::max();
};

#ifdef DEBUG
#define CHECK_NOISE_LEVEL(noise_level_expr, msg_mod, carry_mod)                \
  do {                                                                         \
    if ((msg_mod) == 2 && (carry_mod) == 2) {                                  \
      constexpr int max_noise_level = 3;                                       \
      if ((noise_level_expr) > max_noise_level)                                \
        PANIC("Cuda error: noise exceeds maximum authorized value for 1_1 "    \
              "parameters");                                                   \
    } else if ((msg_mod) == 4 && (carry_mod) == 4) {                           \
      constexpr int max_noise_level = 5;                                       \
      if ((noise_level_expr) > max_noise_level)                                \
        PANIC("Cuda error: noise exceeds maximum authorized value for 2_2 "    \
              "parameters");                                                   \
    } else if ((msg_mod) == 8 && (carry_mod) == 8) {                           \
      constexpr int max_noise_level = 9;                                       \
      if ((noise_level_expr) > max_noise_level)                                \
        PANIC("Cuda error: noise exceeds maximum authorized value for 3_3 "    \
              "parameters");                                                   \
    } else if ((msg_mod) == 0 && (carry_mod) == 0) {                           \
      break;                                                                   \
    } else {                                                                   \
      PANIC("Invalid message modulus or carry modulus")                        \
    }                                                                          \
  } while (0)
#else
#define CHECK_NOISE_LEVEL(noise_level_expr, message_modulus, carry_modulus)    \
  do {                                                                         \
  } while (0)
#endif

template <typename Torus>
__global__ void radix_blocks_rotate_right(Torus *dst, Torus *src,
                                          uint32_t value, uint32_t blocks_count,
                                          uint32_t lwe_size);
void generate_ids_update_degrees(uint64_t *terms_degree, size_t *h_lwe_idx_in,
                                 size_t *h_lwe_idx_out,
                                 int32_t *h_smart_copy_in,
                                 int32_t *h_smart_copy_out, size_t ch_amount,
                                 uint32_t num_radix, uint32_t num_blocks,
                                 size_t chunk_size, size_t message_max,
                                 size_t &total_count, size_t &message_count,
                                 size_t &carry_count, size_t &sm_copy_count);
/*
 *  generate bivariate accumulator (lut) for device pointer
 *    stream - cuda stream
 *    acc_bivariate - device pointer for bivariate accumulator
 *    ...
 *    f - wrapping function with two Torus inputs
 */
template <typename Torus>
void generate_device_accumulator_bivariate(
    cudaStream_t stream, uint32_t gpu_index, Torus *acc_bivariate,
    uint64_t *degree, uint64_t *max_degree, uint32_t glwe_dimension,
    uint32_t polynomial_size, uint32_t message_modulus, uint32_t carry_modulus,
    std::function<Torus(Torus, Torus)> f, bool gpu_memory_allocated);

template <typename Torus>
void generate_device_accumulator_bivariate_with_factor(
    cudaStream_t stream, uint32_t gpu_index, Torus *acc_bivariate,
    uint64_t *degree, uint64_t *max_degree, uint32_t glwe_dimension,
    uint32_t polynomial_size, uint32_t message_modulus, uint32_t carry_modulus,
    std::function<Torus(Torus, Torus)> f, int factor,
    bool gpu_memory_allocated);

template <typename Torus>
void generate_device_accumulator_with_encoding(
    cudaStream_t stream, uint32_t gpu_index, Torus *acc, uint64_t *degree,
    uint64_t *max_degree, uint32_t glwe_dimension, uint32_t polynomial_size,
    uint32_t input_message_modulus, uint32_t input_carry_modulus,
    uint32_t output_message_modulus, uint32_t output_carry_modulus,
    std::function<Torus(Torus)> f, bool gpu_memory_allocated);

template <typename Torus>
void generate_device_accumulator_no_encoding(
    cudaStream_t stream, uint32_t gpu_index, Torus *acc, uint64_t *degree,
    uint32_t message_modulus, uint32_t carry_modulus, uint32_t glwe_dimension,
    uint32_t polynomial_size, std::function<Torus(uint32_t)> f,
    bool gpu_memory_allocated);

/*
 *  generate univariate accumulator (lut) for device pointer
 *    stream - cuda stream
 *    acc - device pointer for univariate accumulator
 *    ...
 *    f - evaluating function with one Torus input
 */
template <typename Torus>
void generate_device_accumulator(
    cudaStream_t stream, uint32_t gpu_index, Torus *acc, uint64_t *degree,
    uint64_t *max_degree, uint32_t glwe_dimension, uint32_t polynomial_size,
    uint32_t message_modulus, uint32_t carry_modulus,
    std::function<Torus(Torus)> f, bool gpu_memory_allocated);

template <typename Torus>
void generate_many_lut_device_accumulator(
    cudaStream_t stream, uint32_t gpu_index, Torus *acc, uint64_t *degrees,
    uint64_t *max_degree, uint32_t glwe_dimension, uint32_t polynomial_size,
    uint32_t message_modulus, uint32_t carry_modulus,
    std::vector<std::function<Torus(Torus)>> &f, bool gpu_memory_allocated);

struct radix_columns {
  std::vector<uint32_t> columns_counter;
  uint32_t num_blocks;
  uint32_t num_radix_in_vec;
  uint32_t chunk_size;
  radix_columns(const uint64_t *const input_degrees, uint32_t num_blocks,
                uint32_t num_radix_in_vec, uint32_t chunk_size,
                bool &needs_processing)
      : num_blocks(num_blocks), num_radix_in_vec(num_radix_in_vec),
        chunk_size(chunk_size) {
    needs_processing = false;
    columns_counter.resize(num_blocks, 0);
    for (uint32_t i = 0; i < num_radix_in_vec; ++i) {
      for (uint32_t j = 0; j < num_blocks; ++j) {
        if (input_degrees[i * num_blocks + j])
          columns_counter[j] += 1;
      }
    }

    for (uint32_t i = 0; i < num_blocks; ++i) {
      if (columns_counter[i] > chunk_size) {
        needs_processing = true;
        break;
      }
    }
  }

  void next_accumulation(uint32_t &total_ciphertexts,
                         uint32_t &message_ciphertexts,
                         bool &needs_processing) {
    message_ciphertexts = 0;
    total_ciphertexts = 0;
    needs_processing = false;
    for (int i = num_blocks - 1; i > 0; --i) {
      uint32_t cur_count = columns_counter[i];
      uint32_t prev_count = columns_counter[i - 1];
      uint32_t new_count = 0;

      // accumulated_blocks from current columns
      new_count += cur_count / chunk_size;
      // all accumulated message blocks needs pbs
      message_ciphertexts += new_count;
      // carry blocks from previous columns
      new_count += prev_count / chunk_size;
      // both carry and message blocks that needs pbs
      total_ciphertexts += new_count;
      // now add remaining non accumulated blocks that does not require pbs
      new_count += cur_count % chunk_size;

      columns_counter[i] = new_count;

      if (new_count > chunk_size)
        needs_processing = true;
    }

    // now do it for 0th block
    uint32_t new_count = columns_counter[0] / chunk_size;
    message_ciphertexts += new_count;
    total_ciphertexts += new_count;
    new_count += columns_counter[0] % chunk_size;
    columns_counter[0] = new_count;

    if (new_count > chunk_size) {
      needs_processing = true;
    }
  }
};

inline void calculate_final_degrees(uint64_t *const out_degrees,
                                    const uint64_t *const input_degrees,
                                    uint32_t num_blocks,
                                    uint32_t num_radix_in_vec,
                                    uint32_t chunk_size,
                                    uint64_t message_modulus) {

  auto get_degree = [message_modulus](uint64_t degree) -> uint64_t {
    return std::min(message_modulus - 1, degree);
  };
  std::vector<std::queue<uint64_t>> columns(num_blocks);
  for (uint32_t i = 0; i < num_radix_in_vec; ++i) {
    for (uint32_t j = 0; j < num_blocks; ++j) {
      if (input_degrees[i * num_blocks + j])
        columns[j].push(input_degrees[i * num_blocks + j]);
    }
  }

  for (uint32_t i = 0; i < num_blocks; ++i) {
    auto &col = columns[i];
    while (col.size() > 1) {
      uint32_t cur_degree = 0;
      uint32_t mn = std::min(chunk_size, (uint32_t)col.size());
      for (int j = 0; j < mn; ++j) {
        cur_degree += col.front();
        col.pop();
      }
      const uint64_t new_degree = get_degree(cur_degree);
      col.push(new_degree);
      if ((i + 1) < num_blocks) {
        columns[i + 1].push(new_degree);
      }
    }
  }

  for (int i = 0; i < num_blocks; i++) {
    out_degrees[i] = (columns[i].empty()) ? 0 : columns[i].front();
  }
}

struct int_radix_params {
  PBS_TYPE pbs_type;
  uint32_t glwe_dimension;
  uint32_t polynomial_size;
  uint32_t big_lwe_dimension;
  uint32_t small_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 message_modulus;
  uint32_t carry_modulus;
  PBS_MS_REDUCTION_T noise_reduction_type;

  int_radix_params(PBS_TYPE pbs_type, uint32_t glwe_dimension,
                   uint32_t polynomial_size, uint32_t big_lwe_dimension,
                   uint32_t small_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 message_modulus, uint32_t carry_modulus,
                   PBS_MS_REDUCTION_T noise_reduction_type)

      : pbs_type(pbs_type), glwe_dimension(glwe_dimension),
        polynomial_size(polynomial_size), big_lwe_dimension(big_lwe_dimension),
        small_lwe_dimension(small_lwe_dimension), ks_level(ks_level),
        ks_base_log(ks_base_log), pbs_level(pbs_level),
        pbs_base_log(pbs_base_log), grouping_factor(grouping_factor),
        message_modulus(message_modulus), carry_modulus(carry_modulus),
        noise_reduction_type(noise_reduction_type){};

  int_radix_params() = default;

  void print() {
    printf("pbs_type: %u, glwe_dimension: %u, "
           "polynomial_size: %u, "
           "big_lwe_dimension: %u, "
           "small_lwe_dimension: %u, ks_level: %u, ks_base_log: %u, pbs_level: "
           "%u, pbs_base_log: "
           "%u, grouping_factor: %u, message_modulus: %u, carry_modulus: %u\n",
           pbs_type, glwe_dimension, polynomial_size, big_lwe_dimension,
           small_lwe_dimension, ks_level, ks_base_log, pbs_level, pbs_base_log,
           grouping_factor, message_modulus, carry_modulus);
  };
};

// Store things needed to apply LUTs
template <typename Torus> struct int_radix_lut {
  int_radix_params params;
  uint32_t num_blocks;
  uint32_t num_luts;
  uint32_t num_many_lut = 1;
  // Tracks the degree of each LUT and the max degree on CPU
  // The max degree is (message_modulus * carry_modulus - 1) except for many lut
  // for which it's different
  uint64_t *degrees;
  uint64_t *max_degrees;

  CudaStreams active_streams;
  bool mem_reuse = false;

  // There will be one buffer on each GPU in multi-GPU computations
  // (same for tmp lwe arrays)
  std::vector<int8_t *> buffer;

  // These arrays will reside on all GPUs
  // lut could actually be allocated & initialized GPU per GPU but this is not
  // done at the moment
  std::vector<Torus *> lut_vec;
  std::vector<Torus *> lut_indexes_vec;
  Torus *h_lut_indexes;
  // All tmp lwe arrays and index arrays for lwe contain the total
  // amount of blocks to be computed on, there is no split between GPUs
  // for the moment
  Torus *lwe_indexes_in;
  Torus *lwe_indexes_out;
  Torus *h_lwe_indexes_in;
  Torus *h_lwe_indexes_out;
  // Enable optimizations if lwe_indexes_(in/out) are trivial
  bool using_trivial_lwe_indexes = true;
  // lwe_trivial_indexes is the intermediary index we need in case
  // lwe_indexes_in != lwe_indexes_out
  Torus *lwe_trivial_indexes;
  CudaRadixCiphertextFFI *tmp_lwe_before_ks;

  /// For multi GPU execution we create vectors of pointers for inputs and
  /// outputs
  std::vector<Torus *> lwe_array_in_vec;
  std::vector<Torus *> lwe_after_ks_vec;
  std::vector<Torus *> lwe_after_pbs_vec;
  std::vector<Torus *> lwe_trivial_indexes_vec;
  std::vector<Torus *> lwe_aligned_vec;

  bool gpu_memory_allocated;

  cudaEvent_t event_scatter_in;
  cudaEvent_t *event_scatter_out;
  cudaEvent_t event_broadcast;

  int_radix_lut(CudaStreams streams, int_radix_params params, uint32_t num_luts,
                uint32_t num_radix_blocks, bool allocate_gpu_memory,
                uint64_t &size_tracker) {

    this->params = params;
    this->num_blocks = num_radix_blocks;
    this->num_luts = num_luts;
    gpu_memory_allocated = allocate_gpu_memory;
    uint64_t lut_indexes_size = num_radix_blocks * sizeof(Torus);
    uint64_t lut_buffer_size =
        (params.glwe_dimension + 1) * params.polynomial_size * sizeof(Torus);

    active_streams = streams.active_gpu_subset(num_radix_blocks);

    ///////////////
    for (uint i = 0; i < active_streams.count(); i++) {
      cuda_set_device(streams.gpu_index(i));
      int8_t *gpu_pbs_buffer;
      auto num_blocks_on_gpu = std::max(
          THRESHOLD_MULTI_GPU,
          get_num_inputs_on_gpu(num_radix_blocks, i, active_streams.count()));

      uint64_t size = 0;
      execute_scratch_pbs<Torus>(
          streams.stream(i), streams.gpu_index(i), &gpu_pbs_buffer,
          params.glwe_dimension, params.small_lwe_dimension,
          params.polynomial_size, params.pbs_level, params.grouping_factor,
          num_blocks_on_gpu, params.pbs_type, allocate_gpu_memory,
          params.noise_reduction_type, size);
      if (i == 0) {
        size_tracker += size;
      }
      buffer.push_back(gpu_pbs_buffer);
    }

    // We create the events only if we have multiple GPUs
    if (active_streams.count() > 1) {
      event_scatter_in = cuda_create_event(streams.gpu_index(0));
      event_broadcast = cuda_create_event(streams.gpu_index(0));

      event_scatter_out =
          (cudaEvent_t *)malloc(active_streams.count() * sizeof(cudaEvent_t));
      for (int i = 0; i < active_streams.count(); i++) {
        event_scatter_out[i] = cuda_create_event(active_streams.gpu_index(i));
      }
    }

    // Allocate LUT
    // LUT is used as a trivial encryption and must be initialized outside
    // this constructor
    for (uint i = 0; i < active_streams.count(); i++) {
      auto lut = (Torus *)cuda_malloc_with_size_tracking_async(
          num_luts * lut_buffer_size, streams.stream(i), streams.gpu_index(i),
          size_tracker, allocate_gpu_memory);
      auto lut_indexes = (Torus *)cuda_malloc_with_size_tracking_async(
          lut_indexes_size, streams.stream(i), streams.gpu_index(i),
          size_tracker, allocate_gpu_memory);
      // lut_indexes is initialized to 0 by default
      // if a different behavior is wanted, it should be rewritten later
      cuda_memset_with_size_tracking_async(
          lut_indexes, 0, lut_indexes_size, streams.stream(i),
          streams.gpu_index(i), allocate_gpu_memory);

      lut_vec.push_back(lut);
      lut_indexes_vec.push_back(lut_indexes);
    }

    // lwe_(input/output)_indexes are initialized to range(num_radix_blocks)
    // by default
    lwe_indexes_in = (Torus *)cuda_malloc_with_size_tracking_async(
        num_radix_blocks * sizeof(Torus), streams.stream(0),
        streams.gpu_index(0), size_tracker, allocate_gpu_memory);
    lwe_indexes_out = (Torus *)cuda_malloc_with_size_tracking_async(
        num_radix_blocks * sizeof(Torus), streams.stream(0),
        streams.gpu_index(0), size_tracker, allocate_gpu_memory);
    lwe_trivial_indexes = (Torus *)cuda_malloc_with_size_tracking_async(
        num_radix_blocks * sizeof(Torus), streams.stream(0),
        streams.gpu_index(0), size_tracker, allocate_gpu_memory);

    h_lwe_indexes_in = (Torus *)malloc(num_radix_blocks * sizeof(Torus));
    h_lwe_indexes_out = (Torus *)malloc(num_radix_blocks * sizeof(Torus));
    for (int i = 0; i < num_radix_blocks; i++)
      h_lwe_indexes_in[i] = i;

    cuda_memcpy_with_size_tracking_async_to_gpu(
        lwe_indexes_in, h_lwe_indexes_in, num_radix_blocks * sizeof(Torus),
        streams.stream(0), streams.gpu_index(0), allocate_gpu_memory);
    cuda_memcpy_with_size_tracking_async_to_gpu(
        lwe_indexes_out, h_lwe_indexes_in, num_radix_blocks * sizeof(Torus),
        streams.stream(0), streams.gpu_index(0), allocate_gpu_memory);
    cuda_memcpy_with_size_tracking_async_to_gpu(
        lwe_trivial_indexes, h_lwe_indexes_in, num_radix_blocks * sizeof(Torus),
        streams.stream(0), streams.gpu_index(0), allocate_gpu_memory);
    memcpy(h_lwe_indexes_out, h_lwe_indexes_in,
           num_radix_blocks * sizeof(Torus));

    /// With multiple GPUs we allocate arrays to be pushed to the vectors and
    /// copy data on each GPU then when we gather data to GPU 0 we can copy
    /// back to the original indexing
    multi_gpu_alloc_lwe_async(streams, lwe_array_in_vec, num_radix_blocks,
                              params.big_lwe_dimension + 1, size_tracker,
                              allocate_gpu_memory);
    multi_gpu_alloc_lwe_async(streams, lwe_after_ks_vec, num_radix_blocks,
                              params.small_lwe_dimension + 1, size_tracker,
                              allocate_gpu_memory);
    multi_gpu_alloc_lwe_async(streams, lwe_after_pbs_vec, num_radix_blocks,
                              params.big_lwe_dimension + 1, size_tracker,
                              allocate_gpu_memory);
    multi_gpu_alloc_array_async(streams, lwe_trivial_indexes_vec,
                                num_radix_blocks, size_tracker,
                                allocate_gpu_memory);
    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
    multi_gpu_copy_array_async(active_streams, lwe_trivial_indexes_vec,
                               lwe_trivial_indexes, num_radix_blocks,
                               allocate_gpu_memory);

    // Keyswitch
    tmp_lwe_before_ks = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_lwe_before_ks,
        num_radix_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    h_lut_indexes = (Torus *)(calloc(num_radix_blocks, sizeof(Torus)));
    degrees = (uint64_t *)malloc(num_luts * sizeof(uint64_t));
    max_degrees = (uint64_t *)malloc(num_luts * sizeof(uint64_t));
  }

  // constructor to reuse memory
  int_radix_lut(CudaStreams streams, int_radix_params params, uint32_t num_luts,
                uint32_t num_radix_blocks, int_radix_lut *base_lut_object,
                bool allocate_gpu_memory, uint64_t &size_tracker) {

    this->params = params;
    this->num_blocks = num_radix_blocks;
    this->num_luts = num_luts;
    gpu_memory_allocated = allocate_gpu_memory;
    uint64_t lut_indexes_size = num_radix_blocks * sizeof(Torus);
    uint64_t lut_buffer_size =
        (params.glwe_dimension + 1) * params.polynomial_size * sizeof(Torus);

    // base lut object should have bigger or equal memory than current one
    if (num_radix_blocks > base_lut_object->num_blocks)
      PANIC("Cuda error: lut does not have enough blocks")
    // pbs
    buffer = base_lut_object->buffer;
    // Keyswitch
    tmp_lwe_before_ks = base_lut_object->tmp_lwe_before_ks;

    /// With multiple GPUs we allocate arrays to be pushed to the vectors and
    /// copy data on each GPU then when we gather data to GPU 0 we can copy back
    /// to the original indexing
    lwe_array_in_vec = base_lut_object->lwe_array_in_vec;
    lwe_after_ks_vec = base_lut_object->lwe_after_ks_vec;
    lwe_after_pbs_vec = base_lut_object->lwe_after_pbs_vec;
    lwe_trivial_indexes_vec = base_lut_object->lwe_trivial_indexes_vec;

    mem_reuse = true;

    // Allocate LUT
    // LUT is used as a trivial encryption and must be initialized outside
    // this constructor
    active_streams = streams.active_gpu_subset(num_radix_blocks);
    for (uint i = 0; i < active_streams.count(); i++) {
      auto lut = (Torus *)cuda_malloc_with_size_tracking_async(
          num_luts * lut_buffer_size, streams.stream(i), streams.gpu_index(i),
          size_tracker, allocate_gpu_memory);
      auto lut_indexes = (Torus *)cuda_malloc_with_size_tracking_async(
          lut_indexes_size, streams.stream(i), streams.gpu_index(i),
          size_tracker, allocate_gpu_memory);
      // lut_indexes is initialized to 0 by default
      // if a different behavior is wanted, it should be rewritten later
      cuda_memset_with_size_tracking_async(
          lut_indexes, 0, lut_indexes_size, streams.stream(i),
          streams.gpu_index(i), allocate_gpu_memory);
      lut_vec.push_back(lut);
      lut_indexes_vec.push_back(lut_indexes);
    }

    // lwe_(input/output)_indexes are initialized to range(num_radix_blocks)
    // by default
    lwe_indexes_in = (Torus *)cuda_malloc_with_size_tracking_async(
        num_radix_blocks * sizeof(Torus), streams.stream(0),
        streams.gpu_index(0), size_tracker, allocate_gpu_memory);
    lwe_indexes_out = (Torus *)cuda_malloc_with_size_tracking_async(
        num_radix_blocks * sizeof(Torus), streams.stream(0),
        streams.gpu_index(0), size_tracker, allocate_gpu_memory);
    lwe_trivial_indexes = (Torus *)cuda_malloc_with_size_tracking_async(
        num_radix_blocks * sizeof(Torus), streams.stream(0),
        streams.gpu_index(0), size_tracker, allocate_gpu_memory);

    h_lwe_indexes_in = (Torus *)malloc(num_radix_blocks * sizeof(Torus));
    h_lwe_indexes_out = (Torus *)malloc(num_radix_blocks * sizeof(Torus));

    for (int i = 0; i < num_radix_blocks; i++)
      h_lwe_indexes_in[i] = i;

    cuda_memcpy_with_size_tracking_async_to_gpu(
        lwe_indexes_in, h_lwe_indexes_in, num_radix_blocks * sizeof(Torus),
        streams.stream(0), streams.gpu_index(0), allocate_gpu_memory);
    cuda_memcpy_with_size_tracking_async_to_gpu(
        lwe_indexes_out, h_lwe_indexes_in, num_radix_blocks * sizeof(Torus),
        streams.stream(0), streams.gpu_index(0), allocate_gpu_memory);
    cuda_memcpy_with_size_tracking_async_to_gpu(
        lwe_trivial_indexes, h_lwe_indexes_in, num_radix_blocks * sizeof(Torus),
        streams.stream(0), streams.gpu_index(0), allocate_gpu_memory);
    memcpy(h_lwe_indexes_out, h_lwe_indexes_in,
           num_radix_blocks * sizeof(Torus));
    h_lut_indexes = (Torus *)(calloc(num_radix_blocks, sizeof(Torus)));
    degrees = (uint64_t *)malloc(num_luts * sizeof(uint64_t));
    max_degrees = (uint64_t *)malloc(num_luts * sizeof(uint64_t));
  }

  // Construction for many luts
  int_radix_lut(CudaStreams streams, int_radix_params params, uint32_t num_luts,
                uint32_t num_radix_blocks, uint32_t num_many_lut,
                bool allocate_gpu_memory, uint64_t &size_tracker) {

    this->num_many_lut = num_many_lut;
    this->params = params;
    this->num_blocks = num_radix_blocks;
    this->num_luts = num_luts;
    gpu_memory_allocated = allocate_gpu_memory;
    uint64_t lut_indexes_size = num_radix_blocks * sizeof(Torus);
    uint64_t lut_buffer_size =
        (params.glwe_dimension + 1) * params.polynomial_size * sizeof(Torus);

    ///////////////
    active_streams = streams.active_gpu_subset(num_radix_blocks);
    for (uint i = 0; i < active_streams.count(); i++) {
      cuda_set_device(streams.gpu_index(i));
      int8_t *gpu_pbs_buffer;
      auto num_blocks_on_gpu = std::max(
          THRESHOLD_MULTI_GPU,
          get_num_inputs_on_gpu(num_radix_blocks, i, active_streams.count()));

      uint64_t size = 0;
      execute_scratch_pbs<Torus>(
          streams.stream(i), streams.gpu_index(i), &gpu_pbs_buffer,
          params.glwe_dimension, params.small_lwe_dimension,
          params.polynomial_size, params.pbs_level, params.grouping_factor,
          num_blocks_on_gpu, params.pbs_type, allocate_gpu_memory,
          params.noise_reduction_type, size);
      if (i == 0) {
        size_tracker += size;
      }
      buffer.push_back(gpu_pbs_buffer);
    }
    // We create the events only if we have multiple GPUs
    if (active_streams.count() > 1) {
      event_scatter_in = cuda_create_event(streams.gpu_index(0));
      event_broadcast = cuda_create_event(streams.gpu_index(0));

      event_scatter_out =
          (cudaEvent_t *)malloc(active_streams.count() * sizeof(cudaEvent_t));
      for (int i = 0; i < active_streams.count(); i++) {
        event_scatter_out[i] = cuda_create_event(active_streams.gpu_index(i));
      }
    }
    // Allocate LUT
    // LUT is used as a trivial encryption and must be initialized outside
    // this constructor
    for (uint i = 0; i < active_streams.count(); i++) {
      auto lut = (Torus *)cuda_malloc_with_size_tracking_async(
          num_luts * lut_buffer_size, streams.stream(i), streams.gpu_index(i),
          size_tracker, allocate_gpu_memory);
      auto lut_indexes = (Torus *)cuda_malloc_with_size_tracking_async(
          lut_indexes_size, streams.stream(i), streams.gpu_index(i),
          size_tracker, allocate_gpu_memory);
      // lut_indexes is initialized to 0 by default
      // if a different behavior is wanted, it should be rewritten later
      cuda_memset_with_size_tracking_async(
          lut_indexes, 0, lut_indexes_size, streams.stream(i),
          streams.gpu_index(i), allocate_gpu_memory);
      lut_vec.push_back(lut);
      lut_indexes_vec.push_back(lut_indexes);
    }

    // lwe_(input/output)_indexes are initialized to range(num_radix_blocks)
    // by default
    lwe_indexes_in = (Torus *)cuda_malloc_with_size_tracking_async(
        num_radix_blocks * sizeof(Torus), streams.stream(0),
        streams.gpu_index(0), size_tracker, allocate_gpu_memory);
    lwe_indexes_out = (Torus *)cuda_malloc_with_size_tracking_async(
        num_radix_blocks * sizeof(Torus), streams.stream(0),
        streams.gpu_index(0), size_tracker, allocate_gpu_memory);
    lwe_trivial_indexes = (Torus *)cuda_malloc_with_size_tracking_async(
        num_radix_blocks * sizeof(Torus), streams.stream(0),
        streams.gpu_index(0), size_tracker, allocate_gpu_memory);

    h_lwe_indexes_in = (Torus *)malloc(num_radix_blocks * sizeof(Torus));
    h_lwe_indexes_out = (Torus *)malloc(num_radix_blocks * sizeof(Torus));

    for (int i = 0; i < num_radix_blocks; i++)
      h_lwe_indexes_in[i] = i;

    cuda_memcpy_with_size_tracking_async_to_gpu(
        lwe_indexes_in, h_lwe_indexes_in, num_radix_blocks * sizeof(Torus),
        streams.stream(0), streams.gpu_index(0), allocate_gpu_memory);
    cuda_memcpy_with_size_tracking_async_to_gpu(
        lwe_indexes_out, h_lwe_indexes_in, num_radix_blocks * sizeof(Torus),
        streams.stream(0), streams.gpu_index(0), allocate_gpu_memory);
    cuda_memcpy_with_size_tracking_async_to_gpu(
        lwe_trivial_indexes, h_lwe_indexes_in, num_radix_blocks * sizeof(Torus),
        streams.stream(0), streams.gpu_index(0), allocate_gpu_memory);
    memcpy(h_lwe_indexes_out, h_lwe_indexes_in,
           num_radix_blocks * sizeof(Torus));

    /// With multiple GPUs we allocate arrays to be pushed to the vectors and
    /// copy data on each GPU then when we gather data to GPU 0 we can copy
    /// back to the original indexing
    multi_gpu_alloc_lwe_async(active_streams, lwe_array_in_vec,
                              num_radix_blocks, params.big_lwe_dimension + 1,
                              size_tracker, allocate_gpu_memory);
    multi_gpu_alloc_lwe_async(active_streams, lwe_after_ks_vec,
                              num_radix_blocks, params.small_lwe_dimension + 1,
                              size_tracker, allocate_gpu_memory);
    multi_gpu_alloc_lwe_many_lut_output_async(
        active_streams, lwe_after_pbs_vec, num_radix_blocks, num_many_lut,
        params.big_lwe_dimension + 1, size_tracker, allocate_gpu_memory);
    multi_gpu_alloc_array_async(active_streams, lwe_trivial_indexes_vec,
                                num_radix_blocks, size_tracker,
                                allocate_gpu_memory);
    multi_gpu_copy_array_from_cpu_async(active_streams, lwe_trivial_indexes_vec,
                                        h_lwe_indexes_in, num_radix_blocks,
                                        allocate_gpu_memory);
    // Keyswitch
    tmp_lwe_before_ks = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_lwe_before_ks,
        num_radix_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    h_lut_indexes = (Torus *)(calloc(num_radix_blocks, sizeof(Torus)));
    degrees = (uint64_t *)malloc(num_many_lut * num_luts * sizeof(uint64_t));
    max_degrees = (uint64_t *)malloc(num_luts * sizeof(uint64_t));
  }

  // Return a pointer to idx-ith lut at gpu_index's global memory
  Torus *get_lut(uint32_t gpu_index, size_t idx) {
    if (!gpu_memory_allocated)
      return nullptr;
    auto lut = lut_vec[gpu_index];
    size_t lut_size = (params.glwe_dimension + 1) * params.polynomial_size;

    if (lut == nullptr)
      PANIC("Cuda error: invalid lut pointer")
    return &lut[idx * lut_size];
  }

  // Return a pointer to idx-ith degree
  uint64_t *get_degree(size_t idx) { return &degrees[num_many_lut * idx]; }

  // Return a pointer to idx-ith max degree
  uint64_t *get_max_degree(size_t idx) { return &max_degrees[idx]; }

  // Return a pointer to idx-ith lut indexes at gpu_index's global memory
  Torus *get_lut_indexes(uint32_t gpu_index, size_t ind) {
    if (!gpu_memory_allocated)
      return nullptr;
    auto lut_indexes = lut_indexes_vec[gpu_index];
    return &lut_indexes[ind];
  }

  // If this function is called we assume the lwe_indexes_(in/out) are not the
  // trivial anymore and thus we disable optimizations
  void set_lwe_indexes(cudaStream_t stream, uint32_t gpu_index,
                       Torus *h_indexes_in, Torus *h_indexes_out) {

    memcpy(h_lwe_indexes_in, h_indexes_in, num_blocks * sizeof(Torus));
    memcpy(h_lwe_indexes_out, h_indexes_out, num_blocks * sizeof(Torus));

    cuda_memcpy_with_size_tracking_async_to_gpu(
        lwe_indexes_in, h_lwe_indexes_in, num_blocks * sizeof(Torus), stream,
        gpu_index, gpu_memory_allocated);
    cuda_memcpy_with_size_tracking_async_to_gpu(
        lwe_indexes_out, h_lwe_indexes_out, num_blocks * sizeof(Torus), stream,
        gpu_index, gpu_memory_allocated);

    using_trivial_lwe_indexes = false;
  }

  // Broadcast luts from device gpu_indexes[0] to all active gpus
  void broadcast_lut(CudaStreams new_active_streams,
                     bool broadcast_lut_values = true) {
    // We only do broadcast if there are more than 1 active GPU
    if (new_active_streams.count() > 1) {
      int active_device = cuda_get_device();

      uint64_t lut_size = (params.glwe_dimension + 1) * params.polynomial_size;

      auto src_lut = lut_vec[0];
      auto src_lut_indexes = lut_indexes_vec[0];
      if (active_streams.count() > 1)
        cuda_event_record(event_broadcast, new_active_streams.stream(0),
                          new_active_streams.gpu_index(0));
      for (uint i = 0; i < new_active_streams.count(); i++) {
        PANIC_IF_FALSE(
            new_active_streams.gpu_index(i) == active_streams.gpu_index(i),
            "Broadcasting LUTs can only be done to the LUT streams or to new "
            "streams that reside on the same GPUs as the source LUTs");

#ifndef DEBUG_FAKE_MULTI_GPU
        if (new_active_streams.gpu_index(i) == new_active_streams.gpu_index(0))
          continue;
#endif

        cuda_stream_wait_event(new_active_streams.stream(i), event_broadcast,
                               new_active_streams.gpu_index(i));
        if (broadcast_lut_values) {
          auto dst_lut = lut_vec[i];
          cuda_memcpy_with_size_tracking_async_gpu_to_gpu(
              dst_lut, src_lut, num_luts * lut_size * sizeof(Torus),
              new_active_streams.stream(i), new_active_streams.gpu_index(i),
              gpu_memory_allocated);
        }
        auto dst_lut_indexes = lut_indexes_vec[i];
        cuda_memcpy_with_size_tracking_async_gpu_to_gpu(
            dst_lut_indexes, src_lut_indexes, num_blocks * sizeof(Torus),
            new_active_streams.stream(i), new_active_streams.gpu_index(i),
            gpu_memory_allocated);
      }
      // Ensure the device set at the end of this method is the same as it was
      // set at the beginning
      cuda_set_device(active_device);
    }
  }

  void allocate_lwe_vector_for_non_trivial_indexes(
      CudaStreams streams, uint64_t max_num_radix_blocks,
      uint64_t &size_tracker, bool allocate_gpu_memory) {
    // We need to create the auxiliary array only in GPU 0
    if (active_streams.count() > 1) {
      lwe_aligned_vec.resize(active_streams.count());
      for (uint i = 0; i < active_streams.count(); i++) {
        uint64_t size_tracker_on_array_i = 0;
        auto inputs_on_gpu = std::max(
            THRESHOLD_MULTI_GPU, get_num_inputs_on_gpu(max_num_radix_blocks, i,
                                                       active_streams.count()));
        Torus *d_array = (Torus *)cuda_malloc_with_size_tracking_async(
            inputs_on_gpu * (params.big_lwe_dimension + 1) * sizeof(Torus),
            streams.stream(0), streams.gpu_index(0), size_tracker_on_array_i,
            allocate_gpu_memory);
        lwe_aligned_vec[i] = d_array;
        size_tracker += size_tracker_on_array_i;
      }
    }
  }

  void release(CudaStreams streams) {
    PANIC_IF_FALSE(lut_indexes_vec.size() == lut_vec.size(),
                   "Lut vec and Lut vec indexes must have the same size");
    for (uint i = 0; i < lut_vec.size(); i++) {
      cuda_drop_with_size_tracking_async(lut_vec[i], streams.stream(i),
                                         streams.gpu_index(i),
                                         gpu_memory_allocated);
      cuda_drop_with_size_tracking_async(lut_indexes_vec[i], streams.stream(i),
                                         streams.gpu_index(i),
                                         gpu_memory_allocated);
    }
    cuda_drop_with_size_tracking_async(lwe_indexes_in, streams.stream(0),
                                       streams.gpu_index(0),
                                       gpu_memory_allocated);
    cuda_drop_with_size_tracking_async(lwe_indexes_out, streams.stream(0),
                                       streams.gpu_index(0),
                                       gpu_memory_allocated);
    cuda_drop_with_size_tracking_async(lwe_trivial_indexes, streams.stream(0),
                                       streams.gpu_index(0),
                                       gpu_memory_allocated);

    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
    lut_vec.clear();
    lut_indexes_vec.clear();
    free(h_lwe_indexes_in);
    free(h_lwe_indexes_out);

    if (!mem_reuse) {
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     tmp_lwe_before_ks, gpu_memory_allocated);
      for (int i = 0; i < buffer.size(); i++) {
        switch (params.pbs_type) {
        case MULTI_BIT:
          cleanup_cuda_multi_bit_programmable_bootstrap(
              streams.stream(i), streams.gpu_index(i), &buffer[i]);
          break;
        case CLASSICAL:
          cleanup_cuda_programmable_bootstrap(streams.stream(i),
                                              streams.gpu_index(i), &buffer[i]);
          break;
        default:
          PANIC("Cuda error (PBS): unknown PBS type. ")
        }
        cuda_synchronize_stream(streams.stream(i), streams.gpu_index(i));
      }
      delete tmp_lwe_before_ks;
      buffer.clear();

      if (gpu_memory_allocated) {
        multi_gpu_release_async(streams, lwe_array_in_vec);
        multi_gpu_release_async(streams, lwe_after_ks_vec);
        multi_gpu_release_async(streams, lwe_after_pbs_vec);
        multi_gpu_release_async(streams, lwe_trivial_indexes_vec);
        streams.synchronize();
      }
      lwe_array_in_vec.clear();
      lwe_after_ks_vec.clear();
      lwe_after_pbs_vec.clear();
      lwe_trivial_indexes_vec.clear();
      if (active_streams.count() > 1) {
        active_streams.synchronize();
        for (uint i = 0; i < active_streams.count(); i++) {
          cuda_event_destroy(event_scatter_out[i], active_streams.gpu_index(i));
        }
        cuda_event_destroy(event_scatter_in, active_streams.gpu_index(0));
        cuda_event_destroy(event_broadcast, active_streams.gpu_index(0));
        free(event_scatter_out);
      }
      if (lwe_aligned_vec.size() > 0) {
        for (uint i = 0; i < active_streams.count(); i++) {
          cuda_drop_with_size_tracking_async(
              lwe_aligned_vec[i], streams.stream(0), streams.gpu_index(0),
              gpu_memory_allocated);
        }
        lwe_aligned_vec.clear();
      }
    }
    free(h_lut_indexes);
    free(degrees);
    free(max_degrees);
  }
};

template <typename InputTorus> struct int_noise_squashing_lut {

  int_radix_params params;
  uint32_t input_glwe_dimension;
  uint32_t input_polynomial_size;
  uint32_t input_big_lwe_dimension;
  uint32_t num_blocks;
  // Tracks the degree of each LUT and the max degree on CPU
  // The max degree is (message_modulus * carry_modulus - 1) except for many lut
  // for which it's different
  uint64_t *degrees;
  uint64_t *max_degrees;

  CudaStreams active_streams;

  // There will be one buffer on each GPU in multi-GPU computations
  // (same for tmp lwe arrays)
  std::vector<int8_t *> pbs_buffer;

  std::vector<__uint128_t *> lut_vec;

  CudaRadixCiphertextFFI *tmp_lwe_before_ks;

  // All tmp lwe arrays and index arrays for lwe contain the total
  // amount of blocks to be computed on, there is no split between GPUs
  // for the moment
  InputTorus *lwe_indexes_in;

  InputTorus *h_lwe_indexes_in;
  InputTorus *lwe_trivial_indexes;

  /// For multi GPU execution we create vectors of pointers for inputs and
  /// outputs
  std::vector<InputTorus *> lwe_array_in_vec;
  std::vector<InputTorus *> lwe_after_ks_vec;
  std::vector<__uint128_t *> lwe_after_pbs_vec;
  std::vector<InputTorus *> lwe_trivial_indexes_vec;

  bool using_trivial_lwe_indexes = true;
  bool gpu_memory_allocated;
  std::vector<InputTorus *> lwe_aligned_scatter_vec;
  std::vector<__uint128_t *> lwe_aligned_gather_vec;
  // noise squashing constructor
  int_noise_squashing_lut(CudaStreams streams, int_radix_params params,
                          uint32_t input_glwe_dimension,
                          uint32_t input_polynomial_size,
                          uint32_t num_radix_blocks,
                          uint32_t original_num_blocks,
                          bool allocate_gpu_memory, uint64_t &size_tracker) {
    this->params = params;
    this->num_blocks = num_radix_blocks;
    gpu_memory_allocated = allocate_gpu_memory;
    // This are the glwe dimension and polynomial size before squashing
    this->input_glwe_dimension = input_glwe_dimension;
    this->input_polynomial_size = input_polynomial_size;
    uint32_t input_big_lwe_dimension =
        input_glwe_dimension * input_polynomial_size;
    this->input_big_lwe_dimension = input_big_lwe_dimension;

    uint64_t lut_buffer_size = (params.glwe_dimension + 1) *
                               params.polynomial_size * sizeof(__uint128_t);

    ///////////////
    active_streams = streams.active_gpu_subset(num_radix_blocks);
    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
    for (uint i = 0; i < active_streams.count(); i++) {
      cuda_set_device(streams.gpu_index(i));
      auto num_radix_blocks_on_gpu = std::max(
          THRESHOLD_MULTI_GPU,
          get_num_inputs_on_gpu(num_radix_blocks, i, active_streams.count()));
      int8_t *gpu_pbs_buffer;
      uint64_t size = 0;
      execute_scratch_pbs<__uint128_t>(
          streams.stream(i), streams.gpu_index(i), &gpu_pbs_buffer,
          params.glwe_dimension, params.small_lwe_dimension,
          params.polynomial_size, params.pbs_level, params.grouping_factor,
          num_radix_blocks_on_gpu, params.pbs_type, allocate_gpu_memory,
          params.noise_reduction_type, size);
      cuda_synchronize_stream(streams.stream(i), streams.gpu_index(i));
      if (i == 0) {
        size_tracker += size;
      }
      pbs_buffer.push_back(gpu_pbs_buffer);
    }
    lwe_indexes_in = (InputTorus *)cuda_malloc_with_size_tracking_async(
        num_radix_blocks * sizeof(InputTorus), streams.stream(0),
        streams.gpu_index(0), size_tracker, allocate_gpu_memory);
    lwe_trivial_indexes = (InputTorus *)cuda_malloc_with_size_tracking_async(
        num_radix_blocks * sizeof(InputTorus), streams.stream(0),
        streams.gpu_index(0), size_tracker, allocate_gpu_memory);
    h_lwe_indexes_in =
        (InputTorus *)malloc(num_radix_blocks * sizeof(InputTorus));
    for (int i = 0; i < num_radix_blocks; i++)
      h_lwe_indexes_in[i] = i;

    cuda_memcpy_with_size_tracking_async_to_gpu(
        lwe_indexes_in, h_lwe_indexes_in, num_radix_blocks * sizeof(InputTorus),
        streams.stream(0), streams.gpu_index(0), allocate_gpu_memory);
    cuda_memcpy_with_size_tracking_async_to_gpu(
        lwe_trivial_indexes, h_lwe_indexes_in,
        num_radix_blocks * sizeof(InputTorus), streams.stream(0),
        streams.gpu_index(0), allocate_gpu_memory);

    multi_gpu_alloc_lwe_async(active_streams, lwe_array_in_vec,
                              num_radix_blocks, params.big_lwe_dimension + 1,
                              size_tracker, allocate_gpu_memory);

    multi_gpu_alloc_lwe_async<InputTorus>(
        active_streams, lwe_after_ks_vec, num_radix_blocks,
        params.small_lwe_dimension + 1, size_tracker, allocate_gpu_memory);
    multi_gpu_alloc_lwe_async<__uint128_t>(
        active_streams, lwe_after_pbs_vec, num_radix_blocks,
        params.big_lwe_dimension + 1, size_tracker, allocate_gpu_memory);
    multi_gpu_alloc_array_async<InputTorus>(
        active_streams, lwe_trivial_indexes_vec, num_radix_blocks, size_tracker,
        allocate_gpu_memory);
    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));

    multi_gpu_copy_array_async(active_streams, lwe_trivial_indexes_vec,
                               lwe_trivial_indexes, num_radix_blocks,
                               allocate_gpu_memory);
    if (allocate_gpu_memory) {
      // Allocate LUT
      // LUT is used as a trivial encryption and must be initialized outside
      // this constructor
      for (uint i = 0; i < active_streams.count(); i++) {
        auto lut = (__uint128_t *)cuda_malloc_with_size_tracking_async(
            lut_buffer_size, streams.stream(i), streams.gpu_index(i),
            size_tracker, allocate_gpu_memory);
        lut_vec.push_back(lut);
        cuda_synchronize_stream(streams.stream(i), streams.gpu_index(i));
      }
    }
    // Keyswitch
    tmp_lwe_before_ks = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<InputTorus>(
        streams.stream(0), streams.gpu_index(0), tmp_lwe_before_ks,
        original_num_blocks, input_big_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    degrees = (uint64_t *)malloc(sizeof(uint64_t));
    max_degrees = (uint64_t *)malloc(sizeof(uint64_t));

    // lut for the squashing
    auto f_squash = [](__uint128_t block) -> __uint128_t { return block; };

    // Generate the identity LUT, for now we only use one GPU
    for (uint i = 0; i < active_streams.count(); i++) {
      auto squash_lut = lut_vec[i];
      generate_device_accumulator<__uint128_t>(
          streams.stream(i), streams.gpu_index(i), squash_lut, degrees,
          max_degrees, params.glwe_dimension, params.polynomial_size,
          params.message_modulus, params.carry_modulus, f_squash,
          allocate_gpu_memory);
    }
  }
  void release(CudaStreams streams) {
    for (uint i = 0; i < lut_vec.size(); i++) {
      cuda_drop_with_size_tracking_async(lut_vec[i], streams.stream(i),
                                         streams.gpu_index(i),
                                         gpu_memory_allocated);
    }
    cuda_drop_with_size_tracking_async(lwe_indexes_in, streams.stream(0),
                                       streams.gpu_index(0),
                                       gpu_memory_allocated);
    cuda_drop_with_size_tracking_async(lwe_trivial_indexes, streams.stream(0),
                                       streams.gpu_index(0),
                                       gpu_memory_allocated);
    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
    lut_vec.clear();
    free(h_lwe_indexes_in);

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_lwe_before_ks, gpu_memory_allocated);
    for (int i = 0; i < pbs_buffer.size(); i++) {
      switch (params.pbs_type) {
      case MULTI_BIT:
        cleanup_cuda_multi_bit_programmable_bootstrap_128(
            streams.stream(i), streams.gpu_index(i), &pbs_buffer[i]);
        break;
      case CLASSICAL:
        cleanup_cuda_programmable_bootstrap_128(
            streams.stream(i), streams.gpu_index(i), &pbs_buffer[i]);
        break;
      default:
        PANIC("Cuda error (PBS): unknown PBS type. ")
      }
      cuda_synchronize_stream(streams.stream(i), streams.gpu_index(i));
    }
    if (lwe_aligned_gather_vec.size() > 0) {
      multi_gpu_release_async(streams, lwe_aligned_gather_vec);
      multi_gpu_release_async(streams, lwe_aligned_scatter_vec);
    }
    multi_gpu_release_async(streams, lwe_array_in_vec);
    multi_gpu_release_async(streams, lwe_after_ks_vec);
    multi_gpu_release_async(streams, lwe_after_pbs_vec);
    multi_gpu_release_async(streams, lwe_trivial_indexes_vec);
    streams.synchronize();

    lwe_array_in_vec.clear();
    lwe_after_ks_vec.clear();
    lwe_after_pbs_vec.clear();
    lwe_trivial_indexes_vec.clear();

    delete tmp_lwe_before_ks;
    pbs_buffer.clear();
  }
};

// Forward declarations for operation buffers
template <typename Torus> struct int_sub_and_propagate;
template <typename Torus> struct int_bitop_buffer;

template <typename Torus> struct int_bit_extract_luts_buffer {
  int_radix_params params;
  int_radix_lut<Torus> *lut;
  bool gpu_memory_allocated;

  // With offset
  int_bit_extract_luts_buffer(CudaStreams streams, int_radix_params params,
                              uint32_t bits_per_block, uint32_t final_offset,
                              uint32_t num_radix_blocks,
                              bool allocate_gpu_memory,
                              uint64_t &size_tracker) {
    this->params = params;
    gpu_memory_allocated = allocate_gpu_memory;

    lut = new int_radix_lut<Torus>(streams, params, bits_per_block,
                                   bits_per_block * num_radix_blocks,
                                   allocate_gpu_memory, size_tracker);

    for (int i = 0; i < bits_per_block; i++) {

      auto operator_f = [i, final_offset](Torus x) -> Torus {
        Torus y = (x >> i) & 1;
        return y << final_offset;
      };

      generate_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0), lut->get_lut(0, i),
          lut->get_degree(i), lut->get_max_degree(i), params.glwe_dimension,
          params.polynomial_size, params.message_modulus, params.carry_modulus,
          operator_f, gpu_memory_allocated);
    }

    /**
     * we have bits_per_blocks LUTs that should be used for all bits in all
     * blocks
     */
    Torus *h_lut_indexes = lut->h_lut_indexes;
    for (int j = 0; j < num_radix_blocks; j++) {
      for (int i = 0; i < bits_per_block; i++)
        h_lut_indexes[i + j * bits_per_block] = i;
    }
    cuda_memcpy_with_size_tracking_async_to_gpu(
        lut->get_lut_indexes(0, 0), h_lut_indexes,
        num_radix_blocks * bits_per_block * sizeof(Torus), streams.stream(0),
        streams.gpu_index(0), allocate_gpu_memory);

    auto active_streams =
        streams.active_gpu_subset(bits_per_block * num_radix_blocks);
    lut->broadcast_lut(active_streams);

    /**
     * the input indexes should take the first bits_per_block PBS to target
     * the block 0, then block 1, etc...
     */
    Torus *h_lwe_indexes_in =
        (Torus *)malloc(num_radix_blocks * bits_per_block * sizeof(Torus));

    for (int j = 0; j < num_radix_blocks; j++) {
      for (int i = 0; i < bits_per_block; i++)
        h_lwe_indexes_in[i + j * bits_per_block] = j;
    }

    /**
     * the output should aim different lwe ciphertexts, so lwe_indexes_out =
     * range(num_luts)
     */
    Torus *h_lwe_indexes_out =
        (Torus *)malloc(num_radix_blocks * bits_per_block * sizeof(Torus));

    for (int i = 0; i < num_radix_blocks * bits_per_block; i++)
      h_lwe_indexes_out[i] = i;

    lut->set_lwe_indexes(streams.stream(0), streams.gpu_index(0),
                         h_lwe_indexes_in, h_lwe_indexes_out);
    lut->allocate_lwe_vector_for_non_trivial_indexes(
        active_streams, num_radix_blocks * bits_per_block, size_tracker,
        allocate_gpu_memory);

    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
    free(h_lwe_indexes_in);
    free(h_lwe_indexes_out);
  }

  // Without offset
  int_bit_extract_luts_buffer(CudaStreams streams, int_radix_params params,
                              uint32_t bits_per_block,
                              uint32_t num_radix_blocks,
                              bool allocate_gpu_memory, uint64_t &size_tracker)
      : int_bit_extract_luts_buffer(streams, params, bits_per_block, 0,
                                    num_radix_blocks, allocate_gpu_memory,
                                    size_tracker) {}

  void release(CudaStreams streams) {
    lut->release(streams);
    delete (lut);
  }
};

template <typename Torus> struct int_shift_and_rotate_buffer {
  int_radix_params params;
  SHIFT_OR_ROTATE_TYPE shift_type;
  bool is_signed;

  CudaRadixCiphertextFFI *tmp_bits;
  CudaRadixCiphertextFFI *tmp_shift_bits;
  CudaRadixCiphertextFFI *tmp_rotated;
  CudaRadixCiphertextFFI *tmp_input_bits_a;
  CudaRadixCiphertextFFI *tmp_input_bits_b;
  CudaRadixCiphertextFFI *tmp_mux_inputs;

  int_bit_extract_luts_buffer<Torus> *bit_extract_luts;
  int_bit_extract_luts_buffer<Torus> *bit_extract_luts_with_offset_2;
  int_radix_lut<Torus> *mux_lut;
  int_radix_lut<Torus> *cleaning_lut;

  Torus offset;
  bool gpu_memory_allocated;

  int_shift_and_rotate_buffer(CudaStreams streams,
                              SHIFT_OR_ROTATE_TYPE shift_type, bool is_signed,
                              int_radix_params params,
                              uint32_t num_radix_blocks,
                              bool allocate_gpu_memory,
                              uint64_t &size_tracker) {
    this->shift_type = shift_type;
    this->is_signed = is_signed;
    this->params = params;
    gpu_memory_allocated = allocate_gpu_memory;

    uint32_t bits_per_block = std::log2(params.message_modulus);
    uint32_t total_nb_bits =
        std::log2(params.message_modulus) * num_radix_blocks;
    uint32_t max_num_bits_that_tell_shift = std::log2(total_nb_bits);

    auto is_power_of_two = [](uint32_t n) {
      return (n > 0) && ((n & (n - 1)) == 0);
    };

    if (!is_power_of_two(total_nb_bits))
      max_num_bits_that_tell_shift += 1;

    offset = (shift_type == LEFT_SHIFT ? 0 : total_nb_bits);

    bit_extract_luts = new int_bit_extract_luts_buffer<Torus>(
        streams, params, bits_per_block, num_radix_blocks, allocate_gpu_memory,
        size_tracker);
    bit_extract_luts_with_offset_2 = new int_bit_extract_luts_buffer<Torus>(
        streams, params, bits_per_block, 2, num_radix_blocks,
        allocate_gpu_memory, size_tracker);

    mux_lut = new int_radix_lut<Torus>(streams, params, 1,
                                       bits_per_block * num_radix_blocks,
                                       allocate_gpu_memory, size_tracker);
    cleaning_lut =
        new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                 allocate_gpu_memory, size_tracker);

    tmp_bits = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_bits,
        bits_per_block * num_radix_blocks, params.big_lwe_dimension,
        size_tracker, allocate_gpu_memory);

    tmp_shift_bits = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_shift_bits,
        max_num_bits_that_tell_shift * num_radix_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

    tmp_rotated = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_rotated,
        bits_per_block * num_radix_blocks, params.big_lwe_dimension,
        size_tracker, allocate_gpu_memory);

    tmp_input_bits_a = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_input_bits_a,
        bits_per_block * num_radix_blocks, params.big_lwe_dimension,
        size_tracker, allocate_gpu_memory);

    tmp_input_bits_b = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_input_bits_b,
        bits_per_block * num_radix_blocks, params.big_lwe_dimension,
        size_tracker, allocate_gpu_memory);

    tmp_mux_inputs = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_mux_inputs,
        bits_per_block * num_radix_blocks, params.big_lwe_dimension,
        size_tracker, allocate_gpu_memory);

    auto mux_lut_f = [](Torus x) -> Torus {
      // x is expected to be x = 0bcba
      // where
      // - c is the control bit
      // - b the bit value returned if c is 1
      // - a the bit value returned if c is 0
      // (any bit above c is ignored)
      x = x & 7;
      auto control_bit = x >> 2;
      auto previous_bit = (x & 2) >> 1;
      auto current_bit = x & 1;

      if (control_bit == 1)
        return previous_bit;
      else
        return current_bit;
    };

    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0), mux_lut->get_lut(0, 0),
        mux_lut->get_degree(0), mux_lut->get_max_degree(0),
        params.glwe_dimension, params.polynomial_size, params.message_modulus,
        params.carry_modulus, mux_lut_f, gpu_memory_allocated);
    auto active_gpu_count_mux =
        streams.active_gpu_subset(bits_per_block * num_radix_blocks);
    mux_lut->broadcast_lut(active_gpu_count_mux);

    auto cleaning_lut_f = [params](Torus x) -> Torus {
      return x % params.message_modulus;
    };
    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0), cleaning_lut->get_lut(0, 0),
        cleaning_lut->get_degree(0), cleaning_lut->get_max_degree(0),
        params.glwe_dimension, params.polynomial_size, params.message_modulus,
        params.carry_modulus, cleaning_lut_f, gpu_memory_allocated);
    auto active_gpu_count_cleaning =
        streams.active_gpu_subset(num_radix_blocks);
    cleaning_lut->broadcast_lut(active_gpu_count_cleaning);
  }

  void release(CudaStreams streams) {
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_bits, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_shift_bits, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_rotated, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_input_bits_a, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_input_bits_b, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_mux_inputs, gpu_memory_allocated);

    bit_extract_luts->release(streams);
    bit_extract_luts_with_offset_2->release(streams);
    mux_lut->release(streams);
    cleaning_lut->release(streams);

    delete tmp_bits;
    delete tmp_shift_bits;
    delete tmp_rotated;
    delete tmp_input_bits_a;
    delete tmp_input_bits_b;
    delete tmp_mux_inputs;
    delete bit_extract_luts;
    delete bit_extract_luts_with_offset_2;
    delete mux_lut;
    delete cleaning_lut;
  }
};

template <typename Torus> struct int_fullprop_buffer {
  int_radix_params params;

  int_radix_lut<Torus> *lut;

  CudaRadixCiphertextFFI *tmp_small_lwe_vector;
  CudaRadixCiphertextFFI *tmp_big_lwe_vector;
  bool gpu_memory_allocated;

  int_fullprop_buffer(CudaStreams streams, int_radix_params params,
                      bool allocate_gpu_memory, uint64_t &size_tracker) {
    this->params = params;
    gpu_memory_allocated = allocate_gpu_memory;
    lut = new int_radix_lut<Torus>(streams.get_ith(0), params, 2, 2,
                                   allocate_gpu_memory, size_tracker);

    // LUTs
    auto lut_f_message = [params](Torus x) -> Torus {
      return x % params.message_modulus;
    };
    auto lut_f_carry = [params](Torus x) -> Torus {
      return x / params.message_modulus;
    };

    //
    Torus *lut_buffer_message = lut->get_lut(0, 0);
    uint64_t *message_degree = lut->get_degree(0);
    uint64_t *message_max_degree = lut->get_max_degree(0);
    Torus *lut_buffer_carry = lut->get_lut(0, 1);
    uint64_t *carry_degree = lut->get_degree(1);
    uint64_t *carry_max_degree = lut->get_max_degree(1);

    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0), lut_buffer_message,
        message_degree, message_max_degree, params.glwe_dimension,
        params.polynomial_size, params.message_modulus, params.carry_modulus,
        lut_f_message, gpu_memory_allocated);

    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0), lut_buffer_carry, carry_degree,
        carry_max_degree, params.glwe_dimension, params.polynomial_size,
        params.message_modulus, params.carry_modulus, lut_f_carry,
        gpu_memory_allocated);

    uint64_t lwe_indexes_size = 2 * sizeof(Torus);
    Torus *h_lwe_indexes = (Torus *)malloc(lwe_indexes_size);
    for (int i = 0; i < 2; i++)
      h_lwe_indexes[i] = i;
    Torus *lwe_indexes = lut->get_lut_indexes(0, 0);
    cuda_memcpy_with_size_tracking_async_to_gpu(
        lwe_indexes, h_lwe_indexes, lwe_indexes_size, streams.stream(0),
        streams.gpu_index(0), allocate_gpu_memory);
    auto active_streams = streams.active_gpu_subset(2);
    lut->broadcast_lut(active_streams);

    tmp_small_lwe_vector = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_small_lwe_vector, 2,
        params.small_lwe_dimension, size_tracker, allocate_gpu_memory);
    tmp_big_lwe_vector = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_big_lwe_vector, 2,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
    free(h_lwe_indexes);
  }

  void release(CudaStreams streams) {

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_small_lwe_vector, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_big_lwe_vector, gpu_memory_allocated);
    lut->release(streams.get_ith(0));
    delete tmp_small_lwe_vector;
    delete tmp_big_lwe_vector;
    delete lut;
  }
};

template <typename Torus> struct int_overflowing_sub_memory {
  Torus *generates_or_propagates;
  Torus *step_output;

  int_radix_lut<Torus> *luts_array;
  int_radix_lut<Torus> *luts_borrow_propagation_sum;
  int_radix_lut<Torus> *message_acc;

  int_radix_params params;
  bool gpu_memory_allocated;

  int_overflowing_sub_memory(CudaStreams streams, int_radix_params params,
                             uint32_t num_radix_blocks,
                             bool allocate_gpu_memory, uint64_t &size_tracker) {
    this->params = params;
    gpu_memory_allocated = allocate_gpu_memory;
    auto glwe_dimension = params.glwe_dimension;
    auto polynomial_size = params.polynomial_size;
    auto message_modulus = params.message_modulus;
    auto carry_modulus = params.carry_modulus;
    auto big_lwe_size = (polynomial_size * glwe_dimension + 1);
    auto big_lwe_size_bytes = big_lwe_size * sizeof(Torus);

    // allocate memory for intermediate calculations
    generates_or_propagates = (Torus *)cuda_malloc_with_size_tracking_async(
        num_radix_blocks * big_lwe_size_bytes, streams.stream(0),
        streams.gpu_index(0), size_tracker, allocate_gpu_memory);
    step_output = (Torus *)cuda_malloc_with_size_tracking_async(
        num_radix_blocks * big_lwe_size_bytes, streams.stream(0),
        streams.gpu_index(0), size_tracker, allocate_gpu_memory);
    cuda_memset_with_size_tracking_async(
        generates_or_propagates, 0, num_radix_blocks * big_lwe_size_bytes,
        streams.stream(0), streams.gpu_index(0), allocate_gpu_memory);
    cuda_memset_with_size_tracking_async(
        step_output, 0, num_radix_blocks * big_lwe_size_bytes,
        streams.stream(0), streams.gpu_index(0), allocate_gpu_memory);

    // declare functions for lut generation
    auto f_lut_does_block_generate_carry = [message_modulus](Torus x) -> Torus {
      if (x < message_modulus)
        return OUTPUT_CARRY::GENERATED;
      return OUTPUT_CARRY::NONE;
    };

    auto f_lut_does_block_generate_or_propagate =
        [message_modulus](Torus x) -> Torus {
      if (x < message_modulus)
        return OUTPUT_CARRY::GENERATED;
      else if (x == message_modulus)
        return OUTPUT_CARRY::PROPAGATED;
      return OUTPUT_CARRY::NONE;
    };

    auto f_luts_borrow_propagation_sum = [](Torus msb, Torus lsb) -> Torus {
      if (msb == OUTPUT_CARRY::PROPAGATED)
        return lsb;
      return msb;
    };

    auto f_message_acc = [message_modulus](Torus x) -> Torus {
      return x % message_modulus;
    };

    // create lut objects
    luts_array = new int_radix_lut<Torus>(streams, params, 2, num_radix_blocks,
                                          allocate_gpu_memory, size_tracker);
    luts_borrow_propagation_sum = new int_radix_lut<Torus>(
        streams, params, 1, num_radix_blocks, luts_array, size_tracker,
        allocate_gpu_memory, size_tracker);
    message_acc = new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                           luts_array, size_tracker,
                                           allocate_gpu_memory, size_tracker);

    auto lut_does_block_generate_carry = luts_array->get_lut(0, 0);
    auto lut_does_block_generate_or_propagate = luts_array->get_lut(0, 1);

    // generate luts (aka accumulators)
    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0), lut_does_block_generate_carry,
        luts_array->get_degree(0), luts_array->get_max_degree(0),
        glwe_dimension, polynomial_size, message_modulus, carry_modulus,
        f_lut_does_block_generate_carry, gpu_memory_allocated);
    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0),
        lut_does_block_generate_or_propagate, luts_array->get_degree(1),
        luts_array->get_max_degree(1), glwe_dimension, polynomial_size,
        message_modulus, carry_modulus, f_lut_does_block_generate_or_propagate,
        gpu_memory_allocated);
    if (allocate_gpu_memory)
      cuda_set_value_async<Torus>(streams.stream(0), streams.gpu_index(0),
                                  luts_array->get_lut_indexes(0, 1), 1,
                                  num_radix_blocks - 1);

    generate_device_accumulator_bivariate<Torus>(
        streams.stream(0), streams.gpu_index(0),
        luts_borrow_propagation_sum->get_lut(0, 0),
        luts_borrow_propagation_sum->get_degree(0),
        luts_borrow_propagation_sum->get_max_degree(0), glwe_dimension,
        polynomial_size, message_modulus, carry_modulus,
        f_luts_borrow_propagation_sum, gpu_memory_allocated);

    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0), message_acc->get_lut(0, 0),
        message_acc->get_degree(0), message_acc->get_max_degree(0),
        glwe_dimension, polynomial_size, message_modulus, carry_modulus,
        f_message_acc, gpu_memory_allocated);

    auto active_streams = streams.active_gpu_subset(num_radix_blocks);
    luts_array->broadcast_lut(active_streams);
    luts_borrow_propagation_sum->broadcast_lut(active_streams);
    message_acc->broadcast_lut(active_streams);
  }

  void release(CudaStreams streams) {
    cuda_drop_with_size_tracking_async(generates_or_propagates,
                                       streams.stream(0), streams.gpu_index(0),
                                       gpu_memory_allocated);
    cuda_drop_with_size_tracking_async(step_output, streams.stream(0),
                                       streams.gpu_index(0),
                                       gpu_memory_allocated);

    luts_array->release(streams);
    luts_borrow_propagation_sum->release(streams);
    message_acc->release(streams);

    delete luts_array;
    delete luts_borrow_propagation_sum;
    delete message_acc;
  }
};

template <typename Torus> struct int_sum_ciphertexts_vec_memory {

  int_radix_params params;
  uint32_t max_total_blocks_in_vec;
  uint32_t num_blocks_in_radix;
  uint32_t max_num_radix_in_vec;
  uint32_t chunk_size;
  bool gpu_memory_allocated;
  bool reduce_degrees_for_single_carry_propagation;

  // temporary buffers
  CudaRadixCiphertextFFI *current_blocks;
  CudaRadixCiphertextFFI *small_lwe_vector;

  uint32_t *d_columns_data;
  uint32_t *d_columns_counter;
  uint32_t **d_columns;

  uint32_t *d_new_columns_data;
  uint32_t *d_new_columns_counter;
  uint32_t **d_new_columns;

  uint64_t *d_degrees;

  // lookup table for extracting message and carry
  int_radix_lut<Torus> *luts_message_carry;

  bool mem_reuse = false;
  bool allocated_luts_message_carry;

  void setup_index_buffers(CudaStreams streams, uint64_t &size_tracker) {

    d_degrees = (uint64_t *)cuda_malloc_with_size_tracking_async(
        max_total_blocks_in_vec * sizeof(uint64_t), streams.stream(0),
        streams.gpu_index(0), size_tracker, gpu_memory_allocated);

    auto num_blocks_in_radix = this->num_blocks_in_radix;
    auto max_num_radix_in_vec = this->max_num_radix_in_vec;
    auto setup_columns = [num_blocks_in_radix, max_num_radix_in_vec, streams](
                             uint32_t **&columns, uint32_t *&columns_data,
                             uint32_t *&columns_counter, uint64_t &size_tracker,
                             bool gpu_memory_allocated) {
      columns_data = (uint32_t *)cuda_malloc_with_size_tracking_async(
          num_blocks_in_radix * max_num_radix_in_vec * sizeof(uint32_t),
          streams.stream(0), streams.gpu_index(0), size_tracker,
          gpu_memory_allocated);
      columns_counter = (uint32_t *)cuda_malloc_with_size_tracking_async(
          num_blocks_in_radix * sizeof(uint32_t), streams.stream(0),
          streams.gpu_index(0), size_tracker, gpu_memory_allocated);
      cuda_memset_with_size_tracking_async(
          columns_counter, 0, num_blocks_in_radix * sizeof(uint32_t),
          streams.stream(0), streams.gpu_index(0), gpu_memory_allocated);
      uint32_t **h_columns = new uint32_t *[num_blocks_in_radix];
      for (int i = 0; i < num_blocks_in_radix; ++i) {
        h_columns[i] = columns_data + i * max_num_radix_in_vec;
      }
      columns = (uint32_t **)cuda_malloc_with_size_tracking_async(
          num_blocks_in_radix * sizeof(uint32_t *), streams.stream(0),
          streams.gpu_index(0), size_tracker, gpu_memory_allocated);
      if (gpu_memory_allocated) {
        cuda_memcpy_async_to_gpu(columns, h_columns,
                                 num_blocks_in_radix * sizeof(uint32_t *),
                                 streams.stream(0), streams.gpu_index(0));
      }
      cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
      delete[] h_columns;
    };

    setup_columns(d_columns, d_columns_data, d_columns_counter, size_tracker,
                  gpu_memory_allocated);
    setup_columns(d_new_columns, d_new_columns_data, d_new_columns_counter,
                  size_tracker, gpu_memory_allocated);
  }

  void setup_lookup_tables(CudaStreams streams, uint32_t num_radix_in_vec,
                           const uint64_t *const degrees) {
    uint32_t message_modulus = params.message_modulus;
    bool _needs_processing = false;
    radix_columns current_columns(degrees, num_blocks_in_radix,
                                  num_radix_in_vec, chunk_size,
                                  _needs_processing);
    uint32_t total_ciphertexts = 0;
    uint32_t total_messages = 0;
    current_columns.next_accumulation(total_ciphertexts, total_messages,
                                      _needs_processing);
    uint32_t pbs_count = std::max(total_ciphertexts, 2 * num_blocks_in_radix);
    if (!mem_reuse) {
      if (total_ciphertexts > 0 ||
          reduce_degrees_for_single_carry_propagation) {
        uint64_t size_tracker = 0;
        luts_message_carry = new int_radix_lut<Torus>(
            streams, params, 2, pbs_count, true, size_tracker);
        allocated_luts_message_carry = true;
        uint64_t message_modulus_bits =
            (uint64_t)std::log2(params.message_modulus);
        uint64_t carry_modulus_bits = (uint64_t)std::log2(params.carry_modulus);
        uint64_t total_bits_per_block =
            message_modulus_bits + carry_modulus_bits;
        uint64_t denominator =
            (uint64_t)std::ceil((pow(2, total_bits_per_block) - 1) /
                                (pow(2, message_modulus_bits) - 1));

        uint64_t upper_bound_num_blocks =
            max_total_blocks_in_vec * 2 / denominator;
        luts_message_carry->allocate_lwe_vector_for_non_trivial_indexes(
            streams, upper_bound_num_blocks, size_tracker, true);
      }
    }
    if (allocated_luts_message_carry) {
      auto message_acc = luts_message_carry->get_lut(0, 0);
      auto carry_acc = luts_message_carry->get_lut(0, 1);

      // define functions for each accumulator
      auto lut_f_message = [message_modulus](Torus x) -> Torus {
        return x % message_modulus;
      };
      auto lut_f_carry = [message_modulus](Torus x) -> Torus {
        return x / message_modulus;
      };

      // generate accumulators
      generate_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0), message_acc,
          luts_message_carry->get_degree(0),
          luts_message_carry->get_max_degree(0), params.glwe_dimension,
          params.polynomial_size, message_modulus, params.carry_modulus,
          lut_f_message, gpu_memory_allocated);
      generate_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0), carry_acc,
          luts_message_carry->get_degree(1),
          luts_message_carry->get_max_degree(1), params.glwe_dimension,
          params.polynomial_size, message_modulus, params.carry_modulus,
          lut_f_carry, gpu_memory_allocated);
      auto active_gpu_count_mc = streams.active_gpu_subset(pbs_count);
      luts_message_carry->broadcast_lut(active_gpu_count_mc);
    }
  }
  int_sum_ciphertexts_vec_memory(
      CudaStreams streams, int_radix_params params,
      uint32_t num_blocks_in_radix, uint32_t max_num_radix_in_vec,
      bool reduce_degrees_for_single_carry_propagation,
      bool allocate_gpu_memory, uint64_t &size_tracker) {
    this->params = params;
    this->mem_reuse = false;
    this->max_total_blocks_in_vec = num_blocks_in_radix * max_num_radix_in_vec;
    this->num_blocks_in_radix = num_blocks_in_radix;
    this->max_num_radix_in_vec = max_num_radix_in_vec;
    this->gpu_memory_allocated = allocate_gpu_memory;
    this->chunk_size = (params.message_modulus * params.carry_modulus - 1) /
                       (params.message_modulus - 1);
    this->allocated_luts_message_carry = false;
    this->reduce_degrees_for_single_carry_propagation =
        reduce_degrees_for_single_carry_propagation;

    setup_index_buffers(streams, size_tracker);
    // because we setup_lut in host function for sum_ciphertexts to save memory
    // the size_tracker is topped up here to have a max bound on the used memory
    uint32_t max_pbs_count = std::max(
        2 * (max_total_blocks_in_vec / chunk_size), 2 * num_blocks_in_radix);
    if (max_pbs_count > 0) {
      int_radix_lut<Torus> *luts_message_carry_dry_run =
          new int_radix_lut<Torus>(streams, params, 2, max_pbs_count, false,
                                   size_tracker);
      luts_message_carry_dry_run->release(streams);
      delete luts_message_carry_dry_run;
    }

    // create and allocate intermediate buffers
    current_blocks = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), current_blocks,
        max_total_blocks_in_vec, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    small_lwe_vector = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), small_lwe_vector,
        max_total_blocks_in_vec, params.small_lwe_dimension, size_tracker,
        allocate_gpu_memory);
  }

  int_sum_ciphertexts_vec_memory(
      CudaStreams streams, int_radix_params params,
      uint32_t num_blocks_in_radix, uint32_t max_num_radix_in_vec,
      CudaRadixCiphertextFFI *current_blocks,
      CudaRadixCiphertextFFI *small_lwe_vector,
      int_radix_lut<Torus> *reused_lut,
      bool reduce_degrees_for_single_carry_propagation,
      bool allocate_gpu_memory, uint64_t &size_tracker) {
    this->mem_reuse = true;
    this->params = params;
    this->max_total_blocks_in_vec = num_blocks_in_radix * max_num_radix_in_vec;
    this->num_blocks_in_radix = num_blocks_in_radix;
    this->max_num_radix_in_vec = max_num_radix_in_vec;
    this->gpu_memory_allocated = allocate_gpu_memory;
    this->chunk_size = (params.message_modulus * params.carry_modulus - 1) /
                       (params.message_modulus - 1);
    this->allocated_luts_message_carry = true;
    this->reduce_degrees_for_single_carry_propagation =
        reduce_degrees_for_single_carry_propagation;

    this->current_blocks = current_blocks;
    this->small_lwe_vector = small_lwe_vector;
    this->luts_message_carry = reused_lut;

    uint64_t message_modulus_bits = (uint64_t)std::log2(params.message_modulus);
    uint64_t carry_modulus_bits = (uint64_t)std::log2(params.carry_modulus);
    uint64_t total_bits_per_block = message_modulus_bits + carry_modulus_bits;
    uint64_t denominator =
        (uint64_t)std::ceil((pow(2, total_bits_per_block) - 1) /
                            (pow(2, message_modulus_bits) - 1));

    uint64_t upper_bound_num_blocks = max_total_blocks_in_vec * 2 / denominator;
    this->luts_message_carry->allocate_lwe_vector_for_non_trivial_indexes(
        streams, upper_bound_num_blocks, size_tracker, allocate_gpu_memory);
    setup_index_buffers(streams, size_tracker);
  }

  void release(CudaStreams streams) {
    cuda_drop_with_size_tracking_async(d_degrees, streams.stream(0),
                                       streams.gpu_index(0),
                                       gpu_memory_allocated);
    cuda_drop_with_size_tracking_async(d_columns_data, streams.stream(0),
                                       streams.gpu_index(0),
                                       gpu_memory_allocated);
    cuda_drop_with_size_tracking_async(d_columns_counter, streams.stream(0),
                                       streams.gpu_index(0),
                                       gpu_memory_allocated);
    cuda_drop_with_size_tracking_async(d_columns, streams.stream(0),
                                       streams.gpu_index(0),
                                       gpu_memory_allocated);

    cuda_drop_with_size_tracking_async(d_new_columns_data, streams.stream(0),
                                       streams.gpu_index(0),
                                       gpu_memory_allocated);
    cuda_drop_with_size_tracking_async(d_new_columns_counter, streams.stream(0),
                                       streams.gpu_index(0),
                                       gpu_memory_allocated);
    cuda_drop_with_size_tracking_async(d_new_columns, streams.stream(0),
                                       streams.gpu_index(0),
                                       gpu_memory_allocated);

    if (!mem_reuse) {
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     current_blocks, gpu_memory_allocated);
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     small_lwe_vector, gpu_memory_allocated);
      if (allocated_luts_message_carry) {
        luts_message_carry->release(streams);
        delete luts_message_carry;
      }
      delete current_blocks;
      delete small_lwe_vector;
    }
  }
};

// For sequential algorithm in group propagation
template <typename Torus> struct int_seq_group_prop_memory {

  CudaRadixCiphertextFFI *group_resolved_carries;
  int_radix_lut<Torus> *lut_sequential_algorithm;
  uint32_t grouping_size;
  bool gpu_memory_allocated;

  int_seq_group_prop_memory(CudaStreams streams, int_radix_params params,
                            uint32_t group_size, uint32_t big_lwe_size_bytes,
                            bool allocate_gpu_memory, uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;
    auto glwe_dimension = params.glwe_dimension;
    auto polynomial_size = params.polynomial_size;
    auto message_modulus = params.message_modulus;
    auto carry_modulus = params.carry_modulus;

    grouping_size = group_size;
    group_resolved_carries = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), group_resolved_carries,
        grouping_size, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    int num_seq_luts = grouping_size - 1;
    Torus *h_seq_lut_indexes = (Torus *)malloc(num_seq_luts * sizeof(Torus));
    lut_sequential_algorithm =
        new int_radix_lut<Torus>(streams, params, num_seq_luts, num_seq_luts,
                                 allocate_gpu_memory, size_tracker);
    for (int index = 0; index < num_seq_luts; index++) {
      auto f_lut_sequential = [index](Torus propa_cum_sum_block) {
        return (propa_cum_sum_block >> (index + 1)) & 1;
      };
      auto seq_lut = lut_sequential_algorithm->get_lut(0, index);
      generate_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0), seq_lut,
          lut_sequential_algorithm->get_degree(index),
          lut_sequential_algorithm->get_max_degree(index), glwe_dimension,
          polynomial_size, message_modulus, carry_modulus, f_lut_sequential,
          gpu_memory_allocated);
      h_seq_lut_indexes[index] = index;
    }
    Torus *seq_lut_indexes = lut_sequential_algorithm->get_lut_indexes(0, 0);
    cuda_memcpy_with_size_tracking_async_to_gpu(
        seq_lut_indexes, h_seq_lut_indexes, num_seq_luts * sizeof(Torus),
        streams.stream(0), streams.gpu_index(0), allocate_gpu_memory);
    auto active_streams = streams.active_gpu_subset(num_seq_luts);
    lut_sequential_algorithm->broadcast_lut(active_streams);
    free(h_seq_lut_indexes);
  };
  void release(CudaStreams streams) {
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   group_resolved_carries,
                                   gpu_memory_allocated);
    lut_sequential_algorithm->release(streams);
    delete group_resolved_carries;
    delete lut_sequential_algorithm;
  };
};

// For hillis steele algorithm in group propagation
template <typename Torus> struct int_hs_group_prop_memory {

  int_radix_lut<Torus> *lut_hillis_steele;
  bool gpu_memory_allocated;

  int_hs_group_prop_memory(CudaStreams streams, int_radix_params params,
                           uint32_t num_groups, uint32_t big_lwe_size_bytes,
                           bool allocate_gpu_memory, uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;
    auto glwe_dimension = params.glwe_dimension;
    auto polynomial_size = params.polynomial_size;
    auto message_modulus = params.message_modulus;
    auto carry_modulus = params.carry_modulus;

    auto f_lut_hillis_steele = [](Torus msb, Torus lsb) -> Torus {
      if (msb == 2) {
        return 1; // Remap Generate to 1
      } else if (msb == 3) {
        // MSB propagates
        if (lsb == 2) {
          return 1;
        } else {
          return lsb;
        } // also remap here
      } else {
        return msb;
      }
    };

    lut_hillis_steele = new int_radix_lut<Torus>(
        streams, params, 1, num_groups, allocate_gpu_memory, size_tracker);

    generate_device_accumulator_bivariate<Torus>(
        streams.stream(0), streams.gpu_index(0),
        lut_hillis_steele->get_lut(0, 0), lut_hillis_steele->get_degree(0),
        lut_hillis_steele->get_max_degree(0), glwe_dimension, polynomial_size,
        message_modulus, carry_modulus, f_lut_hillis_steele,
        gpu_memory_allocated);
    auto active_streams = streams.active_gpu_subset(num_groups);
    lut_hillis_steele->broadcast_lut(active_streams);
  };
  void release(CudaStreams streams) {

    lut_hillis_steele->release(streams);
    delete lut_hillis_steele;
  }
};

// compute_shifted_blocks_and_block_states
template <typename Torus> struct int_shifted_blocks_and_states_memory {
  CudaRadixCiphertextFFI *shifted_blocks_and_states;
  CudaRadixCiphertextFFI *shifted_blocks;
  CudaRadixCiphertextFFI *block_states;

  int_radix_lut<Torus> *luts_array_first_step;
  bool gpu_memory_allocated;

  int_shifted_blocks_and_states_memory(
      CudaStreams streams, int_radix_params params, uint32_t num_radix_blocks,
      uint32_t num_many_lut, uint32_t grouping_size, bool allocate_gpu_memory,
      uint64_t &size_tracker) {

    gpu_memory_allocated = allocate_gpu_memory;
    auto glwe_dimension = params.glwe_dimension;
    auto polynomial_size = params.polynomial_size;
    auto message_modulus = params.message_modulus;
    auto carry_modulus = params.carry_modulus;

    shifted_blocks_and_states = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), shifted_blocks_and_states,
        num_many_lut * num_radix_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    shifted_blocks = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), shifted_blocks,
        num_radix_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    block_states = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), block_states, num_radix_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

    uint32_t num_luts_first_step = 2 * grouping_size + 1;

    luts_array_first_step = new int_radix_lut<Torus>(
        streams, params, num_luts_first_step, num_radix_blocks, num_many_lut,
        allocate_gpu_memory, size_tracker);

    auto f_shift_block = [message_modulus](Torus block) -> Torus {
      return (block % message_modulus) << 1;
    };

    auto f_first_block_state = [message_modulus](Torus block) -> Torus {
      if (block >= message_modulus)
        return OUTPUT_CARRY::GENERATED;
      else {
        return OUTPUT_CARRY::NONE;
      }
    };
    std::vector<std::function<Torus(Torus)>> f_first_grouping_luts = {
        f_first_block_state, f_shift_block};

    auto first_block_lut = luts_array_first_step->get_lut(0, 0);
    auto first_block_lut_degrees = luts_array_first_step->get_degree(0);
    auto first_block_lut_max_degree = luts_array_first_step->get_max_degree(0);
    generate_many_lut_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0), first_block_lut,
        first_block_lut_degrees, first_block_lut_max_degree, glwe_dimension,
        polynomial_size, message_modulus, carry_modulus, f_first_grouping_luts,
        gpu_memory_allocated);

    // luts for other blocks of the first grouping
    for (int lut_id = 1; lut_id < grouping_size; lut_id++) {
      auto f_state = [message_modulus, lut_id](Torus block) -> Torus {
        uint64_t r = 0;
        if (block >= message_modulus) {
          r = 2; // Generates Carry
        } else if (block == (message_modulus - 1)) {
          r = 1; // Propagates a carry
        } else {
          r = 0; // Does not generate carry
        }
        return r << (lut_id - 1);
      };
      std::vector<std::function<Torus(Torus)>> f_grouping_luts = {
          f_state, f_shift_block};
      auto lut = luts_array_first_step->get_lut(0, lut_id);
      auto lut_degrees = luts_array_first_step->get_degree(lut_id);
      auto lut_max_degree = luts_array_first_step->get_max_degree(lut_id);
      generate_many_lut_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0), lut, lut_degrees,
          lut_max_degree, glwe_dimension, polynomial_size, message_modulus,
          carry_modulus, f_grouping_luts, gpu_memory_allocated);
    }

    // luts for the rest of groupings (except for the last block)
    for (int i = 0; i < grouping_size; i++) {
      uint32_t lut_id = i + grouping_size;
      auto f_state = [message_modulus, i](Torus block) -> Torus {
        uint64_t r = 0;
        if (block >= message_modulus) {
          r = 2; // Generates Carry
        } else if (block == (message_modulus - 1)) {
          r = 1; // Propagates a carry
        } else {
          r = 0; // Does not borrow
        }
        return r << i;
      };
      std::vector<std::function<Torus(Torus)>> f_grouping_luts = {
          f_state, f_shift_block};

      auto lut = luts_array_first_step->get_lut(0, lut_id);
      auto lut_degrees = luts_array_first_step->get_degree(lut_id);
      auto lut_max_degree = luts_array_first_step->get_max_degree(lut_id);
      generate_many_lut_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0), lut, lut_degrees,
          lut_max_degree, glwe_dimension, polynomial_size, message_modulus,
          carry_modulus, f_grouping_luts, gpu_memory_allocated);
    }

    // For the last block we need to generate a new lut
    auto f_last_block_state = [message_modulus](Torus block) -> Torus {
      if (block >= message_modulus)
        return 2 << 1; // Generates
      else
        return 0; // Nothing
    };

    uint32_t lut_id = num_luts_first_step - 1; // The last lut of the first step

    auto last_block_lut = luts_array_first_step->get_lut(0, lut_id);
    auto last_block_lut_degrees = luts_array_first_step->get_degree(lut_id);
    auto last_block_lut_max_degree =
        luts_array_first_step->get_max_degree(lut_id);

    std::vector<std::function<Torus(Torus)>> f_last_grouping_luts = {
        f_last_block_state, f_shift_block};

    generate_many_lut_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0), last_block_lut,
        last_block_lut_degrees, last_block_lut_max_degree, glwe_dimension,
        polynomial_size, message_modulus, carry_modulus, f_last_grouping_luts,
        gpu_memory_allocated);

    // Generate the indexes to switch between luts within the pbs
    uint64_t lut_indexes_size = num_radix_blocks * sizeof(Torus);

    Torus *h_lut_indexes = luts_array_first_step->h_lut_indexes;
    for (int index = 0; index < num_radix_blocks; index++) {
      uint32_t grouping_index = index / grouping_size;
      bool is_in_first_grouping = (grouping_index == 0);
      uint32_t index_in_grouping = index % grouping_size;
      bool is_last_index = (index == (num_radix_blocks - 1));
      if (is_last_index) {
        if (num_radix_blocks == 1) {
          h_lut_indexes[index] = 2 * grouping_size;
        } else {
          h_lut_indexes[index] = 2;
        }
      } else if (is_in_first_grouping) {
        h_lut_indexes[index] = index_in_grouping;
      } else {
        h_lut_indexes[index] = index_in_grouping + grouping_size;
      }
    }

    // copy the indexes to the gpu
    Torus *lut_indexes = luts_array_first_step->get_lut_indexes(0, 0);
    cuda_memcpy_with_size_tracking_async_to_gpu(
        lut_indexes, h_lut_indexes, lut_indexes_size, streams.stream(0),
        streams.gpu_index(0), allocate_gpu_memory);
    // Do I need to do something else for the multi-gpu?
    auto active_streams = streams.active_gpu_subset(num_radix_blocks);
    luts_array_first_step->broadcast_lut(active_streams);
  };
  void release(CudaStreams streams) {

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   shifted_blocks_and_states,
                                   gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   shifted_blocks, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   block_states, gpu_memory_allocated);

    luts_array_first_step->release(streams);
    delete luts_array_first_step;
    delete shifted_blocks_and_states;
    delete shifted_blocks;
    delete block_states;
  };
};

// compute_propagation simulator and group carries
template <typename Torus> struct int_prop_simu_group_carries_memory {
  CudaRadixCiphertextFFI *propagation_cum_sums;
  CudaRadixCiphertextFFI *simulators;
  CudaRadixCiphertextFFI *prepared_blocks;
  CudaRadixCiphertextFFI *grouping_pgns;
  CudaRadixCiphertextFFI *resolved_carries;

  Torus *scalar_array_cum_sum;
  Torus *h_scalar_array_cum_sum;

  int_radix_lut<Torus> *luts_array_second_step;

  int_seq_group_prop_memory<Torus> *seq_group_prop_mem;
  int_hs_group_prop_memory<Torus> *hs_group_prop_mem;

  uint32_t group_size;
  bool use_sequential_algorithm_to_resolve_group_carries;
  bool gpu_memory_allocated;

  int_prop_simu_group_carries_memory(
      CudaStreams streams, int_radix_params params, uint32_t num_radix_blocks,
      uint32_t grouping_size, uint32_t num_groups, bool allocate_gpu_memory,
      uint64_t &size_tracker) {

    gpu_memory_allocated = allocate_gpu_memory;
    auto glwe_dimension = params.glwe_dimension;
    auto polynomial_size = params.polynomial_size;
    auto message_modulus = params.message_modulus;
    auto carry_modulus = params.carry_modulus;
    auto big_lwe_size = (polynomial_size * glwe_dimension + 1);
    auto big_lwe_size_bytes = big_lwe_size * sizeof(Torus);

    uint32_t block_modulus = message_modulus * carry_modulus;
    uint32_t num_bits_in_block = std::log2(block_modulus);

    group_size = grouping_size;

    propagation_cum_sums = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), propagation_cum_sums,
        num_radix_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    simulators = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), simulators, num_radix_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    prepared_blocks = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), prepared_blocks,
        num_radix_blocks + 1, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    resolved_carries = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), resolved_carries,
        num_groups + 1, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    grouping_pgns = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), grouping_pgns, num_groups,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    scalar_array_cum_sum = (Torus *)cuda_malloc_with_size_tracking_async(
        num_radix_blocks * sizeof(Torus), streams.stream(0),
        streams.gpu_index(0), size_tracker, allocate_gpu_memory);
    cuda_memset_with_size_tracking_async(
        scalar_array_cum_sum, 0, num_radix_blocks * sizeof(Torus),
        streams.stream(0), streams.gpu_index(0), allocate_gpu_memory);
    h_scalar_array_cum_sum = new Torus[num_radix_blocks]();

    // create lut objects for step 2
    uint64_t lut_indexes_size = num_radix_blocks * sizeof(Torus);
    uint32_t num_carry_to_resolve = num_groups - 1;
    uint32_t saturated_sub =
        ((num_carry_to_resolve > 1) ? num_carry_to_resolve - 1 : 0);
    uint32_t sequential_depth = saturated_sub / (grouping_size - 1);
    uint32_t hillis_steel_depth;

    if (num_carry_to_resolve == 0) {
      hillis_steel_depth = 0;
    } else {
      hillis_steel_depth = std::ceil(std::log2(num_carry_to_resolve));
    }

    use_sequential_algorithm_to_resolve_group_carries =
        sequential_depth <= hillis_steel_depth;
    uint32_t num_extra_luts = 0;
    if (use_sequential_algorithm_to_resolve_group_carries) {
      num_extra_luts = (grouping_size - 1);
    } else {
      num_extra_luts = 1;
    }

    uint32_t num_luts_second_step = 2 * grouping_size + num_extra_luts;
    luts_array_second_step = new int_radix_lut<Torus>(
        streams, params, num_luts_second_step, num_radix_blocks,
        allocate_gpu_memory, size_tracker);

    // luts for first group inner propagation
    for (int lut_id = 0; lut_id < grouping_size - 1; lut_id++) {
      auto f_first_grouping_inner_propagation =
          [lut_id](Torus propa_cum_sum_block) -> Torus {
        uint64_t carry = (propa_cum_sum_block >> lut_id) & 1;

        if (carry != 0) {
          return 2ull; // Generates Carry
        } else {
          return 0ull; // Does not generate carry
        }
      };

      generate_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0),
          luts_array_second_step->get_lut(0, lut_id),
          luts_array_second_step->get_degree(lut_id),
          luts_array_second_step->get_max_degree(lut_id), glwe_dimension,
          polynomial_size, message_modulus, carry_modulus,
          f_first_grouping_inner_propagation, gpu_memory_allocated);
    }

    auto f_first_grouping_outer_propagation =
        [num_bits_in_block](Torus block) -> Torus {
      return (block >> (num_bits_in_block - 1)) & 1;
    };

    int lut_id = grouping_size - 1;
    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0),
        luts_array_second_step->get_lut(0, lut_id),
        luts_array_second_step->get_degree(lut_id),
        luts_array_second_step->get_max_degree(lut_id), glwe_dimension,
        polynomial_size, message_modulus, carry_modulus,
        f_first_grouping_outer_propagation, gpu_memory_allocated);

    // for other groupings inner propagation
    for (int index = 0; index < grouping_size; index++) {
      uint32_t lut_id = index + grouping_size;

      auto f_other_groupings_inner_propagation =
          [index](Torus propa_cum_sum_block) -> Torus {
        uint64_t mask = (2 << index) - 1;
        if (propa_cum_sum_block >= (2 << index)) {
          return 2ull; // Generates
        } else if ((propa_cum_sum_block & mask) == mask) {
          return 1ull; // Propagate
        } else {
          return 0ull; // Nothing
        }
      };

      generate_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0),
          luts_array_second_step->get_lut(0, lut_id),
          luts_array_second_step->get_degree(lut_id),
          luts_array_second_step->get_max_degree(lut_id), glwe_dimension,
          polynomial_size, message_modulus, carry_modulus,
          f_other_groupings_inner_propagation, gpu_memory_allocated);
    }

    if (use_sequential_algorithm_to_resolve_group_carries) {
      for (int index = 0; index < grouping_size - 1; index++) {
        uint32_t lut_id = index + 2 * grouping_size;

        auto f_group_propagation = [index, block_modulus,
                                    num_bits_in_block](Torus block) -> Torus {
          if (block == (block_modulus - 1)) {
            return 0ull;
          } else {
            return ((UINT64_MAX << index) % (1ull << (num_bits_in_block + 1)));
          }
        };

        generate_device_accumulator<Torus>(
            streams.stream(0), streams.gpu_index(0),
            luts_array_second_step->get_lut(0, lut_id),
            luts_array_second_step->get_degree(lut_id),
            luts_array_second_step->get_max_degree(lut_id), glwe_dimension,
            polynomial_size, message_modulus, carry_modulus,
            f_group_propagation, gpu_memory_allocated);
      }
    } else {
      uint32_t lut_id = 2 * grouping_size;
      auto f_group_propagation = [block_modulus](Torus block) {
        if (block == (block_modulus - 1)) {
          return 2ull;
        } else {
          return UINT64_MAX % (block_modulus * 2ull);
        }
      };

      generate_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0),
          luts_array_second_step->get_lut(0, lut_id),
          luts_array_second_step->get_degree(lut_id),
          luts_array_second_step->get_max_degree(lut_id), glwe_dimension,
          polynomial_size, message_modulus, carry_modulus, f_group_propagation,
          gpu_memory_allocated);
    }

    Torus *h_second_lut_indexes = (Torus *)malloc(lut_indexes_size);

    for (int index = 0; index < num_radix_blocks; index++) {
      uint32_t grouping_index = index / grouping_size;
      bool is_in_first_grouping = (grouping_index == 0);
      uint32_t index_in_grouping = index % grouping_size;

      if (is_in_first_grouping) {
        h_second_lut_indexes[index] = index_in_grouping;
      } else if (index_in_grouping == (grouping_size - 1)) {
        if (use_sequential_algorithm_to_resolve_group_carries) {
          int inner_index = (grouping_index - 1) % (grouping_size - 1);
          h_second_lut_indexes[index] = inner_index + 2 * grouping_size;
        } else {
          h_second_lut_indexes[index] = 2 * grouping_size;
        }
      } else {
        h_second_lut_indexes[index] = index_in_grouping + grouping_size;
      }

      bool may_have_its_padding_bit_set =
          !is_in_first_grouping && (index_in_grouping == grouping_size - 1);

      if (may_have_its_padding_bit_set) {
        if (use_sequential_algorithm_to_resolve_group_carries) {
          h_scalar_array_cum_sum[index] =
              1 << ((grouping_index - 1) % (grouping_size - 1));
        } else {
          h_scalar_array_cum_sum[index] = 1;
        }
      } else {
        h_scalar_array_cum_sum[index] = 0;
      }
    }

    // copy the indexes to the gpu
    Torus *second_lut_indexes = luts_array_second_step->get_lut_indexes(0, 0);
    cuda_memcpy_with_size_tracking_async_to_gpu(
        second_lut_indexes, h_second_lut_indexes, lut_indexes_size,
        streams.stream(0), streams.gpu_index(0), allocate_gpu_memory);

    cuda_memcpy_with_size_tracking_async_to_gpu(
        scalar_array_cum_sum, h_scalar_array_cum_sum,
        num_radix_blocks * sizeof(Torus), streams.stream(0),
        streams.gpu_index(0), allocate_gpu_memory);
    auto active_streams = streams.active_gpu_subset(num_radix_blocks);
    luts_array_second_step->broadcast_lut(active_streams);

    if (use_sequential_algorithm_to_resolve_group_carries) {

      seq_group_prop_mem = new int_seq_group_prop_memory<Torus>(
          streams, params, grouping_size, big_lwe_size_bytes,
          allocate_gpu_memory, size_tracker);

    } else {
      hs_group_prop_mem = new int_hs_group_prop_memory<Torus>(
          streams, params, num_groups, big_lwe_size_bytes, allocate_gpu_memory,
          size_tracker);
    }

    free(h_second_lut_indexes);
  };

  // needed for the division to update the lut indexes
  void update_lut_indexes(CudaStreams streams, Torus *new_lut_indexes,
                          Torus *new_scalars, uint32_t new_num_blocks) {
    Torus *lut_indexes = luts_array_second_step->get_lut_indexes(0, 0);
    cuda_memcpy_with_size_tracking_async_gpu_to_gpu(
        lut_indexes, new_lut_indexes, new_num_blocks * sizeof(Torus),
        streams.stream(0), streams.gpu_index(0), gpu_memory_allocated);
    auto new_active_streams = streams.active_gpu_subset(new_num_blocks);
    // We just need to update the lut indexes so we use false here
    luts_array_second_step->broadcast_lut(new_active_streams, false);

    cuda_memcpy_with_size_tracking_async_gpu_to_gpu(
        scalar_array_cum_sum, new_scalars, new_num_blocks * sizeof(Torus),
        streams.stream(0), streams.gpu_index(0), gpu_memory_allocated);
  }

  void release(CudaStreams streams) {
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   propagation_cum_sums, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   simulators, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   grouping_pgns, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   prepared_blocks, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   resolved_carries, gpu_memory_allocated);
    cuda_drop_with_size_tracking_async(scalar_array_cum_sum, streams.stream(0),
                                       streams.gpu_index(0),
                                       gpu_memory_allocated);
    luts_array_second_step->release(streams);

    if (use_sequential_algorithm_to_resolve_group_carries) {
      seq_group_prop_mem->release(streams);
      delete seq_group_prop_mem;
    } else {
      hs_group_prop_mem->release(streams);
      delete hs_group_prop_mem;
    }

    delete propagation_cum_sums;
    delete simulators;
    delete grouping_pgns;
    delete prepared_blocks;
    delete resolved_carries;
    delete luts_array_second_step;
    delete[] h_scalar_array_cum_sum;
  };
};

template <typename Torus> struct int_sc_prop_memory {
  uint32_t num_many_lut;
  uint32_t lut_stride;

  uint32_t num_groups;
  CudaRadixCiphertextFFI *output_flag;
  CudaRadixCiphertextFFI *last_lhs;
  CudaRadixCiphertextFFI *last_rhs;
  int_radix_lut<Torus> *lut_message_extract;

  int_radix_lut<Torus> *lut_overflow_flag_prep;

  int_shifted_blocks_and_states_memory<Torus> *shifted_blocks_state_mem;
  int_prop_simu_group_carries_memory<Torus> *prop_simu_group_carries_mem;

  int_radix_params params;
  uint32_t requested_flag;
  bool gpu_memory_allocated;

  int_sc_prop_memory(CudaStreams streams, int_radix_params params,
                     uint32_t num_radix_blocks, uint32_t requested_flag_in,
                     bool allocate_gpu_memory, uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;
    this->params = params;
    auto glwe_dimension = params.glwe_dimension;
    auto polynomial_size = params.polynomial_size;
    auto message_modulus = params.message_modulus;
    auto carry_modulus = params.carry_modulus;
    requested_flag = requested_flag_in;
    // for compute shifted blocks and block states
    uint32_t block_modulus = message_modulus * carry_modulus;
    uint32_t num_bits_in_block = std::log2(block_modulus);
    uint32_t grouping_size = num_bits_in_block;
    num_groups = (num_radix_blocks + grouping_size - 1) / grouping_size;

    num_many_lut = 2; // many luts apply 2 luts
    uint32_t box_size = polynomial_size / block_modulus;
    lut_stride = (block_modulus / num_many_lut) * box_size;

    shifted_blocks_state_mem = new int_shifted_blocks_and_states_memory<Torus>(
        streams, params, num_radix_blocks, num_many_lut, grouping_size,
        allocate_gpu_memory, size_tracker);

    prop_simu_group_carries_mem = new int_prop_simu_group_carries_memory<Torus>(
        streams, params, num_radix_blocks, grouping_size, num_groups,
        allocate_gpu_memory, size_tracker);

    //  Step 3 elements
    int num_luts_message_extract =
        requested_flag == outputFlag::FLAG_NONE ? 1 : 2;
    lut_message_extract = new int_radix_lut<Torus>(
        streams, params, num_luts_message_extract, num_radix_blocks + 1,
        allocate_gpu_memory, size_tracker);
    // lut for the first block in the first grouping
    auto f_message_extract = [message_modulus](Torus block) -> Torus {
      return (block >> 1) % message_modulus;
    };

    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0),
        lut_message_extract->get_lut(0, 0), lut_message_extract->get_degree(0),
        lut_message_extract->get_max_degree(0), glwe_dimension, polynomial_size,
        message_modulus, carry_modulus, f_message_extract,
        gpu_memory_allocated);

    // This store a single block that with be used to store the overflow or
    // carry results
    output_flag = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), output_flag,
        num_radix_blocks + 1, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    if (requested_flag == outputFlag::FLAG_OVERFLOW) {
      last_lhs = new CudaRadixCiphertextFFI;
      last_rhs = new CudaRadixCiphertextFFI;
      create_zero_radix_ciphertext_async<Torus>(
          streams.stream(0), streams.gpu_index(0), last_lhs, 1,
          params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
      create_zero_radix_ciphertext_async<Torus>(
          streams.stream(0), streams.gpu_index(0), last_rhs, 1,
          params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

      // For step 1 overflow should be enable only if flag overflow
      uint32_t num_bits_in_message = std::log2(message_modulus);
      lut_overflow_flag_prep = new int_radix_lut<Torus>(
          streams, params, 1, 1, allocate_gpu_memory, size_tracker);

      auto f_overflow_fp = [num_bits_in_message](Torus lhs,
                                                 Torus rhs) -> Torus {
        Torus mask = (1 << (num_bits_in_message - 1)) - 1;
        Torus lhs_except_last_bit = lhs & mask;
        Torus rhs_except_last_bit = rhs & mask;
        Torus input_carry1 = 1;
        Torus input_carry2 = 0;

        Torus output_carry1 =
            ((lhs + rhs + input_carry1) >> num_bits_in_message) & 1;
        Torus output_carry2 =
            ((lhs + rhs + input_carry2) >> num_bits_in_message) & 1;
        Torus input_carry_last_bit1 =
            ((lhs_except_last_bit + rhs_except_last_bit + input_carry1) >>
             (num_bits_in_message - 1)) &
            1;
        Torus input_carry_last_bit2 =
            ((lhs_except_last_bit + rhs_except_last_bit + input_carry2) >>
             (num_bits_in_message - 1)) &
            1;

        Torus output1 = (Torus)(input_carry_last_bit1 != output_carry1);
        Torus output2 = (Torus)(input_carry_last_bit2 != output_carry2);

        return output1 << 3 | output2 << 2;
      };

      generate_device_accumulator_bivariate<Torus>(
          streams.stream(0), streams.gpu_index(0),
          lut_overflow_flag_prep->get_lut(0, 0),
          lut_overflow_flag_prep->get_degree(0),
          lut_overflow_flag_prep->get_max_degree(0), glwe_dimension,
          polynomial_size, message_modulus, carry_modulus, f_overflow_fp,
          gpu_memory_allocated);

      auto active_streams = streams.active_gpu_subset(1);
      lut_overflow_flag_prep->broadcast_lut(active_streams);
    }

    // For the final cleanup in case of overflow or carry (it seems that I can)
    // It seems that this lut could be apply together with the other one but for
    // now we won't do it
    if (requested_flag == outputFlag::FLAG_OVERFLOW) { // Overflow case
      auto f_overflow_last = [num_radix_blocks,
                              requested_flag_in](Torus block) -> Torus {
        uint32_t position = (num_radix_blocks == 1 &&
                             requested_flag_in == outputFlag::FLAG_OVERFLOW)
                                ? 0
                                : 1;
        Torus input_carry = (block >> position) & 1;
        Torus does_overflow_if_carry_is_1 = (block >> 3) & 1;
        Torus does_overflow_if_carry_is_0 = (block >> 2) & 1;
        if (input_carry == outputFlag::FLAG_OVERFLOW) {
          return does_overflow_if_carry_is_1;
        } else {
          return does_overflow_if_carry_is_0;
        }
      };

      generate_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0),
          lut_message_extract->get_lut(0, 1),
          lut_message_extract->get_degree(1),
          lut_message_extract->get_max_degree(1), glwe_dimension,
          polynomial_size, message_modulus, carry_modulus, f_overflow_last,
          gpu_memory_allocated);

      Torus *h_lut_indexes = lut_message_extract->h_lut_indexes;
      for (int index = 0; index < num_radix_blocks + 1; index++) {
        if (index < num_radix_blocks) {
          h_lut_indexes[index] = 0;
        } else {
          h_lut_indexes[index] = 1;
        }
      }
      cuda_memcpy_with_size_tracking_async_to_gpu(
          lut_message_extract->get_lut_indexes(0, 0), h_lut_indexes,
          (num_radix_blocks + 1) * sizeof(Torus), streams.stream(0),
          streams.gpu_index(0), allocate_gpu_memory);
    }
    if (requested_flag == outputFlag::FLAG_CARRY) { // Carry case

      auto f_carry_last = [](Torus block) -> Torus {
        return ((block >> 2) & 1);
      };

      generate_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0),
          lut_message_extract->get_lut(0, 1),
          lut_message_extract->get_degree(1),
          lut_message_extract->get_max_degree(1), glwe_dimension,
          polynomial_size, message_modulus, carry_modulus, f_carry_last,
          gpu_memory_allocated);

      Torus *h_lut_indexes = lut_message_extract->h_lut_indexes;
      for (int index = 0; index < num_radix_blocks + 1; index++) {
        if (index < num_radix_blocks) {
          h_lut_indexes[index] = 0;
        } else {
          h_lut_indexes[index] = 1;
        }
      }
      cuda_memcpy_with_size_tracking_async_to_gpu(
          lut_message_extract->get_lut_indexes(0, 0), h_lut_indexes,
          (num_radix_blocks + 1) * sizeof(Torus), streams.stream(0),
          streams.gpu_index(0), allocate_gpu_memory);
    }
    auto active_streams = streams.active_gpu_subset(num_radix_blocks + 1);
    lut_message_extract->broadcast_lut(active_streams);
  };

  void release(CudaStreams streams) {

    shifted_blocks_state_mem->release(streams);
    prop_simu_group_carries_mem->release(streams);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   output_flag, gpu_memory_allocated);
    lut_message_extract->release(streams);
    delete shifted_blocks_state_mem;
    delete prop_simu_group_carries_mem;
    delete output_flag;
    delete lut_message_extract;

    if (requested_flag == outputFlag::FLAG_OVERFLOW) { // In case of overflow
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     last_lhs, gpu_memory_allocated);
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     last_rhs, gpu_memory_allocated);
      lut_overflow_flag_prep->release(streams);
      delete lut_overflow_flag_prep;
      delete last_lhs;
      delete last_rhs;
    }
  };
};

template <typename Torus> struct int_shifted_blocks_and_borrow_states_memory {
  CudaRadixCiphertextFFI *shifted_blocks_and_borrow_states;
  CudaRadixCiphertextFFI *shifted_blocks;
  CudaRadixCiphertextFFI *borrow_states;

  int_radix_lut<Torus> *luts_array_first_step;
  bool gpu_memory_allocated;

  int_shifted_blocks_and_borrow_states_memory(
      CudaStreams streams, int_radix_params params, uint32_t num_radix_blocks,
      uint32_t num_many_lut, uint32_t grouping_size, bool allocate_gpu_memory,
      uint64_t &size_tracker) {

    gpu_memory_allocated = allocate_gpu_memory;
    auto glwe_dimension = params.glwe_dimension;
    auto polynomial_size = params.polynomial_size;
    auto message_modulus = params.message_modulus;
    auto carry_modulus = params.carry_modulus;

    shifted_blocks_and_borrow_states = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0),
        shifted_blocks_and_borrow_states, num_radix_blocks * num_many_lut,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    shifted_blocks = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), shifted_blocks,
        num_radix_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    borrow_states = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), borrow_states,
        num_radix_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    uint32_t num_luts_first_step = 2 * grouping_size + 1;

    luts_array_first_step = new int_radix_lut<Torus>(
        streams, params, num_luts_first_step, num_radix_blocks, num_many_lut,
        allocate_gpu_memory, size_tracker);

    auto f_shift_block = [message_modulus](Torus block) -> Torus {
      uint64_t overflow_guard = message_modulus;
      uint64_t block_mod = block % message_modulus;
      return (overflow_guard | block_mod) << 1;
    };

    auto f_first_block_state = [message_modulus](Torus block) -> Torus {
      if (block < message_modulus)
        return 1; // Borrows
      else {
        return 0; // Nothing
      }
    };
    std::vector<std::function<Torus(Torus)>> f_first_grouping_luts = {
        f_first_block_state, f_shift_block};

    auto first_block_lut = luts_array_first_step->get_lut(0, 0);
    auto first_block_lut_degrees = luts_array_first_step->get_degree(0);
    auto first_block_lut_max_degree = luts_array_first_step->get_max_degree(0);

    generate_many_lut_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0), first_block_lut,
        first_block_lut_degrees, first_block_lut_max_degree, glwe_dimension,
        polynomial_size, message_modulus, carry_modulus, f_first_grouping_luts,
        gpu_memory_allocated);

    // luts for other blocks of the first grouping
    for (int lut_id = 1; lut_id < grouping_size; lut_id++) {
      auto f_state = [message_modulus, lut_id](Torus block) -> Torus {
        uint64_t r = 0;
        if (block < message_modulus) {
          r = 2; // Borrows
        } else if (block == message_modulus) {
          r = 1; // Propagates a borrow
        } else {
          r = 0; // Does not borrow
        }
        return r << (lut_id - 1);
      };
      std::vector<std::function<Torus(Torus)>> f_grouping_luts = {
          f_state, f_shift_block};
      auto lut = luts_array_first_step->get_lut(0, lut_id);
      auto lut_degrees = luts_array_first_step->get_degree(lut_id);
      auto lut_max_degree = luts_array_first_step->get_max_degree(lut_id);
      generate_many_lut_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0), lut, lut_degrees,
          lut_max_degree, glwe_dimension, polynomial_size, message_modulus,
          carry_modulus, f_grouping_luts, gpu_memory_allocated);
    }

    // luts for the rest of groupings (except for the last block)
    for (int i = 0; i < grouping_size; i++) {
      uint32_t lut_id = i + grouping_size;
      auto f_state = [message_modulus, i](Torus block) -> Torus {
        uint64_t r = 0;
        if (block < message_modulus) {
          r = 2; // Generates borrow
        } else if (block == message_modulus) {
          r = 1; // Propagates a borrow
        } else {
          r = 0; // Does not borrow
        }
        return r << i;
      };
      std::vector<std::function<Torus(Torus)>> f_grouping_luts = {
          f_state, f_shift_block};

      auto lut = luts_array_first_step->get_lut(0, lut_id);
      auto lut_degrees = luts_array_first_step->get_degree(lut_id);
      auto lut_max_degree = luts_array_first_step->get_max_degree(lut_id);
      generate_many_lut_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0), lut, lut_degrees,
          lut_max_degree, glwe_dimension, polynomial_size, message_modulus,
          carry_modulus, f_grouping_luts, gpu_memory_allocated);
    }

    auto f_last_block_state = [message_modulus](Torus block) -> Torus {
      if (block < message_modulus)
        return 2 << 1; // Generates a borrow
      else
        return 0; // Nothing
    };

    uint32_t lut_id = num_luts_first_step - 1; // The last lut of the first step

    auto last_block_lut = luts_array_first_step->get_lut(0, lut_id);
    auto last_block_lut_degrees = luts_array_first_step->get_degree(lut_id);
    auto last_block_lut_max_degree =
        luts_array_first_step->get_max_degree(lut_id);

    std::vector<std::function<Torus(Torus)>> f_last_grouping_luts = {
        f_last_block_state, f_shift_block};

    generate_many_lut_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0), last_block_lut,
        last_block_lut_degrees, last_block_lut_max_degree, glwe_dimension,
        polynomial_size, message_modulus, carry_modulus, f_last_grouping_luts,
        gpu_memory_allocated);

    // Generate the indexes to switch between luts within the pbs
    uint64_t lut_indexes_size = num_radix_blocks * sizeof(Torus);
    Torus *h_lut_indexes = luts_array_first_step->h_lut_indexes;

    for (int index = 0; index < num_radix_blocks; index++) {
      uint32_t grouping_index = index / grouping_size;
      bool is_in_first_grouping = (grouping_index == 0);
      uint32_t index_in_grouping = index % grouping_size;
      bool is_last_index = (index == (num_radix_blocks - 1));
      if (is_last_index) {
        if (num_radix_blocks == 1) {
          h_lut_indexes[index] = 2 * grouping_size;
        } else {
          h_lut_indexes[index] = 2;
        }
      } else if (is_in_first_grouping) {
        h_lut_indexes[index] = index_in_grouping;
      } else {
        h_lut_indexes[index] = index_in_grouping + grouping_size;
      }
    }
    // copy the indexes to the gpu
    Torus *lut_indexes = luts_array_first_step->get_lut_indexes(0, 0);
    cuda_memcpy_with_size_tracking_async_to_gpu(
        lut_indexes, h_lut_indexes, lut_indexes_size, streams.stream(0),
        streams.gpu_index(0), allocate_gpu_memory);
    // Do I need to do something else for the multi-gpu?
    auto active_streams = streams.active_gpu_subset(num_radix_blocks);
    luts_array_first_step->broadcast_lut(active_streams);
  };

  // needed for the division to update the lut indexes
  void update_lut_indexes(CudaStreams streams, Torus *new_lut_indexes,
                          uint32_t new_num_blocks) {
    Torus *lut_indexes = luts_array_first_step->get_lut_indexes(0, 0);
    cuda_memcpy_with_size_tracking_async_gpu_to_gpu(
        lut_indexes, new_lut_indexes, new_num_blocks * sizeof(Torus),
        streams.stream(0), streams.gpu_index(0), gpu_memory_allocated);
    auto new_active_streams = streams.active_gpu_subset(new_num_blocks);
    // We just need to update the lut indexes so we use false here
    luts_array_first_step->broadcast_lut(new_active_streams, false);
  }
  void release(CudaStreams streams) {

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   shifted_blocks_and_borrow_states,
                                   gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   shifted_blocks, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   borrow_states, gpu_memory_allocated);

    luts_array_first_step->release(streams);
    delete luts_array_first_step;
    delete shifted_blocks_and_borrow_states;
    delete shifted_blocks;
    delete borrow_states;
  };
};

template <typename Torus> struct int_borrow_prop_memory {
  uint32_t num_many_lut;
  uint32_t lut_stride;

  uint32_t group_size;
  uint32_t num_groups;
  CudaRadixCiphertextFFI *overflow_block;

  int_radix_lut<Torus> *lut_message_extract;
  int_radix_lut<Torus> *lut_borrow_flag;

  int_shifted_blocks_and_borrow_states_memory<Torus>
      *shifted_blocks_borrow_state_mem;
  int_prop_simu_group_carries_memory<Torus> *prop_simu_group_carries_mem;

  int_radix_params params;

  CudaStreams active_streams;
  CudaStreams sub_streams_1;
  CudaStreams sub_streams_2;

  cudaEvent_t *incoming_events;
  cudaEvent_t *outgoing_events1;
  cudaEvent_t *outgoing_events2;

  uint32_t compute_overflow;
  bool gpu_memory_allocated;
  int_borrow_prop_memory(CudaStreams streams, int_radix_params params,
                         uint32_t num_radix_blocks,
                         uint32_t compute_overflow_in, bool allocate_gpu_memory,
                         uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;
    this->params = params;
    auto glwe_dimension = params.glwe_dimension;
    auto polynomial_size = params.polynomial_size;
    auto message_modulus = params.message_modulus;
    auto carry_modulus = params.carry_modulus;
    compute_overflow = compute_overflow_in;
    // for compute shifted blocks and block states
    uint32_t block_modulus = message_modulus * carry_modulus;
    uint32_t num_bits_in_block = std::log2(block_modulus);
    uint32_t grouping_size = num_bits_in_block;
    group_size = grouping_size;
    num_groups = (num_radix_blocks + grouping_size - 1) / grouping_size;

    num_many_lut = 2; // many luts apply 2 luts
    uint32_t box_size = polynomial_size / block_modulus;
    lut_stride = (block_modulus / num_many_lut) * box_size;

    shifted_blocks_borrow_state_mem =
        new int_shifted_blocks_and_borrow_states_memory<Torus>(
            streams, params, num_radix_blocks, num_many_lut, grouping_size,
            allocate_gpu_memory, size_tracker);

    prop_simu_group_carries_mem = new int_prop_simu_group_carries_memory<Torus>(
        streams, params, num_radix_blocks, grouping_size, num_groups,
        allocate_gpu_memory, size_tracker);

    overflow_block = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), overflow_block, 1,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

    lut_message_extract =
        new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                 allocate_gpu_memory, size_tracker);
    // lut for the first block in the first grouping
    auto f_message_extract = [message_modulus](Torus block) -> Torus {
      return (block >> 1) % message_modulus;
    };

    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0),
        lut_message_extract->get_lut(0, 0), lut_message_extract->get_degree(0),
        lut_message_extract->get_max_degree(0), glwe_dimension, polynomial_size,
        message_modulus, carry_modulus, f_message_extract,
        gpu_memory_allocated);
    active_streams = streams.active_gpu_subset(num_radix_blocks);

    lut_message_extract->broadcast_lut(active_streams);

    if (compute_overflow) {
      lut_borrow_flag =
          new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                   allocate_gpu_memory, size_tracker);
      // lut for the first block in the first grouping
      auto f_borrow_flag = [](Torus block) -> Torus {
        return ((block >> 2) & 1);
      };

      generate_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0),
          lut_borrow_flag->get_lut(0, 0), lut_borrow_flag->get_degree(0),
          lut_borrow_flag->get_max_degree(0), glwe_dimension, polynomial_size,
          message_modulus, carry_modulus, f_borrow_flag, gpu_memory_allocated);
      lut_borrow_flag->broadcast_lut(active_streams);
    }

    active_streams = streams.active_gpu_subset(num_radix_blocks);
    sub_streams_1.create_on_same_gpus(active_streams);
    sub_streams_2.create_on_same_gpus(active_streams);

    incoming_events =
        (cudaEvent_t *)malloc(active_streams.count() * sizeof(cudaEvent_t));
    outgoing_events1 =
        (cudaEvent_t *)malloc(active_streams.count() * sizeof(cudaEvent_t));
    outgoing_events2 =
        (cudaEvent_t *)malloc(active_streams.count() * sizeof(cudaEvent_t));
    for (uint j = 0; j < active_streams.count(); j++) {
      incoming_events[j] = cuda_create_event(active_streams.gpu_index(j));
      outgoing_events1[j] = cuda_create_event(active_streams.gpu_index(j));
      outgoing_events2[j] = cuda_create_event(active_streams.gpu_index(j));
    }
  };

  // needed for the division to update the lut indexes
  void update_lut_indexes(CudaStreams streams, Torus *first_indexes_for_div,
                          Torus *second_indexes_for_div, Torus *scalars_for_div,
                          uint32_t new_num_blocks) {
    shifted_blocks_borrow_state_mem->update_lut_indexes(
        streams, first_indexes_for_div, new_num_blocks);
    prop_simu_group_carries_mem->update_lut_indexes(
        streams, second_indexes_for_div, scalars_for_div, new_num_blocks);
  }
  void release(CudaStreams streams) {

    shifted_blocks_borrow_state_mem->release(streams);
    prop_simu_group_carries_mem->release(streams);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   overflow_block, gpu_memory_allocated);

    lut_message_extract->release(streams);
    delete lut_message_extract;
    delete overflow_block;
    if (compute_overflow) {
      lut_borrow_flag->release(streams);
      delete lut_borrow_flag;
    }

    // The substreams have to be synchronized before destroying events
    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));

    // release events
    for (uint j = 0; j < active_streams.count(); j++) {
      cuda_event_destroy(incoming_events[j], active_streams.gpu_index(j));
      cuda_event_destroy(outgoing_events1[j], active_streams.gpu_index(j));
      cuda_event_destroy(outgoing_events2[j], active_streams.gpu_index(j));
    }
    free(incoming_events);
    free(outgoing_events1);
    free(outgoing_events2);

    sub_streams_1.release();
    sub_streams_2.release();
  };
};

template <typename Torus> struct int_zero_out_if_buffer {

  int_radix_params params;

  CudaRadixCiphertextFFI *tmp;

  bool gpu_memory_allocated;

  int_zero_out_if_buffer(CudaStreams streams, int_radix_params params,
                         uint32_t num_radix_blocks, bool allocate_gpu_memory,
                         uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;
    this->params = params;
    auto active_streams = streams.active_gpu_subset(num_radix_blocks);

    tmp = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp, num_radix_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
  }
  void release(CudaStreams streams) {
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0), tmp,
                                   gpu_memory_allocated);
    delete tmp;
    tmp = nullptr;
  }
};

template <typename Torus> struct int_mul_memory {
  CudaRadixCiphertextFFI *vector_result_sb;
  CudaRadixCiphertextFFI *block_mul_res;
  CudaRadixCiphertextFFI *small_lwe_vector;

  int_radix_lut<Torus> *luts_array; // lsb msb
  int_radix_lut<Torus> *zero_out_predicate_lut;

  int_sum_ciphertexts_vec_memory<Torus> *sum_ciphertexts_mem;
  int_sc_prop_memory<Torus> *sc_prop_mem;
  int_zero_out_if_buffer<Torus> *zero_out_mem;

  int_radix_params params;
  bool boolean_mul = false;
  bool gpu_memory_allocated;

  int_mul_memory(CudaStreams streams, int_radix_params params,
                 bool const is_boolean_left, bool const is_boolean_right,
                 uint32_t num_radix_blocks, bool allocate_gpu_memory,
                 uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;
    this->boolean_mul = is_boolean_left || is_boolean_right;
    this->params = params;

    if (boolean_mul) {
      auto zero_out_predicate_lut_f = [](Torus block,
                                         Torus condition) -> Torus {
        if (condition == 0)
          return 0;
        else
          return block;
      };
      zero_out_predicate_lut =
          new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                   allocate_gpu_memory, size_tracker);
      generate_device_accumulator_bivariate<Torus>(
          streams.stream(0), streams.gpu_index(0),
          zero_out_predicate_lut->get_lut(0, 0),
          zero_out_predicate_lut->get_degree(0),
          zero_out_predicate_lut->get_max_degree(0), params.glwe_dimension,
          params.polynomial_size, params.message_modulus, params.carry_modulus,
          zero_out_predicate_lut_f, gpu_memory_allocated);

      auto active_streams = streams.active_gpu_subset(num_radix_blocks);
      zero_out_predicate_lut->broadcast_lut(active_streams);

      zero_out_mem = new int_zero_out_if_buffer<Torus>(
          streams, params, num_radix_blocks, allocate_gpu_memory, size_tracker);

      return;
    }

    auto glwe_dimension = params.glwe_dimension;
    auto polynomial_size = params.polynomial_size;
    auto message_modulus = params.message_modulus;
    auto carry_modulus = params.carry_modulus;

    // 'vector_result_lsb' contains blocks from all possible shifts of
    // radix_lwe_left excluding zero ciphertext blocks
    int lsb_vector_block_count = num_radix_blocks * (num_radix_blocks + 1) / 2;

    // 'vector_result_msb' contains blocks from all possible shifts of
    // radix_lwe_left except the last blocks of each shift
    int msb_vector_block_count = num_radix_blocks * (num_radix_blocks - 1) / 2;

    int total_block_count = num_radix_blocks * num_radix_blocks;

    // allocate memory for intermediate buffers
    vector_result_sb = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), vector_result_sb,
        2 * total_block_count, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    block_mul_res = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), block_mul_res,
        2 * total_block_count, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    small_lwe_vector = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), small_lwe_vector,
        2 * total_block_count, params.small_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    // create int_radix_lut objects for lsb, msb, message, carry
    // luts_array -> lut = {lsb_acc, msb_acc}
    luts_array = new int_radix_lut<Torus>(streams, params, 2, total_block_count,
                                          allocate_gpu_memory, size_tracker);
    auto lsb_acc = luts_array->get_lut(0, 0);
    auto msb_acc = luts_array->get_lut(0, 1);

    // define functions for each accumulator
    auto lut_f_lsb = [message_modulus](Torus x, Torus y) -> Torus {
      return (x * y) % message_modulus;
    };
    auto lut_f_msb = [message_modulus](Torus x, Torus y) -> Torus {
      return (x * y) / message_modulus;
    };

    // generate accumulators
    generate_device_accumulator_bivariate<Torus>(
        streams.stream(0), streams.gpu_index(0), lsb_acc,
        luts_array->get_degree(0), luts_array->get_max_degree(0),
        glwe_dimension, polynomial_size, message_modulus, carry_modulus,
        lut_f_lsb, gpu_memory_allocated);
    generate_device_accumulator_bivariate<Torus>(
        streams.stream(0), streams.gpu_index(0), msb_acc,
        luts_array->get_degree(1), luts_array->get_max_degree(1),
        glwe_dimension, polynomial_size, message_modulus, carry_modulus,
        lut_f_msb, gpu_memory_allocated);

    // lut_indexes_vec for luts_array should be reinitialized
    // first lsb_vector_block_count value should reference to lsb_acc
    // last msb_vector_block_count values should reference to msb_acc
    // for message and carry default lut_indexes_vec is fine
    if (allocate_gpu_memory)
      cuda_set_value_async<Torus>(
          streams.stream(0), streams.gpu_index(0),
          luts_array->get_lut_indexes(0, lsb_vector_block_count), 1,
          msb_vector_block_count);
    auto active_streams = streams.active_gpu_subset(total_block_count);
    luts_array->broadcast_lut(active_streams);
    // create memory object for sum ciphertexts
    sum_ciphertexts_mem = new int_sum_ciphertexts_vec_memory<Torus>(
        streams, params, num_radix_blocks, 2 * num_radix_blocks,
        vector_result_sb, small_lwe_vector, luts_array, true,
        allocate_gpu_memory, size_tracker);
    uint32_t requested_flag = outputFlag::FLAG_NONE;
    sc_prop_mem = new int_sc_prop_memory<Torus>(
        streams, params, num_radix_blocks, requested_flag, allocate_gpu_memory,
        size_tracker);
  }

  void release(CudaStreams streams) {

    if (boolean_mul) {
      zero_out_predicate_lut->release(streams);
      zero_out_mem->release(streams);
      delete zero_out_mem;
      delete zero_out_predicate_lut;

      return;
    }
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   vector_result_sb, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   block_mul_res, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   small_lwe_vector, gpu_memory_allocated);

    luts_array->release(streams);
    sum_ciphertexts_mem->release(streams);
    sc_prop_mem->release(streams);

    delete vector_result_sb;
    delete block_mul_res;
    delete small_lwe_vector;
    delete luts_array;
    delete sum_ciphertexts_mem;
    delete sc_prop_mem;
  }
};

template <typename Torus> struct int_logical_scalar_shift_buffer {
  int_radix_params params;
  std::vector<int_radix_lut<Torus> *> lut_buffers_bivariate;

  SHIFT_OR_ROTATE_TYPE shift_type;

  CudaRadixCiphertextFFI *tmp_rotated;

  bool reuse_memory = false;
  bool gpu_memory_allocated;

  int_logical_scalar_shift_buffer(CudaStreams streams,
                                  SHIFT_OR_ROTATE_TYPE shift_type,
                                  int_radix_params params,
                                  uint32_t num_radix_blocks,
                                  bool allocate_gpu_memory,
                                  uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;
    this->shift_type = shift_type;
    this->params = params;

    uint32_t max_amount_of_pbs = num_radix_blocks;

    tmp_rotated = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_rotated,
        max_amount_of_pbs + 2, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    uint32_t num_bits_in_block = (uint32_t)std::log2(params.message_modulus);

    // LUT
    // pregenerate lut vector and indexes
    // lut for left shift
    // here we generate 'num_bits_in_block' times lut
    // one for each 'shift_within_block' = 'shift' % 'num_bits_in_block'
    // even though lut_left contains 'num_bits_in_block' lut
    // lut_indexes_vec will have indexes for single lut only and those indexes
    // will be 0 it means for pbs corresponding lut should be selected and
    // pass along lut_indexes_vec filled with zeros

    // calculate bivariate lut for each 'shift_within_block'
    // so that in case an application calls scratches only once for a whole
    // circuit it can reuse memory for different shift values
    for (int s_w_b = 1; s_w_b < num_bits_in_block; s_w_b++) {
      auto cur_lut_bivariate =
          new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                   allocate_gpu_memory, size_tracker);

      uint32_t shift_within_block = s_w_b;

      std::function<Torus(Torus, Torus)> shift_lut_f;

      if (shift_type == LEFT_SHIFT) {
        shift_lut_f = [shift_within_block, params](
                          Torus current_block, Torus previous_block) -> Torus {
          current_block = current_block << shift_within_block;
          previous_block = previous_block << shift_within_block;

          Torus message_of_current_block =
              current_block % params.message_modulus;
          Torus carry_of_previous_block =
              previous_block / params.message_modulus;
          return message_of_current_block + carry_of_previous_block;
        };
      } else {
        shift_lut_f = [num_bits_in_block, shift_within_block,
                       params](Torus current_block, Torus next_block) -> Torus {
          // left shift so as not to lose
          // bits when shifting right afterwards
          next_block <<= num_bits_in_block;
          next_block >>= shift_within_block;

          // The way of getting carry / message is reversed compared
          // to the usual way but its normal:
          // The message is in the upper bits, the carry in lower bits
          Torus message_of_current_block = current_block >> shift_within_block;
          Torus carry_of_previous_block = next_block % params.message_modulus;

          return message_of_current_block + carry_of_previous_block;
        };
      }

      // right shift
      generate_device_accumulator_bivariate<Torus>(
          streams.stream(0), streams.gpu_index(0),
          cur_lut_bivariate->get_lut(0, 0), cur_lut_bivariate->get_degree(0),
          cur_lut_bivariate->get_max_degree(0), params.glwe_dimension,
          params.polynomial_size, params.message_modulus, params.carry_modulus,
          shift_lut_f, gpu_memory_allocated);
      auto active_streams = streams.active_gpu_subset(num_radix_blocks);
      cur_lut_bivariate->broadcast_lut(active_streams);

      lut_buffers_bivariate.push_back(cur_lut_bivariate);
    }
  }

  int_logical_scalar_shift_buffer(CudaStreams streams,
                                  SHIFT_OR_ROTATE_TYPE shift_type,
                                  int_radix_params params,
                                  uint32_t num_radix_blocks,
                                  bool allocate_gpu_memory,
                                  CudaRadixCiphertextFFI *pre_allocated_buffer,
                                  uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;
    this->shift_type = shift_type;
    this->params = params;
    tmp_rotated = pre_allocated_buffer;
    reuse_memory = true;

    if (allocate_gpu_memory)
      set_zero_radix_ciphertext_slice_async<Torus>(
          streams.stream(0), streams.gpu_index(0), tmp_rotated, 0,
          tmp_rotated->num_radix_blocks);

    uint32_t num_bits_in_block = (uint32_t)std::log2(params.message_modulus);

    // LUT
    // pregenerate lut vector and indexes
    // lut for left shift
    // here we generate 'num_bits_in_block' times lut
    // one for each 'shift_within_block' = 'shift' % 'num_bits_in_block'
    // even though lut_left contains 'num_bits_in_block' lut
    // lut_indexes_vec will have indexes for single lut only and those indexes
    // will be 0 it means for pbs corresponding lut should be selected and
    // pass along lut_indexes_vec filled with zeros

    // calculate bivariate lut for each 'shift_within_block'
    // so that in case an application calls scratches only once for a whole
    // circuit it can reuse memory for different shift values
    for (int s_w_b = 1; s_w_b < num_bits_in_block; s_w_b++) {
      auto cur_lut_bivariate =
          new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                   allocate_gpu_memory, size_tracker);

      uint32_t shift_within_block = s_w_b;

      std::function<Torus(Torus, Torus)> shift_lut_f;

      if (shift_type == LEFT_SHIFT) {
        shift_lut_f = [shift_within_block, params](
                          Torus current_block, Torus previous_block) -> Torus {
          current_block = current_block << shift_within_block;
          previous_block = previous_block << shift_within_block;

          Torus message_of_current_block =
              current_block % params.message_modulus;
          Torus carry_of_previous_block =
              previous_block / params.message_modulus;
          return message_of_current_block + carry_of_previous_block;
        };
      } else {
        shift_lut_f = [num_bits_in_block, shift_within_block,
                       params](Torus current_block, Torus next_block) -> Torus {
          // left shift so as not to lose
          // bits when shifting right afterwards
          next_block <<= num_bits_in_block;
          next_block >>= shift_within_block;

          // The way of getting carry / message is reversed compared
          // to the usual way but its normal:
          // The message is in the upper bits, the carry in lower bits
          Torus message_of_current_block = current_block >> shift_within_block;
          Torus carry_of_previous_block = next_block % params.message_modulus;

          return message_of_current_block + carry_of_previous_block;
        };
      }

      // right shift
      generate_device_accumulator_bivariate<Torus>(
          streams.stream(0), streams.gpu_index(0),
          cur_lut_bivariate->get_lut(0, 0), cur_lut_bivariate->get_degree(0),
          cur_lut_bivariate->get_max_degree(0), params.glwe_dimension,
          params.polynomial_size, params.message_modulus, params.carry_modulus,
          shift_lut_f, gpu_memory_allocated);
      auto active_streams = streams.active_gpu_subset(num_radix_blocks);
      cur_lut_bivariate->broadcast_lut(active_streams);

      lut_buffers_bivariate.push_back(cur_lut_bivariate);
    }
  }
  void release(CudaStreams streams) {
    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
    for (auto &buffer : lut_buffers_bivariate) {
      buffer->release(streams);
      delete buffer;
    }
    lut_buffers_bivariate.clear();

    if (!reuse_memory) {
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     tmp_rotated, gpu_memory_allocated);
      cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
      delete tmp_rotated;
    }
  }
};

template <typename Torus> struct int_arithmetic_scalar_shift_buffer {
  int_radix_params params;
  std::vector<int_radix_lut<Torus> *> lut_buffers_univariate;
  std::vector<int_radix_lut<Torus> *> lut_buffers_bivariate;

  SHIFT_OR_ROTATE_TYPE shift_type;

  CudaRadixCiphertextFFI *tmp_rotated;

  CudaStreams local_streams_1;
  CudaStreams local_streams_2;
  bool gpu_memory_allocated;

  int_arithmetic_scalar_shift_buffer(CudaStreams streams,
                                     SHIFT_OR_ROTATE_TYPE shift_type,
                                     int_radix_params params,
                                     uint32_t num_radix_blocks,
                                     bool allocate_gpu_memory,
                                     uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;

    auto active_streams = streams.active_gpu_subset(1);
    // In the arithmetic shift, a PBS has to be applied to the last rotated
    // block twice: once to shift it, once to compute the padding block to be
    // copied onto all blocks to the left of the last rotated block
    local_streams_1.create_on_same_gpus(active_streams);
    local_streams_2.create_on_same_gpus(active_streams);
    this->shift_type = shift_type;
    this->params = params;

    tmp_rotated = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_rotated,
        num_radix_blocks + 3, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    uint32_t num_bits_in_block = (uint32_t)std::log2(params.message_modulus);

    // LUT
    // pregenerate lut vector and indexes lut

    // lut to shift the last block
    // calculate lut for each 'shift_within_block'
    // so that in case an application calls scratches only once for a whole
    // circuit it can reuse memory for different shift values
    // With two bits of message this is actually only one LUT.
    for (int s_w_b = 1; s_w_b < num_bits_in_block; s_w_b++) {
      auto shift_last_block_lut_univariate = new int_radix_lut<Torus>(
          streams, params, 1, 1, allocate_gpu_memory, size_tracker);

      uint32_t shift_within_block = s_w_b;

      std::function<Torus(Torus)> last_block_lut_f;
      last_block_lut_f = [num_bits_in_block, shift_within_block,
                          params](Torus x) -> Torus {
        x = x % params.message_modulus;
        uint32_t x_sign_bit = x >> (num_bits_in_block - 1) & 1;
        uint32_t shifted = x >> shift_within_block;
        // padding is a message full of 1 if sign bit is one
        // else padding is a zero message
        uint32_t padding = (params.message_modulus - 1) * x_sign_bit;

        // Make padding have 1s only in places where bits
        // where actually need to be padded
        padding <<= num_bits_in_block - shift_within_block;
        padding %= params.message_modulus;

        return shifted | padding;
      };

      generate_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0),
          shift_last_block_lut_univariate->get_lut(0, 0),
          shift_last_block_lut_univariate->get_degree(0),
          shift_last_block_lut_univariate->get_max_degree(0),
          params.glwe_dimension, params.polynomial_size, params.message_modulus,
          params.carry_modulus, last_block_lut_f, gpu_memory_allocated);
      auto active_streams_shift_last = streams.active_gpu_subset(1);
      shift_last_block_lut_univariate->broadcast_lut(active_streams_shift_last);

      lut_buffers_univariate.push_back(shift_last_block_lut_univariate);
    }

    auto padding_block_lut_univariate = new int_radix_lut<Torus>(
        streams, params, 1, 1, allocate_gpu_memory, size_tracker);

    // lut to compute the padding block
    std::function<Torus(Torus)> padding_block_lut_f;
    padding_block_lut_f = [num_bits_in_block, params](Torus x) -> Torus {
      x = x % params.message_modulus;
      uint32_t x_sign_bit = x >> (num_bits_in_block - 1) & 1;
      // padding is a message full of 1 if sign bit is one
      // else padding is a zero message
      return (params.message_modulus - 1) * x_sign_bit;
    };

    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0),
        padding_block_lut_univariate->get_lut(0, 0),
        padding_block_lut_univariate->get_degree(0),
        padding_block_lut_univariate->get_max_degree(0), params.glwe_dimension,
        params.polynomial_size, params.message_modulus, params.carry_modulus,
        padding_block_lut_f, gpu_memory_allocated);
    // auto active_streams = streams.active_gpu_subset(1);
    padding_block_lut_univariate->broadcast_lut(active_streams);

    lut_buffers_univariate.push_back(padding_block_lut_univariate);

    // lut to shift the first blocks
    // calculate lut for each 'shift_within_block'
    // so that in case an application calls scratches only once for a whole
    // circuit it can reuse memory for different shift values
    // NB: with two bits of message, this is actually only one LUT.
    for (int s_w_b = 1; s_w_b < num_bits_in_block; s_w_b++) {
      auto shift_blocks_lut_bivariate =
          new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                   allocate_gpu_memory, size_tracker);

      uint32_t shift_within_block = s_w_b;

      std::function<Torus(Torus, Torus)> blocks_lut_f;
      blocks_lut_f = [num_bits_in_block, shift_within_block,
                      params](Torus current_block, Torus next_block) -> Torus {
        // left shift so as not to lose
        // bits when shifting right after
        next_block <<= num_bits_in_block;
        next_block >>= shift_within_block;

        // The way of getting carry / message is reversed compared
        // to the usual way but its normal:
        // The message is in the upper bits, the carry in lower bits
        uint32_t message_of_current_block = current_block >> shift_within_block;
        uint32_t carry_of_previous_block = next_block % params.message_modulus;

        return message_of_current_block + carry_of_previous_block;
      };

      generate_device_accumulator_bivariate<Torus>(
          streams.stream(0), streams.gpu_index(0),
          shift_blocks_lut_bivariate->get_lut(0, 0),
          shift_blocks_lut_bivariate->get_degree(0),
          shift_blocks_lut_bivariate->get_max_degree(0), params.glwe_dimension,
          params.polynomial_size, params.message_modulus, params.carry_modulus,
          blocks_lut_f, gpu_memory_allocated);
      auto active_streams_shift_blocks =
          streams.active_gpu_subset(num_radix_blocks);
      shift_blocks_lut_bivariate->broadcast_lut(active_streams_shift_blocks);

      lut_buffers_bivariate.push_back(shift_blocks_lut_bivariate);
    }
  }

  void release(CudaStreams streams) {
    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
    local_streams_1.release();
    local_streams_2.release();
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_rotated, gpu_memory_allocated);
    for (auto &buffer : lut_buffers_bivariate) {
      buffer->release(streams);
      delete buffer;
    }
    for (auto &buffer : lut_buffers_univariate) {
      buffer->release(streams);
      delete buffer;
    }
    lut_buffers_bivariate.clear();
    lut_buffers_univariate.clear();

    delete tmp_rotated;
  }
};

template <typename Torus> struct int_cmux_buffer {
  int_radix_lut<Torus> *predicate_lut;
  int_radix_lut<Torus> *message_extract_lut;

  CudaRadixCiphertextFFI *buffer_in;
  CudaRadixCiphertextFFI *buffer_out;
  CudaRadixCiphertextFFI *condition_array;

  int_radix_params params;
  bool allocate_gpu_memory;
  bool gpu_memory_allocated;
  int_cmux_buffer(CudaStreams streams,
                  std::function<Torus(Torus)> predicate_lut_f,
                  int_radix_params params, uint32_t num_radix_blocks,
                  bool allocate_gpu_memory, uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;

    this->params = params;
    this->allocate_gpu_memory = allocate_gpu_memory;

    buffer_in = new CudaRadixCiphertextFFI;
    buffer_out = new CudaRadixCiphertextFFI;
    condition_array = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), buffer_in,
        2 * num_radix_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), buffer_out,
        2 * num_radix_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), condition_array,
        2 * num_radix_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    auto lut_f = [predicate_lut_f](Torus block, Torus condition) -> Torus {
      return predicate_lut_f(condition) ? 0 : block;
    };
    auto inverted_lut_f = [predicate_lut_f](Torus block,
                                            Torus condition) -> Torus {
      return predicate_lut_f(condition) ? block : 0;
    };
    auto message_extract_lut_f = [params](Torus x) -> Torus {
      return x % params.message_modulus;
    };

    predicate_lut =
        new int_radix_lut<Torus>(streams, params, 2, 2 * num_radix_blocks,
                                 allocate_gpu_memory, size_tracker);

    message_extract_lut =
        new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                 allocate_gpu_memory, size_tracker);

    generate_device_accumulator_bivariate<Torus>(
        streams.stream(0), streams.gpu_index(0), predicate_lut->get_lut(0, 0),
        predicate_lut->get_degree(0), predicate_lut->get_max_degree(0),
        params.glwe_dimension, params.polynomial_size, params.message_modulus,
        params.carry_modulus, inverted_lut_f, gpu_memory_allocated);

    generate_device_accumulator_bivariate<Torus>(
        streams.stream(0), streams.gpu_index(0), predicate_lut->get_lut(0, 1),
        predicate_lut->get_degree(1), predicate_lut->get_max_degree(1),
        params.glwe_dimension, params.polynomial_size, params.message_modulus,
        params.carry_modulus, lut_f, gpu_memory_allocated);

    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0),
        message_extract_lut->get_lut(0, 0), message_extract_lut->get_degree(0),
        message_extract_lut->get_max_degree(0), params.glwe_dimension,
        params.polynomial_size, params.message_modulus, params.carry_modulus,
        message_extract_lut_f, gpu_memory_allocated);
    Torus *h_lut_indexes = predicate_lut->h_lut_indexes;
    for (int index = 0; index < 2 * num_radix_blocks; index++) {
      if (index < num_radix_blocks) {
        h_lut_indexes[index] = 0;
      } else {
        h_lut_indexes[index] = 1;
      }
    }
    cuda_memcpy_with_size_tracking_async_to_gpu(
        predicate_lut->get_lut_indexes(0, 0), h_lut_indexes,
        2 * num_radix_blocks * sizeof(Torus), streams.stream(0),
        streams.gpu_index(0), allocate_gpu_memory);
    auto active_streams_pred = streams.active_gpu_subset(2 * num_radix_blocks);
    predicate_lut->broadcast_lut(active_streams_pred);
    auto active_streams_msg = streams.active_gpu_subset(num_radix_blocks);
    message_extract_lut->broadcast_lut(active_streams_msg);
  }

  void release(CudaStreams streams) {
    predicate_lut->release(streams);
    delete predicate_lut;
    message_extract_lut->release(streams);
    delete message_extract_lut;

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   buffer_in, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   buffer_out, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   condition_array, gpu_memory_allocated);
    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
    delete buffer_in;
    delete buffer_out;
    delete condition_array;
  }
};

template <typename Torus> struct int_are_all_block_true_buffer {
  COMPARISON_TYPE op;
  int_radix_params params;

  CudaRadixCiphertextFFI *tmp_out;
  CudaRadixCiphertextFFI *tmp_block_accumulated;

  // This map store LUTs that checks the equality between some input and values
  // of interest in are_all_block_true(), as with max_value (the maximum message
  // value).
  int_radix_lut<Torus> *is_max_value;
  Torus *preallocated_h_lut;
  bool gpu_memory_allocated;

  int_are_all_block_true_buffer(CudaStreams streams, COMPARISON_TYPE op,
                                int_radix_params params,
                                uint32_t num_radix_blocks,
                                bool allocate_gpu_memory,
                                uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;
    this->params = params;
    this->op = op;

    Torus total_modulus = params.message_modulus * params.carry_modulus;
    uint32_t max_value = (total_modulus - 1) / (params.message_modulus - 1);

    int max_chunks = (num_radix_blocks + max_value - 1) / max_value;
    tmp_out = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_out, num_radix_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    tmp_block_accumulated = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_block_accumulated,
        max_chunks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    is_max_value = new int_radix_lut<Torus>(streams, params, 2, max_chunks,
                                            allocate_gpu_memory, size_tracker);
    auto is_max_value_f = [max_value](Torus x) -> Torus {
      return x == max_value;
    };
    preallocated_h_lut = (Torus *)malloc(
        (params.glwe_dimension + 1) * params.polynomial_size * sizeof(Torus));
    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0), is_max_value->get_lut(0, 0),
        is_max_value->get_degree(0), is_max_value->get_max_degree(0),
        params.glwe_dimension, params.polynomial_size, params.message_modulus,
        params.carry_modulus, is_max_value_f, gpu_memory_allocated);

    auto active_streams = streams.active_gpu_subset(max_chunks);
    is_max_value->broadcast_lut(active_streams);
  }

  void release(CudaStreams streams) {
    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_out, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_block_accumulated, gpu_memory_allocated);
    is_max_value->release(streams);
    free(preallocated_h_lut);
    delete is_max_value;
    delete tmp_out;
    delete tmp_block_accumulated;
  }
};

template <typename Torus> struct int_comparison_eq_buffer {
  int_radix_params params;
  COMPARISON_TYPE op;

  int_radix_lut<Torus> *operator_lut;
  int_radix_lut<Torus> *is_non_zero_lut;
  int_radix_lut<Torus> *scalar_comparison_luts;

  int_are_all_block_true_buffer<Torus> *are_all_block_true_buffer;
  bool gpu_memory_allocated;

  int_comparison_eq_buffer(CudaStreams streams, COMPARISON_TYPE op,
                           int_radix_params params, uint32_t num_radix_blocks,
                           bool allocate_gpu_memory, uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;
    this->params = params;
    this->op = op;
    Torus total_modulus = params.message_modulus * params.carry_modulus;

    are_all_block_true_buffer = new int_are_all_block_true_buffer<Torus>(
        streams, op, params, num_radix_blocks, allocate_gpu_memory,
        size_tracker);

    // f(x) -> x == 0
    auto is_non_zero_lut_f = [total_modulus](Torus x) -> Torus {
      return (x % total_modulus) != 0;
    };

    is_non_zero_lut =
        new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                 allocate_gpu_memory, size_tracker);

    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0), is_non_zero_lut->get_lut(0, 0),
        is_non_zero_lut->get_degree(0), is_non_zero_lut->get_max_degree(0),
        params.glwe_dimension, params.polynomial_size, params.message_modulus,
        params.carry_modulus, is_non_zero_lut_f, gpu_memory_allocated);

    auto active_streams = streams.active_gpu_subset(num_radix_blocks);
    is_non_zero_lut->broadcast_lut(active_streams);

    // Scalar may have up to num_radix_blocks blocks
    scalar_comparison_luts = new int_radix_lut<Torus>(
        streams, params, total_modulus, num_radix_blocks, allocate_gpu_memory,
        size_tracker);

    // Operator LUT
    auto operator_f = [op](Torus lhs, Torus rhs) -> Torus {
      if (op == COMPARISON_TYPE::EQ) {
        return (lhs == rhs);
      } else if (op == COMPARISON_TYPE::NE) {
        return (lhs != rhs);
      } else {
        // For signed scalar comparisons we check equality with zero
        return (lhs == rhs);
      }
    };
    for (int i = 0; i < total_modulus; i++) {
      auto lut_f = [i, operator_f](Torus x) -> Torus {
        return operator_f(i, x);
      };

      generate_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0),
          scalar_comparison_luts->get_lut(0, i),
          scalar_comparison_luts->get_degree(i),
          scalar_comparison_luts->get_max_degree(i), params.glwe_dimension,
          params.polynomial_size, params.message_modulus, params.carry_modulus,
          lut_f, gpu_memory_allocated);
    }
    scalar_comparison_luts->broadcast_lut(active_streams);
    if (op == COMPARISON_TYPE::EQ || op == COMPARISON_TYPE::NE) {
      operator_lut =
          new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                   allocate_gpu_memory, size_tracker);

      generate_device_accumulator_bivariate<Torus>(
          streams.stream(0), streams.gpu_index(0), operator_lut->get_lut(0, 0),
          operator_lut->get_degree(0), operator_lut->get_max_degree(0),
          params.glwe_dimension, params.polynomial_size, params.message_modulus,
          params.carry_modulus, operator_f, gpu_memory_allocated);

      operator_lut->broadcast_lut(active_streams);
    } else {
      operator_lut = nullptr;
    }
  }

  void release(CudaStreams streams) {
    if (op == COMPARISON_TYPE::EQ || op == COMPARISON_TYPE::NE) {
      PANIC_IF_FALSE(operator_lut != nullptr,
                     "Cuda error: no operator lut was created");
      operator_lut->release(streams);
      delete operator_lut;
      operator_lut = nullptr;
    }
    is_non_zero_lut->release(streams);
    delete is_non_zero_lut;
    is_non_zero_lut = nullptr;
    are_all_block_true_buffer->release(streams);
    delete are_all_block_true_buffer;
    are_all_block_true_buffer = nullptr;
    scalar_comparison_luts->release(streams);
    delete scalar_comparison_luts;
    scalar_comparison_luts = nullptr;
  }
};

template <typename Torus> struct int_tree_sign_reduction_buffer {
  int_radix_params params;

  std::function<Torus(Torus, Torus)> block_selector_f;

  int_radix_lut<Torus> *tree_inner_leaf_lut;
  int_radix_lut<Torus> *tree_last_leaf_lut;

  int_radix_lut<Torus> *tree_last_leaf_scalar_lut;

  Torus *preallocated_h_lut;
  CudaRadixCiphertextFFI *tmp_x;
  CudaRadixCiphertextFFI *tmp_y;
  bool gpu_memory_allocated;

  int_tree_sign_reduction_buffer(CudaStreams streams,
                                 std::function<Torus(Torus)> operator_f,
                                 int_radix_params params,
                                 uint32_t num_radix_blocks,
                                 bool allocate_gpu_memory,
                                 uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;
    this->params = params;

    block_selector_f = [](Torus msb, Torus lsb) -> Torus {
      if (msb == IS_EQUAL) // EQUAL
        return lsb;
      else
        return msb;
    };

    tmp_x = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_x, num_radix_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    tmp_y = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_y, num_radix_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    // LUTs
    tree_inner_leaf_lut =
        new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                 allocate_gpu_memory, size_tracker);

    tree_last_leaf_lut = new int_radix_lut<Torus>(
        streams, params, 1, 1, allocate_gpu_memory, size_tracker);

    preallocated_h_lut = (Torus *)malloc(
        (params.glwe_dimension + 1) * params.polynomial_size * sizeof(Torus));

    tree_last_leaf_scalar_lut = new int_radix_lut<Torus>(
        streams, params, 1, 1, allocate_gpu_memory, size_tracker);

    generate_device_accumulator_bivariate<Torus>(
        streams.stream(0), streams.gpu_index(0),
        tree_inner_leaf_lut->get_lut(0, 0), tree_inner_leaf_lut->get_degree(0),
        tree_inner_leaf_lut->get_max_degree(0), params.glwe_dimension,
        params.polynomial_size, params.message_modulus, params.carry_modulus,
        block_selector_f, gpu_memory_allocated);
    auto active_streams = streams.active_gpu_subset(num_radix_blocks);
    tree_inner_leaf_lut->broadcast_lut(active_streams);
  }

  void release(CudaStreams streams) {
    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_x, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_y, gpu_memory_allocated);
    tree_inner_leaf_lut->release(streams);
    delete tree_inner_leaf_lut;
    tree_last_leaf_lut->release(streams);
    delete tree_last_leaf_lut;
    tree_last_leaf_scalar_lut->release(streams);
    delete tree_last_leaf_scalar_lut;

    delete tmp_x;
    delete tmp_y;
    free(preallocated_h_lut);
  }
};

template <typename Torus> struct int_comparison_diff_buffer {
  int_radix_params params;
  COMPARISON_TYPE op;

  CudaRadixCiphertextFFI *tmp_packed;

  std::function<Torus(Torus)> operator_f;

  int_tree_sign_reduction_buffer<Torus> *tree_buffer;

  CudaRadixCiphertextFFI *tmp_signs_a;
  CudaRadixCiphertextFFI *tmp_signs_b;
  int_radix_lut<Torus> *reduce_signs_lut;
  bool gpu_memory_allocated;
  Torus *preallocated_h_lut1;
  Torus *preallocated_h_lut2;
  int_comparison_diff_buffer(CudaStreams streams, COMPARISON_TYPE op,
                             int_radix_params params, uint32_t num_radix_blocks,
                             bool allocate_gpu_memory, uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;
    this->params = params;
    this->op = op;

    operator_f = [op](Torus x) -> Torus {
      switch (op) {
      case GT:
        return x == IS_SUPERIOR;
      case GE:
        return (x == IS_SUPERIOR) || (x == IS_EQUAL);
      case LT:
        return x == IS_INFERIOR;
      case LE:
        return (x == IS_INFERIOR) || (x == IS_EQUAL);
      default:
        PANIC("Cuda error (comparisons): unknown comparison type")
      }
    };

    tmp_packed = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_packed, num_radix_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

    tree_buffer = new int_tree_sign_reduction_buffer<Torus>(
        streams, operator_f, params, num_radix_blocks, allocate_gpu_memory,
        size_tracker);
    tmp_signs_a = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_signs_a, num_radix_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    tmp_signs_b = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_signs_b, num_radix_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    // LUTs
    reduce_signs_lut =
        new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                 allocate_gpu_memory, size_tracker);
    preallocated_h_lut1 = (Torus *)malloc(
        (params.glwe_dimension + 1) * params.polynomial_size * sizeof(Torus));
    preallocated_h_lut2 = (Torus *)malloc(
        (params.glwe_dimension + 1) * params.polynomial_size * sizeof(Torus));
  }

  void release(CudaStreams streams) {
    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_packed, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_signs_a, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_signs_b, gpu_memory_allocated);
    tree_buffer->release(streams);
    delete tree_buffer;
    reduce_signs_lut->release(streams);
    delete reduce_signs_lut;

    delete tmp_packed;
    delete tmp_signs_a;
    delete tmp_signs_b;
    free(preallocated_h_lut1);
    free(preallocated_h_lut2);
  }
};

template <typename Torus> struct int_comparison_buffer {
  COMPARISON_TYPE op;

  int_radix_params params;

  //////////////////
  int_radix_lut<Torus> *identity_lut;
  std::function<Torus(Torus)> identity_lut_f;

  int_radix_lut<Torus> *is_zero_lut;

  int_comparison_eq_buffer<Torus> *eq_buffer;
  int_comparison_diff_buffer<Torus> *diff_buffer;

  CudaRadixCiphertextFFI *tmp_block_comparisons;
  CudaRadixCiphertextFFI *tmp_lwe_array_out;
  CudaRadixCiphertextFFI *tmp_trivial_sign_block;

  // Scalar EQ / NE
  CudaRadixCiphertextFFI *tmp_packed_input;

  // Max Min
  int_cmux_buffer<Torus> *cmux_buffer;

  // Signed LUT
  int_radix_lut<Torus> *signed_lut;
  bool is_signed;

  // Used for scalar comparisons
  int_radix_lut<Torus> *signed_msb_lut;
  CudaStreams lsb_streams;
  CudaStreams msb_streams;
  bool gpu_memory_allocated;
  Torus *preallocated_h_lut;

  int_comparison_buffer(CudaStreams streams, COMPARISON_TYPE op,
                        int_radix_params params, uint32_t num_radix_blocks,
                        bool is_signed, bool allocate_gpu_memory,
                        uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;
    this->params = params;
    this->op = op;
    this->is_signed = is_signed;

    auto active_streams = streams.active_gpu_subset(num_radix_blocks);

    identity_lut_f = [](Torus x) -> Torus { return x; };

    lsb_streams.create_on_same_gpus(active_streams);
    msb_streams.create_on_same_gpus(active_streams);

    // +1 to have space for signed comparison
    tmp_lwe_array_out = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_lwe_array_out,
        num_radix_blocks + 1, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    tmp_packed_input = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_packed_input,
        2 * num_radix_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    // Block comparisons
    tmp_block_comparisons = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_block_comparisons,
        num_radix_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    // Cleaning LUT
    identity_lut =
        new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                 allocate_gpu_memory, size_tracker);

    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0), identity_lut->get_lut(0, 0),
        identity_lut->get_degree(0), identity_lut->get_max_degree(0),
        params.glwe_dimension, params.polynomial_size, params.message_modulus,
        params.carry_modulus, identity_lut_f, gpu_memory_allocated);
    identity_lut->broadcast_lut(active_streams);

    uint32_t total_modulus = params.message_modulus * params.carry_modulus;
    auto is_zero_f = [total_modulus](Torus x) -> Torus {
      return (x % total_modulus) == 0;
    };

    is_zero_lut = new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                           allocate_gpu_memory, size_tracker);

    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0), is_zero_lut->get_lut(0, 0),
        is_zero_lut->get_degree(0), is_zero_lut->get_max_degree(0),
        params.glwe_dimension, params.polynomial_size, params.message_modulus,
        params.carry_modulus, is_zero_f, gpu_memory_allocated);

    is_zero_lut->broadcast_lut(active_streams);

    switch (op) {
    case COMPARISON_TYPE::MAX:
    case COMPARISON_TYPE::MIN:
      cmux_buffer = new int_cmux_buffer<Torus>(
          streams,
          [op](Torus x) -> Torus {
            if (op == COMPARISON_TYPE::MAX)
              return (x == IS_SUPERIOR);
            else
              return (x == IS_INFERIOR);
          },
          params, num_radix_blocks, allocate_gpu_memory, size_tracker);
    case COMPARISON_TYPE::GT:
    case COMPARISON_TYPE::GE:
    case COMPARISON_TYPE::LT:
    case COMPARISON_TYPE::LE:
      diff_buffer = new int_comparison_diff_buffer<Torus>(
          streams, op, params, num_radix_blocks, allocate_gpu_memory,
          size_tracker);
    case COMPARISON_TYPE::EQ:
    case COMPARISON_TYPE::NE:
      eq_buffer = new int_comparison_eq_buffer<Torus>(
          streams, op, params, num_radix_blocks, allocate_gpu_memory,
          size_tracker);
      break;
    default:
      PANIC("Unsupported comparison operation.")
    }

    if (is_signed) {

      tmp_trivial_sign_block = new CudaRadixCiphertextFFI;
      create_zero_radix_ciphertext_async<Torus>(
          streams.stream(0), streams.gpu_index(0), tmp_trivial_sign_block, 1,
          params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

      signed_lut = new int_radix_lut<Torus>(streams, params, 1, 1,
                                            allocate_gpu_memory, size_tracker);
      signed_msb_lut = new int_radix_lut<Torus>(
          streams, params, 1, 1, allocate_gpu_memory, size_tracker);

      auto message_modulus = (int)params.message_modulus;
      uint32_t sign_bit_pos = log2(message_modulus) - 1;
      std::function<Torus(Torus, Torus)> signed_lut_f =
          [sign_bit_pos](Torus x, Torus y) -> Torus {
        auto x_sign_bit = x >> sign_bit_pos;
        auto y_sign_bit = y >> sign_bit_pos;

        // The block that has its sign bit set is going
        // to be ordered as 'greater' by the cmp fn.
        // However, we are dealing with signed number,
        // so in reality, it is the smaller of the two.
        // i.e the cmp result is reversed
        if (x_sign_bit == y_sign_bit) {
          // Both have either sign bit set or unset,
          // cmp will give correct result
          if (x < y)
            return (Torus)(IS_INFERIOR);
          else if (x == y)
            return (Torus)(IS_EQUAL);
          else
            return (Torus)(IS_SUPERIOR);
        } else {
          if (x < y)
            return (Torus)(IS_SUPERIOR);
          else if (x == y)
            return (Torus)(IS_EQUAL);
          else
            return (Torus)(IS_INFERIOR);
        }
        PANIC("Cuda error: sign_lut creation failed due to wrong function.")
      };

      generate_device_accumulator_bivariate<Torus>(
          streams.stream(0), streams.gpu_index(0), signed_lut->get_lut(0, 0),
          signed_lut->get_degree(0), signed_lut->get_max_degree(0),
          params.glwe_dimension, params.polynomial_size, params.message_modulus,
          params.carry_modulus, signed_lut_f, gpu_memory_allocated);
      auto active_streams = streams.active_gpu_subset(1);
      signed_lut->broadcast_lut(active_streams);
    }
    preallocated_h_lut = (Torus *)malloc(
        (params.glwe_dimension + 1) * params.polynomial_size * sizeof(Torus));
  }

  void release(CudaStreams streams) {
    switch (op) {
    case COMPARISON_TYPE::MAX:
    case COMPARISON_TYPE::MIN:
      cmux_buffer->release(streams);
      delete (cmux_buffer);
    case COMPARISON_TYPE::GT:
    case COMPARISON_TYPE::GE:
    case COMPARISON_TYPE::LT:
    case COMPARISON_TYPE::LE:
      diff_buffer->release(streams);
      delete (diff_buffer);
    case COMPARISON_TYPE::EQ:
    case COMPARISON_TYPE::NE:
      eq_buffer->release(streams);
      delete (eq_buffer);
      break;
    default:
      PANIC("Unsupported comparison operation.")
    }
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_lwe_array_out, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_block_comparisons, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_packed_input, gpu_memory_allocated);
    identity_lut->release(streams);
    delete identity_lut;
    is_zero_lut->release(streams);
    delete is_zero_lut;
    delete tmp_lwe_array_out;
    delete tmp_block_comparisons;
    delete tmp_packed_input;

    if (is_signed) {
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     tmp_trivial_sign_block,
                                     gpu_memory_allocated);
      signed_lut->release(streams);
      delete signed_lut;
      signed_msb_lut->release(streams);
      delete signed_msb_lut;
      delete tmp_trivial_sign_block;
    }
    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
    lsb_streams.release();
    msb_streams.release();
    free(preallocated_h_lut);
  }
};

// used only when 4 gpus are available
template <typename Torus> struct unsigned_int_div_rem_2_2_memory {
  int_radix_params params;
  bool gpu_memory_allocated;

  // memory objects for other operations
  int_borrow_prop_memory<Torus> *overflow_sub_mem_1;
  int_borrow_prop_memory<Torus> *overflow_sub_mem_2;
  int_borrow_prop_memory<Torus> *overflow_sub_mem_3;
  int_comparison_buffer<Torus> *comparison_buffer_1;
  int_comparison_buffer<Torus> *comparison_buffer_2;
  int_comparison_buffer<Torus> *comparison_buffer_3;
  int_sub_and_propagate<Torus> *sub_and_propagate_mem;
  int_bitop_buffer<Torus> *bitor_mem_1;
  int_bitop_buffer<Torus> *bitor_mem_2;
  int_bitop_buffer<Torus> *bitor_mem_3;
  int_logical_scalar_shift_buffer<Torus> *shift_mem;

  // lookup tables
  int_radix_lut<Torus> *message_extract_lut_1;
  int_radix_lut<Torus> *message_extract_lut_2;
  int_radix_lut<Torus> *zero_out_if_not_1_lut_1;
  int_radix_lut<Torus> *zero_out_if_not_1_lut_2;
  int_radix_lut<Torus> *zero_out_if_not_2_lut_1;
  int_radix_lut<Torus> *zero_out_if_not_2_lut_2;
  int_radix_lut<Torus> *quotient_lut_1;
  int_radix_lut<Torus> *quotient_lut_2;
  int_radix_lut<Torus> *quotient_lut_3;

  // sub streams
  CudaStreams sub_streams_1;

  // temporary device buffers
  CudaRadixCiphertextFFI *d1;                  // num_blocks + 1
  CudaRadixCiphertextFFI *d2;                  // num_blocks + 1
  CudaRadixCiphertextFFI *d3;                  // num_blocks + 1
  CudaRadixCiphertextFFI *tmp_gpu_0;           // num_blocks + 1
  CudaRadixCiphertextFFI *tmp_gpu_1;           // num_blocks + 1
  CudaRadixCiphertextFFI *tmp_gpu_2;           // num_blocks + 1
  CudaRadixCiphertextFFI *tmp_gpu_3;           // num_blocks + 1
  CudaRadixCiphertextFFI *divisor_gpu_1;       // num_blocks
  CudaRadixCiphertextFFI *divisor_gpu_2;       // num_blocks
  CudaRadixCiphertextFFI *remainder_gpu_1;     // num_blocks
  CudaRadixCiphertextFFI *remainder_gpu_2;     // num_blocks
  CudaRadixCiphertextFFI *remainder_gpu_3;     // num_blocks
  CudaRadixCiphertextFFI *low1;                // num_blocks
  CudaRadixCiphertextFFI *low2;                // num_blocks
  CudaRadixCiphertextFFI *low3;                // num_blocks
  CudaRadixCiphertextFFI *rem0;                // num_blocks
  CudaRadixCiphertextFFI *rem1;                // num_blocks
  CudaRadixCiphertextFFI *rem2;                // num_blocks
  CudaRadixCiphertextFFI *rem3;                // num_blocks
  CudaRadixCiphertextFFI *sub_result_1;        // num_blocks
  CudaRadixCiphertextFFI *sub_result_2;        // num_blocks
  CudaRadixCiphertextFFI *sub_result_3;        // num_blocks
  CudaRadixCiphertextFFI *sub_1_overflowed;    // num_blocks
  CudaRadixCiphertextFFI *sub_2_overflowed;    // num_blocks
  CudaRadixCiphertextFFI *sub_3_overflowed;    // num_blocks
  CudaRadixCiphertextFFI *comparison_blocks_1; // num_blocks
  CudaRadixCiphertextFFI *comparison_blocks_2; // num_blocks
  CudaRadixCiphertextFFI *comparison_blocks_3; // num_blocks
  CudaRadixCiphertextFFI *cmp_1;               // boolean block
  CudaRadixCiphertextFFI *cmp_2;               // boolean block
  CudaRadixCiphertextFFI *cmp_3;               // boolean block
  CudaRadixCiphertextFFI *c0;                  // single block
  CudaRadixCiphertextFFI *q1;                  // single block
  CudaRadixCiphertextFFI *q2;                  // single block
  CudaRadixCiphertextFFI *q3;                  // single block

  Torus *h_buffer; // used for memory copies

  Torus **first_indexes_for_overflow_sub_gpu_0;
  Torus **second_indexes_for_overflow_sub_gpu_0;
  Torus **scalars_for_overflow_sub_gpu_0;

  Torus **first_indexes_for_overflow_sub_gpu_1;
  Torus **second_indexes_for_overflow_sub_gpu_1;
  Torus **scalars_for_overflow_sub_gpu_1;

  Torus **first_indexes_for_overflow_sub_gpu_2;
  Torus **second_indexes_for_overflow_sub_gpu_2;
  Torus **scalars_for_overflow_sub_gpu_2;

  cudaEvent_t create_indexes_done;

  uint32_t max_indexes_to_erase;
  uint64_t tmp_size_tracker = 0;

  // allocate and initialize if needed, temporary arrays used to calculate
  // cuda integer div_rem_2_2 operation
  void init_temporary_buffers(CudaStreams streams, uint32_t num_blocks,
                              bool allocate_gpu_memory,
                              uint64_t &size_tracker) {

    // more than one block temporary arrays
    tmp_gpu_0 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_gpu_0, num_blocks + 1,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    d3 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), d3, num_blocks + 1,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    low3 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), low3, num_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    rem3 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), rem3, num_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    sub_result_1 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), sub_result_1, num_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    sub_1_overflowed = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), sub_1_overflowed, 1,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    comparison_blocks_1 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), comparison_blocks_1,
        num_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    cmp_1 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), cmp_1, 1,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    q3 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), q3, 1,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

    tmp_gpu_1 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(1), streams.gpu_index(1), tmp_gpu_1, num_blocks + 1,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    d2 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(1), streams.gpu_index(1), d2, num_blocks + 1,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    low2 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(1), streams.gpu_index(1), low2, num_blocks,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    rem2 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(1), streams.gpu_index(1), rem2, num_blocks,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    divisor_gpu_1 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(1), streams.gpu_index(1), divisor_gpu_1, num_blocks,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    remainder_gpu_1 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(1), streams.gpu_index(1), remainder_gpu_1, num_blocks,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    sub_result_2 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(1), streams.gpu_index(1), sub_result_2, num_blocks,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    sub_2_overflowed = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(1), streams.gpu_index(1), sub_2_overflowed, 1,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    comparison_blocks_2 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(1), streams.gpu_index(1), comparison_blocks_2,
        num_blocks, params.big_lwe_dimension, tmp_size_tracker,
        allocate_gpu_memory);
    cmp_2 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(1), streams.gpu_index(1), cmp_2, 1,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    q2 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(1), streams.gpu_index(1), q2, 1,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);

    tmp_gpu_2 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(2), streams.gpu_index(2), tmp_gpu_2, num_blocks + 1,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    d1 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(2), streams.gpu_index(2), d1, num_blocks + 1,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    low1 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(2), streams.gpu_index(2), low1, num_blocks,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    rem1 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(2), streams.gpu_index(2), rem1, num_blocks,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    divisor_gpu_2 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(2), streams.gpu_index(2), divisor_gpu_2, num_blocks,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    remainder_gpu_2 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(2), streams.gpu_index(2), remainder_gpu_2, num_blocks,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    sub_result_3 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(2), streams.gpu_index(2), sub_result_3, num_blocks,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    sub_3_overflowed = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(2), streams.gpu_index(2), sub_3_overflowed, 1,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    comparison_blocks_3 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(2), streams.gpu_index(2), comparison_blocks_3,
        num_blocks, params.big_lwe_dimension, tmp_size_tracker,
        allocate_gpu_memory);
    cmp_3 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(2), streams.gpu_index(2), cmp_3, 1,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    q1 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(2), streams.gpu_index(2), q1, 1,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);

    tmp_gpu_3 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(3), streams.gpu_index(3), tmp_gpu_3, num_blocks + 1,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    rem0 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(3), streams.gpu_index(3), rem0, num_blocks,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    remainder_gpu_3 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(3), streams.gpu_index(3), remainder_gpu_3, num_blocks,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
    c0 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(3), streams.gpu_index(3), c0, 1,
        params.big_lwe_dimension, tmp_size_tracker, allocate_gpu_memory);
  }

  // initialize lookup tables for div_rem_2_2 operation
  void init_lookup_tables(CudaStreams streams, uint32_t num_blocks,
                          bool allocate_gpu_memory, uint64_t &size_tracker) {

    zero_out_if_not_1_lut_1 =
        new int_radix_lut<Torus>(streams.get_ith(0), params, 1, num_blocks,
                                 allocate_gpu_memory, size_tracker);

    zero_out_if_not_2_lut_1 =
        new int_radix_lut<Torus>(streams.get_ith(1), params, 1, num_blocks,
                                 allocate_gpu_memory, tmp_size_tracker);

    zero_out_if_not_2_lut_2 =
        new int_radix_lut<Torus>(streams.get_ith(2), params, 1, num_blocks,
                                 allocate_gpu_memory, tmp_size_tracker);

    zero_out_if_not_1_lut_2 =
        new int_radix_lut<Torus>(streams.get_ith(3), params, 1, num_blocks,
                                 allocate_gpu_memory, tmp_size_tracker);

    auto zero_out_if_not_1_lut_f = [](Torus x) -> Torus {
      Torus block = x / 2;
      bool condition = (x & 1) == 1;
      return block * (Torus)condition;
    };
    auto zero_out_if_not_2_lut_f = [](Torus x) -> Torus {
      Torus block = x / 3;
      bool condition = (x % 3) == 2;
      return block * (Torus)condition;
    };

    int_radix_lut<Torus> *luts[2] = {zero_out_if_not_1_lut_1,
                                     zero_out_if_not_1_lut_2};
    size_t lut_gpu_indexes[2] = {0, 3};
    for (int j = 0; j < 2; j++) {
      generate_device_accumulator<Torus>(
          streams.stream(lut_gpu_indexes[j]),
          streams.gpu_index(lut_gpu_indexes[j]), luts[j]->get_lut(0, 0),
          luts[j]->get_degree(0), luts[j]->get_max_degree(0),
          params.glwe_dimension, params.polynomial_size, params.message_modulus,
          params.carry_modulus, zero_out_if_not_1_lut_f, gpu_memory_allocated);
    }

    luts[0] = zero_out_if_not_2_lut_1;
    luts[1] = zero_out_if_not_2_lut_2;
    lut_gpu_indexes[0] = 1;
    lut_gpu_indexes[1] = 2;
    for (int j = 0; j < 2; j++) {
      generate_device_accumulator<Torus>(
          streams.stream(lut_gpu_indexes[j]),
          streams.gpu_index(lut_gpu_indexes[j]), luts[j]->get_lut(0, 0),
          luts[j]->get_degree(0), luts[j]->get_max_degree(0),
          params.glwe_dimension, params.polynomial_size, params.message_modulus,
          params.carry_modulus, zero_out_if_not_2_lut_f, gpu_memory_allocated);
    }

    quotient_lut_1 =
        new int_radix_lut<Torus>(streams.get_ith(2), params, 1, 1,
                                 allocate_gpu_memory, tmp_size_tracker);
    quotient_lut_2 =
        new int_radix_lut<Torus>(streams.get_ith(1), params, 1, 1,
                                 allocate_gpu_memory, tmp_size_tracker);
    quotient_lut_3 = new int_radix_lut<Torus>(
        streams.get_ith(0), params, 1, 1, allocate_gpu_memory, size_tracker);

    auto quotient_lut_1_f = [](Torus cond) -> Torus {
      return (Torus)(cond == 2);
    };
    auto quotient_lut_2_f = [](Torus cond) -> Torus {
      return (Torus)((cond == 2) * 2);
    };
    auto quotient_lut_3_f = [](Torus cond) -> Torus { return cond * 3; };

    generate_device_accumulator<Torus>(
        streams.stream(2), streams.gpu_index(2), quotient_lut_1->get_lut(0, 0),
        quotient_lut_1->get_degree(0), quotient_lut_1->get_max_degree(0),
        params.glwe_dimension, params.polynomial_size, params.message_modulus,
        params.carry_modulus, quotient_lut_1_f, gpu_memory_allocated);
    generate_device_accumulator<Torus>(
        streams.stream(1), streams.gpu_index(1), quotient_lut_2->get_lut(0, 0),
        quotient_lut_2->get_degree(0), quotient_lut_2->get_max_degree(0),
        params.glwe_dimension, params.polynomial_size, params.message_modulus,
        params.carry_modulus, quotient_lut_2_f, gpu_memory_allocated);
    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0), quotient_lut_3->get_lut(0, 0),
        quotient_lut_3->get_degree(0), quotient_lut_3->get_max_degree(0),
        params.glwe_dimension, params.polynomial_size, params.message_modulus,
        params.carry_modulus, quotient_lut_3_f, gpu_memory_allocated);

    message_extract_lut_1 = new int_radix_lut<Torus>(
        streams, params, 1, num_blocks, allocate_gpu_memory, size_tracker);
    message_extract_lut_2 = new int_radix_lut<Torus>(
        streams, params, 1, num_blocks, allocate_gpu_memory, size_tracker);

    auto message_modulus = params.message_modulus;
    auto lut_f_message_extract = [message_modulus](Torus x) -> Torus {
      return x % message_modulus;
    };

    luts[0] = message_extract_lut_1;
    luts[1] = message_extract_lut_2;

    for (int j = 0; j < 2; j++) {
      generate_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0), luts[j]->get_lut(0, 0),
          luts[j]->get_degree(0), luts[j]->get_max_degree(0),
          params.glwe_dimension, params.polynomial_size, params.message_modulus,
          params.carry_modulus, lut_f_message_extract, gpu_memory_allocated);
      auto active_streams = streams.active_gpu_subset(num_blocks);
      luts[j]->broadcast_lut(active_streams);
    }
  }

  unsigned_int_div_rem_2_2_memory(CudaStreams streams, int_radix_params params,
                                  uint32_t num_blocks, bool allocate_gpu_memory,
                                  uint64_t &size_tracker) {
    if (streams.count() < 4) {
      PANIC("GPU count should be greater than 4m when using div_rem_2_2");
    }

    if (params.message_modulus != 4 || params.carry_modulus != 4) {
      PANIC("Only message_modulus == 4 && carry_modulus == 4 parameters are "
            "supported");
    }

    this->params = params;
    gpu_memory_allocated = allocate_gpu_memory;

    sub_and_propagate_mem = new int_sub_and_propagate<Torus>(
        streams.get_ith(0), params, num_blocks + 1, outputFlag::FLAG_NONE,
        allocate_gpu_memory, size_tracker);

    shift_mem = new int_logical_scalar_shift_buffer<Torus>(
        streams.get_ith(1), SHIFT_OR_ROTATE_TYPE::LEFT_SHIFT, params,
        2 * num_blocks, allocate_gpu_memory, tmp_size_tracker);

    uint32_t compute_overflow = 1;
    overflow_sub_mem_1 = new int_borrow_prop_memory<Torus>(
        streams.get_ith(0), params, num_blocks, compute_overflow,
        allocate_gpu_memory, size_tracker);
    overflow_sub_mem_2 = new int_borrow_prop_memory<Torus>(
        streams.get_ith(1), params, num_blocks, compute_overflow,
        allocate_gpu_memory, tmp_size_tracker);
    overflow_sub_mem_3 = new int_borrow_prop_memory<Torus>(
        streams.get_ith(2), params, num_blocks, compute_overflow,
        allocate_gpu_memory, tmp_size_tracker);
    uint32_t group_size = overflow_sub_mem_1->group_size;
    bool use_seq = overflow_sub_mem_1->prop_simu_group_carries_mem
                       ->use_sequential_algorithm_to_resolve_group_carries;

    cuda_set_device(0);
    cudaEventCreateWithFlags(&create_indexes_done, cudaEventDisableTiming);
    create_indexes_for_overflow_sub(streams.get_ith(0), num_blocks, group_size,
                                    use_seq, allocate_gpu_memory, size_tracker);
    cudaEventRecord(create_indexes_done, streams.stream(0));
    cuda_set_device(1);
    cudaStreamWaitEvent(streams.stream(1), create_indexes_done, 0);
    cuda_set_device(2);
    cudaStreamWaitEvent(streams.stream(2), create_indexes_done, 0);

    scatter_indexes_for_overflowing_sub(
        streams.stream(1), streams.gpu_index(1),
        &first_indexes_for_overflow_sub_gpu_1,
        &second_indexes_for_overflow_sub_gpu_1, &scalars_for_overflow_sub_gpu_1,
        num_blocks, allocate_gpu_memory, tmp_size_tracker);
    scatter_indexes_for_overflowing_sub(
        streams.stream(2), streams.gpu_index(2),
        &first_indexes_for_overflow_sub_gpu_2,
        &second_indexes_for_overflow_sub_gpu_2, &scalars_for_overflow_sub_gpu_2,
        num_blocks, allocate_gpu_memory, tmp_size_tracker);

    comparison_buffer_1 = new int_comparison_buffer<Torus>(
        streams.get_ith(0), COMPARISON_TYPE::EQ, params, num_blocks, false,
        allocate_gpu_memory, size_tracker);
    comparison_buffer_2 = new int_comparison_buffer<Torus>(
        streams.get_ith(1), COMPARISON_TYPE::EQ, params, num_blocks, false,
        allocate_gpu_memory, tmp_size_tracker);
    comparison_buffer_3 = new int_comparison_buffer<Torus>(
        streams.get_ith(2), COMPARISON_TYPE::EQ, params, num_blocks, false,
        allocate_gpu_memory, tmp_size_tracker);
    bitor_mem_1 = new int_bitop_buffer<Torus>(
        streams.get_ith(0), BITOP_TYPE::BITOR, params, num_blocks,
        allocate_gpu_memory, size_tracker);
    bitor_mem_2 = new int_bitop_buffer<Torus>(
        streams.get_ith(1), BITOP_TYPE::BITOR, params, num_blocks,
        allocate_gpu_memory, tmp_size_tracker);
    bitor_mem_3 = new int_bitop_buffer<Torus>(
        streams.get_ith(2), BITOP_TYPE::BITOR, params, num_blocks,
        allocate_gpu_memory, tmp_size_tracker);

    init_lookup_tables(streams, num_blocks, allocate_gpu_memory, size_tracker);
    init_temporary_buffers(streams, num_blocks, allocate_gpu_memory,
                           size_tracker);

    sub_streams_1.create_on_same_gpus(streams);
  }

  void scatter_indexes_for_overflowing_sub(
      cudaStream_t const stream, size_t gpu_index, Torus ***first_indexes_ptr,
      Torus ***second_indexes_ptr, Torus ***scalars_ptr, uint32_t num_blocks,
      bool allocate_gpu_memory, uint64_t &size_tracker) {

    auto first_indexes = (Torus **)malloc(num_blocks * sizeof(Torus *));
    auto second_indexes = (Torus **)malloc(num_blocks * sizeof(Torus *));
    auto scalars = (Torus **)malloc(num_blocks * sizeof(Torus *));

    for (int nb = 1; nb <= num_blocks; nb++) {
      first_indexes[nb - 1] = (Torus *)cuda_malloc_with_size_tracking_async(
          nb * sizeof(Torus), stream, gpu_index, size_tracker,
          allocate_gpu_memory);
      second_indexes[nb - 1] = (Torus *)cuda_malloc_with_size_tracking_async(
          nb * sizeof(Torus), stream, gpu_index, size_tracker,
          allocate_gpu_memory);
      scalars[nb - 1] = (Torus *)cuda_malloc_with_size_tracking_async(
          nb * sizeof(Torus), stream, gpu_index, size_tracker,
          allocate_gpu_memory);

      cuda_memcpy_with_size_tracking_async_gpu_to_gpu(
          first_indexes[nb - 1], first_indexes_for_overflow_sub_gpu_0[nb - 1],
          nb * sizeof(Torus), stream, gpu_index, allocate_gpu_memory);
      cuda_memcpy_with_size_tracking_async_gpu_to_gpu(
          second_indexes[nb - 1], second_indexes_for_overflow_sub_gpu_0[nb - 1],
          nb * sizeof(Torus), stream, gpu_index, allocate_gpu_memory);
      cuda_memcpy_with_size_tracking_async_gpu_to_gpu(
          scalars[nb - 1], scalars_for_overflow_sub_gpu_0[nb - 1],
          nb * sizeof(Torus), stream, gpu_index, allocate_gpu_memory);
      *first_indexes_ptr = first_indexes;
      *second_indexes_ptr = second_indexes;
      *scalars_ptr = scalars;
    }
  }

  void create_indexes_for_overflow_sub(CudaStreams streams, uint32_t num_blocks,
                                       uint32_t group_size, bool use_seq,
                                       bool allocate_gpu_memory,
                                       uint64_t &size_tracker) {
    max_indexes_to_erase = num_blocks;

    first_indexes_for_overflow_sub_gpu_0 =
        (Torus **)malloc(num_blocks * sizeof(Torus *));
    second_indexes_for_overflow_sub_gpu_0 =
        (Torus **)malloc(num_blocks * sizeof(Torus *));
    scalars_for_overflow_sub_gpu_0 =
        (Torus **)malloc(num_blocks * sizeof(Torus *));

    Torus *h_lut_indexes = (Torus *)malloc(num_blocks * sizeof(Torus));
    Torus *h_scalar = (Torus *)malloc(num_blocks * sizeof(Torus));

    // Extra indexes for the luts in first step
    for (int nb = 1; nb <= num_blocks; nb++) {
      first_indexes_for_overflow_sub_gpu_0[nb - 1] =
          (Torus *)cuda_malloc_with_size_tracking_async(
              nb * sizeof(Torus), streams.stream(0), streams.gpu_index(0),
              size_tracker, allocate_gpu_memory);
      for (int index = 0; index < nb; index++) {
        uint32_t grouping_index = index / group_size;
        bool is_in_first_grouping = (grouping_index == 0);
        uint32_t index_in_grouping = index % group_size;
        bool is_last_index = (index == (nb - 1));
        if (is_last_index) {
          if (nb == 1) {
            h_lut_indexes[index] = 2 * group_size;
          } else {
            h_lut_indexes[index] = 2;
          }
        } else if (is_in_first_grouping) {
          h_lut_indexes[index] = index_in_grouping;
        } else {
          h_lut_indexes[index] = index_in_grouping + group_size;
        }
      }
      cuda_memcpy_with_size_tracking_async_to_gpu(
          first_indexes_for_overflow_sub_gpu_0[nb - 1], h_lut_indexes,
          nb * sizeof(Torus), streams.stream(0), streams.gpu_index(0),
          allocate_gpu_memory);
    }
    // Extra indexes for the luts in second step
    for (int nb = 1; nb <= num_blocks; nb++) {
      second_indexes_for_overflow_sub_gpu_0[nb - 1] =
          (Torus *)cuda_malloc_with_size_tracking_async(
              nb * sizeof(Torus), streams.stream(0), streams.gpu_index(0),
              size_tracker, allocate_gpu_memory);
      scalars_for_overflow_sub_gpu_0[nb - 1] =
          (Torus *)cuda_malloc_with_size_tracking_async(
              nb * sizeof(Torus), streams.stream(0), streams.gpu_index(0),
              size_tracker, allocate_gpu_memory);

      for (int index = 0; index < nb; index++) {
        uint32_t grouping_index = index / group_size;
        bool is_in_first_grouping = (grouping_index == 0);
        uint32_t index_in_grouping = index % group_size;

        if (is_in_first_grouping) {
          h_lut_indexes[index] = index_in_grouping;
        } else if (index_in_grouping == (group_size - 1)) {
          if (use_seq) {
            int inner_index = (grouping_index - 1) % (group_size - 1);
            h_lut_indexes[index] = inner_index + 2 * group_size;
          } else {
            h_lut_indexes[index] = 2 * group_size;
          }
        } else {
          h_lut_indexes[index] = index_in_grouping + group_size;
        }

        bool may_have_its_padding_bit_set =
            !is_in_first_grouping && (index_in_grouping == group_size - 1);

        if (may_have_its_padding_bit_set) {
          if (use_seq) {
            h_scalar[index] = 1 << ((grouping_index - 1) % (group_size - 1));
          } else {
            h_scalar[index] = 1;
          }
        } else {
          h_scalar[index] = 0;
        }
      }
      cuda_memcpy_with_size_tracking_async_to_gpu(
          second_indexes_for_overflow_sub_gpu_0[nb - 1], h_lut_indexes,
          nb * sizeof(Torus), streams.stream(0), streams.gpu_index(0),
          allocate_gpu_memory);
      cuda_memcpy_with_size_tracking_async_to_gpu(
          scalars_for_overflow_sub_gpu_0[nb - 1], h_scalar, nb * sizeof(Torus),
          streams.stream(0), streams.gpu_index(0), allocate_gpu_memory);
    }
    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
    free(h_lut_indexes);
    free(h_scalar);
  };

  void release(CudaStreams streams) {

    if (streams.count() < 4) {
      PANIC("GPU count should be greater than 4 when using div_rem_2_2");
    }
    if (params.message_modulus != 4 || params.carry_modulus != 4) {
      PANIC("Only message_modulus == 4 && carry_modulus == 4 parameters are "
            "supported");
    }

    // release and delete integer ops memory objects
    sub_and_propagate_mem->release(streams.get_ith(0));
    shift_mem->release(streams.get_ith(1));
    overflow_sub_mem_1->release(streams.get_ith(0));
    overflow_sub_mem_2->release(streams.get_ith(1));
    overflow_sub_mem_3->release(streams.get_ith(2));
    comparison_buffer_1->release(streams.get_ith(0));
    comparison_buffer_2->release(streams.get_ith(1));
    comparison_buffer_3->release(streams.get_ith(2));
    bitor_mem_1->release(streams.get_ith(0));
    bitor_mem_2->release(streams.get_ith(1));
    bitor_mem_3->release(streams.get_ith(2));

    delete sub_and_propagate_mem;
    sub_and_propagate_mem = nullptr;
    delete shift_mem;
    shift_mem = nullptr;
    delete overflow_sub_mem_1;
    overflow_sub_mem_1 = nullptr;
    delete overflow_sub_mem_2;
    overflow_sub_mem_2 = nullptr;
    delete overflow_sub_mem_3;
    overflow_sub_mem_3 = nullptr;
    delete comparison_buffer_1;
    comparison_buffer_1 = nullptr;
    delete comparison_buffer_2;
    comparison_buffer_2 = nullptr;
    delete comparison_buffer_3;
    comparison_buffer_3 = nullptr;
    delete bitor_mem_1;
    bitor_mem_1 = nullptr;
    delete bitor_mem_2;
    bitor_mem_2 = nullptr;
    delete bitor_mem_3;
    bitor_mem_3 = nullptr;

    // release and delete lut objects
    message_extract_lut_1->release(streams);
    message_extract_lut_2->release(streams);
    zero_out_if_not_1_lut_1->release(streams.get_ith(0));
    zero_out_if_not_1_lut_2->release(streams.get_ith(3));
    zero_out_if_not_2_lut_1->release(streams.get_ith(1));
    zero_out_if_not_2_lut_2->release(streams.get_ith(2));
    quotient_lut_1->release(streams.get_ith(2));
    quotient_lut_2->release(streams.get_ith(1));
    quotient_lut_3->release(streams.get_ith(0));

    delete message_extract_lut_1;
    message_extract_lut_1 = nullptr;
    delete message_extract_lut_2;
    message_extract_lut_2 = nullptr;
    delete zero_out_if_not_1_lut_1;
    zero_out_if_not_1_lut_1 = nullptr;
    delete zero_out_if_not_1_lut_2;
    zero_out_if_not_1_lut_2 = nullptr;
    delete zero_out_if_not_2_lut_1;
    zero_out_if_not_2_lut_1 = nullptr;
    delete zero_out_if_not_2_lut_2;
    zero_out_if_not_2_lut_2 = nullptr;
    delete quotient_lut_1;
    quotient_lut_1 = nullptr;
    delete quotient_lut_2;
    quotient_lut_2 = nullptr;
    delete quotient_lut_3;
    quotient_lut_3 = nullptr;

    // release and delete temporary buffers
    release_radix_ciphertext_async(streams.stream(2), streams.gpu_index(2), d1,
                                   gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(1), streams.gpu_index(1), d2,
                                   gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0), d3,
                                   gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(2), streams.gpu_index(2),
                                   low1, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(1), streams.gpu_index(1),
                                   low2, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   low3, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(3), streams.gpu_index(3),
                                   rem0, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(2), streams.gpu_index(2),
                                   rem1, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(1), streams.gpu_index(1),
                                   rem2, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   rem3, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(2), streams.gpu_index(2),
                                   sub_result_3, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(1), streams.gpu_index(1),
                                   sub_result_2, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   sub_result_1, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(2), streams.gpu_index(2),
                                   sub_3_overflowed, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(1), streams.gpu_index(1),
                                   sub_2_overflowed, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   sub_1_overflowed, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(3), streams.gpu_index(3),
                                   tmp_gpu_3, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(2), streams.gpu_index(2),
                                   tmp_gpu_2, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(1), streams.gpu_index(1),
                                   tmp_gpu_1, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_gpu_0, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(1), streams.gpu_index(1),
                                   divisor_gpu_1, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(2), streams.gpu_index(2),
                                   divisor_gpu_2, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(1), streams.gpu_index(1),
                                   remainder_gpu_1, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(2), streams.gpu_index(2),
                                   remainder_gpu_2, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(3), streams.gpu_index(3),
                                   remainder_gpu_3, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(2), streams.gpu_index(2),
                                   comparison_blocks_3, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(1), streams.gpu_index(1),
                                   comparison_blocks_2, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   comparison_blocks_1, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(2), streams.gpu_index(2),
                                   cmp_3, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(1), streams.gpu_index(1),
                                   cmp_2, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   cmp_1, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(3), streams.gpu_index(3), c0,
                                   gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(2), streams.gpu_index(2), q1,
                                   gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(1), streams.gpu_index(1), q2,
                                   gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0), q3,
                                   gpu_memory_allocated);

    delete d1;
    d1 = nullptr;
    delete d2;
    d2 = nullptr;
    delete d3;
    d3 = nullptr;
    delete low1;
    low1 = nullptr;
    delete low2;
    low2 = nullptr;
    delete low3;
    low3 = nullptr;
    delete rem0;
    rem0 = nullptr;
    delete rem1;
    rem1 = nullptr;
    delete rem2;
    rem2 = nullptr;
    delete rem3;
    rem3 = nullptr;
    delete sub_result_1;
    sub_result_1 = nullptr;
    delete sub_result_2;
    sub_result_2 = nullptr;
    delete sub_result_3;
    sub_result_3 = nullptr;
    delete sub_1_overflowed;
    sub_1_overflowed = nullptr;
    delete sub_2_overflowed;
    sub_2_overflowed = nullptr;
    delete sub_3_overflowed;
    sub_3_overflowed = nullptr;
    delete tmp_gpu_0;
    tmp_gpu_0 = nullptr;
    delete tmp_gpu_1;
    tmp_gpu_1 = nullptr;
    delete tmp_gpu_2;
    tmp_gpu_2 = nullptr;
    delete tmp_gpu_3;
    tmp_gpu_3 = nullptr;
    delete divisor_gpu_1;
    divisor_gpu_1 = nullptr;
    delete divisor_gpu_2;
    divisor_gpu_2 = nullptr;
    delete remainder_gpu_1;
    remainder_gpu_1 = nullptr;
    delete remainder_gpu_2;
    remainder_gpu_2 = nullptr;
    delete remainder_gpu_3;
    remainder_gpu_3 = nullptr;
    delete comparison_blocks_1;
    comparison_blocks_1 = nullptr;
    delete comparison_blocks_2;
    comparison_blocks_2 = nullptr;
    delete comparison_blocks_3;
    comparison_blocks_3 = nullptr;
    delete cmp_1;
    cmp_1 = nullptr;
    delete cmp_2;
    cmp_2 = nullptr;
    delete cmp_3;
    cmp_3 = nullptr;
    delete c0;
    c0 = nullptr;
    delete q1;
    q1 = nullptr;
    delete q2;
    q2 = nullptr;
    delete q3;
    q3 = nullptr;

    for (int i = 0; i < max_indexes_to_erase; i++) {
      cuda_drop_with_size_tracking_async(
          first_indexes_for_overflow_sub_gpu_0[i], streams.stream(0),
          streams.gpu_index(0), gpu_memory_allocated);
      cuda_drop_with_size_tracking_async(
          second_indexes_for_overflow_sub_gpu_0[i], streams.stream(0),
          streams.gpu_index(0), gpu_memory_allocated);
      cuda_drop_with_size_tracking_async(
          scalars_for_overflow_sub_gpu_0[i], streams.stream(0),
          streams.gpu_index(0), gpu_memory_allocated);
      cuda_drop_with_size_tracking_async(
          first_indexes_for_overflow_sub_gpu_1[i], streams.stream(1),
          streams.gpu_index(1), gpu_memory_allocated);
      cuda_drop_with_size_tracking_async(
          second_indexes_for_overflow_sub_gpu_1[i], streams.stream(1),
          streams.gpu_index(1), gpu_memory_allocated);
      cuda_drop_with_size_tracking_async(
          scalars_for_overflow_sub_gpu_1[i], streams.stream(1),
          streams.gpu_index(1), gpu_memory_allocated);
      cuda_drop_with_size_tracking_async(
          first_indexes_for_overflow_sub_gpu_2[i], streams.stream(2),
          streams.gpu_index(2), gpu_memory_allocated);
      cuda_drop_with_size_tracking_async(
          second_indexes_for_overflow_sub_gpu_2[i], streams.stream(2),
          streams.gpu_index(2), gpu_memory_allocated);
      cuda_drop_with_size_tracking_async(
          scalars_for_overflow_sub_gpu_2[i], streams.stream(2),
          streams.gpu_index(2), gpu_memory_allocated);
    }

    for (uint j = 0; j < 3; j++) {
      cuda_synchronize_stream(streams.stream(j), streams.gpu_index(j));
    }

    free(first_indexes_for_overflow_sub_gpu_0);
    free(second_indexes_for_overflow_sub_gpu_0);
    free(scalars_for_overflow_sub_gpu_0);
    free(first_indexes_for_overflow_sub_gpu_1);
    free(second_indexes_for_overflow_sub_gpu_1);
    free(scalars_for_overflow_sub_gpu_1);
    free(first_indexes_for_overflow_sub_gpu_2);
    free(second_indexes_for_overflow_sub_gpu_2);
    free(scalars_for_overflow_sub_gpu_2);

    cudaEventDestroy(create_indexes_done);

    // release sub streams
    sub_streams_1.release();
  }
};

template <typename Torus> struct unsigned_int_div_rem_memory {
  int_radix_params params;

  // memory objects for other operations
  int_logical_scalar_shift_buffer<Torus> *shift_mem_1;
  int_logical_scalar_shift_buffer<Torus> *shift_mem_2;
  int_borrow_prop_memory<Torus> *overflow_sub_mem;
  int_comparison_buffer<Torus> *comparison_buffer;
  unsigned_int_div_rem_2_2_memory<Torus> *div_rem_2_2_mem;

  // lookup tables
  int_radix_lut<Torus> **masking_luts_1;
  int_radix_lut<Torus> **masking_luts_2;
  int_radix_lut<Torus> *message_extract_lut_1;
  int_radix_lut<Torus> *message_extract_lut_2;
  int_radix_lut<Torus> **zero_out_if_overflow_did_not_happen;
  int_radix_lut<Torus> **zero_out_if_overflow_happened;
  int_radix_lut<Torus> **merge_overflow_flags_luts;

  // sub streams
  CudaStreams sub_streams_1;
  CudaStreams sub_streams_2;
  CudaStreams sub_streams_3;
  CudaStreams sub_streams_4;

  // temporary device buffers
  CudaRadixCiphertextFFI *remainder1;
  CudaRadixCiphertextFFI *remainder2;
  CudaRadixCiphertextFFI *numerator_block_stack;
  CudaRadixCiphertextFFI *numerator_block_1;
  CudaRadixCiphertextFFI *tmp_radix;
  CudaRadixCiphertextFFI *interesting_remainder1;
  CudaRadixCiphertextFFI *interesting_remainder2;
  CudaRadixCiphertextFFI *interesting_divisor;
  CudaRadixCiphertextFFI *divisor_ms_blocks;
  CudaRadixCiphertextFFI *new_remainder;
  CudaRadixCiphertextFFI *subtraction_overflowed;
  CudaRadixCiphertextFFI *did_not_overflow;
  CudaRadixCiphertextFFI *overflow_sum;
  CudaRadixCiphertextFFI *overflow_sum_radix;
  CudaRadixCiphertextFFI *tmp_1;
  CudaRadixCiphertextFFI *at_least_one_upper_block_is_non_zero;
  CudaRadixCiphertextFFI *cleaned_merged_interesting_remainder;

  Torus **first_indexes_for_overflow_sub;
  Torus **second_indexes_for_overflow_sub;
  Torus **scalars_for_overflow_sub;
  uint32_t max_indexes_to_erase;
  bool gpu_memory_allocated;

  // allocate and initialize if needed, temporary arrays used to calculate
  // cuda integer div_rem operation
  void init_temporary_buffers(CudaStreams streams, uint32_t num_blocks,
                              bool allocate_gpu_memory,
                              uint64_t &size_tracker) {

    // non boolean temporary arrays, with `num_blocks` blocks
    remainder1 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), remainder1, num_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    remainder2 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), remainder2, num_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    numerator_block_stack = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), numerator_block_stack,
        num_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    interesting_remainder2 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), interesting_remainder2,
        num_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    interesting_divisor = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), interesting_divisor,
        num_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    divisor_ms_blocks = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), divisor_ms_blocks, num_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    new_remainder = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), new_remainder, num_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    cleaned_merged_interesting_remainder = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0),
        cleaned_merged_interesting_remainder, num_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    tmp_1 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_1, num_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

    // temporary arrays used as stacks
    tmp_radix = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_radix, num_blocks + 1,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    interesting_remainder1 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), interesting_remainder1,
        num_blocks + 1, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
    numerator_block_1 = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), numerator_block_1, 1,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

    // temporary arrays for boolean blocks
    subtraction_overflowed = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), subtraction_overflowed, 1,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    did_not_overflow = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), did_not_overflow, 1,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    overflow_sum = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), overflow_sum, 1,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    overflow_sum_radix = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), overflow_sum_radix, num_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    at_least_one_upper_block_is_non_zero = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0),
        at_least_one_upper_block_is_non_zero, 1, params.big_lwe_dimension,
        size_tracker, allocate_gpu_memory);
  }

  // initialize lookup tables for div_rem operation
  void init_lookup_tables(CudaStreams streams, uint32_t num_blocks,
                          bool allocate_gpu_memory, uint64_t &size_tracker) {
    uint32_t num_bits_in_message = 31 - __builtin_clz(params.message_modulus);

    // create and generate masking_luts_1[] and masking_lut_2[]
    // both of them are equal but because they are used in two different
    // executions in parallel we need two different pbs_buffers.
    masking_luts_1 = new int_radix_lut<Torus> *[params.message_modulus - 1];
    masking_luts_2 = new int_radix_lut<Torus> *[params.message_modulus - 1];
    for (int i = 0; i < params.message_modulus - 1; i++) {
      uint32_t shifted_mask = i;
      std::function<Torus(Torus)> lut_f_masking =
          [shifted_mask](Torus x) -> Torus { return x & shifted_mask; };

      masking_luts_1[i] = new int_radix_lut<Torus>(
          streams, params, 1, 1, allocate_gpu_memory, size_tracker);
      masking_luts_2[i] = new int_radix_lut<Torus>(
          streams, params, 1, num_blocks, allocate_gpu_memory, size_tracker);

      generate_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0),
          masking_luts_1[i]->get_lut(0, 0), masking_luts_1[i]->get_degree(0),
          masking_luts_1[i]->get_max_degree(0), params.glwe_dimension,
          params.polynomial_size, params.message_modulus, params.carry_modulus,
          lut_f_masking, gpu_memory_allocated);
      auto active_streams_1 = streams.active_gpu_subset(1);
      masking_luts_1[i]->broadcast_lut(active_streams_1);

      generate_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0),
          masking_luts_2[i]->get_lut(0, 0), masking_luts_2[i]->get_degree(0),
          masking_luts_2[i]->get_max_degree(0), params.glwe_dimension,
          params.polynomial_size, params.message_modulus, params.carry_modulus,
          lut_f_masking, gpu_memory_allocated);
      auto active_streams_2 = streams.active_gpu_subset(num_blocks);
      masking_luts_2[i]->broadcast_lut(active_streams_2);
    }

    // create and generate message_extract_lut_1 and message_extract_lut_2
    // both of them are equal but because they are used in two different
    // executions in parallel we need two different pbs_buffers.
    message_extract_lut_1 = new int_radix_lut<Torus>(
        streams, params, 1, num_blocks, allocate_gpu_memory, size_tracker);
    message_extract_lut_2 = new int_radix_lut<Torus>(
        streams, params, 1, num_blocks, allocate_gpu_memory, size_tracker);

    auto message_modulus = params.message_modulus;
    auto lut_f_message_extract = [message_modulus](Torus x) -> Torus {
      return x % message_modulus;
    };

    int_radix_lut<Torus> *luts[2] = {message_extract_lut_1,
                                     message_extract_lut_2};
    auto active_streams = streams.active_gpu_subset(num_blocks);
    for (int j = 0; j < 2; j++) {
      generate_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0), luts[j]->get_lut(0, 0),
          luts[j]->get_degree(0), luts[j]->get_max_degree(0),
          params.glwe_dimension, params.polynomial_size, params.message_modulus,
          params.carry_modulus, lut_f_message_extract, gpu_memory_allocated);
      luts[j]->broadcast_lut(active_streams);
    }

    // Give name to closures to improve readability
    auto overflow_happened = [](uint64_t overflow_sum) {
      return overflow_sum != 0;
    };
    auto overflow_did_not_happen = [&overflow_happened](uint64_t overflow_sum) {
      return !overflow_happened(overflow_sum);
    };

    // create and generate zero_out_if_overflow_did_not_happen
    zero_out_if_overflow_did_not_happen = new int_radix_lut<Torus> *[2];
    zero_out_if_overflow_did_not_happen[0] = new int_radix_lut<Torus>(
        streams, params, 1, num_blocks, allocate_gpu_memory, size_tracker);
    zero_out_if_overflow_did_not_happen[1] = new int_radix_lut<Torus>(
        streams, params, 1, num_blocks, allocate_gpu_memory, size_tracker);

    auto cur_lut_f = [&](Torus block, Torus overflow_sum) -> Torus {
      if (overflow_did_not_happen(overflow_sum)) {
        return 0;
      } else {
        return block;
      }
    };

    generate_device_accumulator_bivariate_with_factor<Torus>(
        streams.stream(0), streams.gpu_index(0),
        zero_out_if_overflow_did_not_happen[0]->get_lut(0, 0),
        zero_out_if_overflow_did_not_happen[0]->get_degree(0),
        zero_out_if_overflow_did_not_happen[0]->get_max_degree(0),
        params.glwe_dimension, params.polynomial_size, params.message_modulus,
        params.carry_modulus, cur_lut_f, params.message_modulus - 2,
        gpu_memory_allocated);
    zero_out_if_overflow_did_not_happen[0]->broadcast_lut(active_streams);
    generate_device_accumulator_bivariate_with_factor<Torus>(
        streams.stream(0), streams.gpu_index(0),
        zero_out_if_overflow_did_not_happen[1]->get_lut(0, 0),
        zero_out_if_overflow_did_not_happen[1]->get_degree(0),
        zero_out_if_overflow_did_not_happen[1]->get_max_degree(0),
        params.glwe_dimension, params.polynomial_size, params.message_modulus,
        params.carry_modulus, cur_lut_f, params.message_modulus - 1,
        gpu_memory_allocated);
    zero_out_if_overflow_did_not_happen[1]->broadcast_lut(active_streams);

    // create and generate zero_out_if_overflow_happened
    zero_out_if_overflow_happened = new int_radix_lut<Torus> *[2];
    zero_out_if_overflow_happened[0] = new int_radix_lut<Torus>(
        streams, params, 1, num_blocks, allocate_gpu_memory, size_tracker);
    zero_out_if_overflow_happened[1] = new int_radix_lut<Torus>(
        streams, params, 1, num_blocks, allocate_gpu_memory, size_tracker);

    auto overflow_happened_f = [&](Torus block, Torus overflow_sum) -> Torus {
      if (overflow_happened(overflow_sum)) {
        return 0;
      } else {
        return block;
      }
    };

    generate_device_accumulator_bivariate_with_factor<Torus>(
        streams.stream(0), streams.gpu_index(0),
        zero_out_if_overflow_happened[0]->get_lut(0, 0),
        zero_out_if_overflow_happened[0]->get_degree(0),
        zero_out_if_overflow_happened[0]->get_max_degree(0),
        params.glwe_dimension, params.polynomial_size, params.message_modulus,
        params.carry_modulus, overflow_happened_f, params.message_modulus - 2,
        gpu_memory_allocated);
    zero_out_if_overflow_happened[0]->broadcast_lut(active_streams);
    generate_device_accumulator_bivariate_with_factor<Torus>(
        streams.stream(0), streams.gpu_index(0),
        zero_out_if_overflow_happened[1]->get_lut(0, 0),
        zero_out_if_overflow_happened[1]->get_degree(0),
        zero_out_if_overflow_happened[1]->get_max_degree(0),
        params.glwe_dimension, params.polynomial_size, params.message_modulus,
        params.carry_modulus, overflow_happened_f, params.message_modulus - 1,
        gpu_memory_allocated);
    zero_out_if_overflow_happened[1]->broadcast_lut(active_streams);

    // merge_overflow_flags_luts
    merge_overflow_flags_luts = new int_radix_lut<Torus> *[num_bits_in_message];
    auto active_gpu_count_for_bits = streams.active_gpu_subset(1);
    for (int i = 0; i < num_bits_in_message; i++) {
      auto lut_f_bit = [i](Torus x, Torus y) -> Torus {
        return (x == 0 && y == 0) << i;
      };

      merge_overflow_flags_luts[i] = new int_radix_lut<Torus>(
          streams, params, 1, 1, allocate_gpu_memory, size_tracker);

      generate_device_accumulator_bivariate<Torus>(
          streams.stream(0), streams.gpu_index(0),
          merge_overflow_flags_luts[i]->get_lut(0, 0),
          merge_overflow_flags_luts[i]->get_degree(0),
          merge_overflow_flags_luts[i]->get_max_degree(0),
          params.glwe_dimension, params.polynomial_size, params.message_modulus,
          params.carry_modulus, lut_f_bit, gpu_memory_allocated);
      merge_overflow_flags_luts[i]->broadcast_lut(active_gpu_count_for_bits);
    }
  }

  unsigned_int_div_rem_memory(CudaStreams streams, int_radix_params params,
                              uint32_t num_blocks, bool allocate_gpu_memory,
                              uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;
    auto active_streams = streams.active_gpu_subset(2 * num_blocks);
    this->params = params;

    if (params.message_modulus == 4 && params.carry_modulus == 4 &&
        streams.count() >= 4) {
      div_rem_2_2_mem = new unsigned_int_div_rem_2_2_memory<Torus>(
          streams, params, num_blocks, allocate_gpu_memory, size_tracker);
      return;
    }

    shift_mem_1 = new int_logical_scalar_shift_buffer<Torus>(
        streams, SHIFT_OR_ROTATE_TYPE::LEFT_SHIFT, params, 2 * num_blocks,
        allocate_gpu_memory, size_tracker);

    shift_mem_2 = new int_logical_scalar_shift_buffer<Torus>(
        streams, SHIFT_OR_ROTATE_TYPE::LEFT_SHIFT, params, 2 * num_blocks,
        allocate_gpu_memory, size_tracker);

    uint32_t compute_overflow = 1;
    overflow_sub_mem = new int_borrow_prop_memory<Torus>(
        streams, params, num_blocks, compute_overflow, allocate_gpu_memory,
        size_tracker);
    uint32_t group_size = overflow_sub_mem->group_size;
    bool use_seq = overflow_sub_mem->prop_simu_group_carries_mem
                       ->use_sequential_algorithm_to_resolve_group_carries;
    create_indexes_for_overflow_sub(streams, num_blocks, group_size, use_seq,
                                    allocate_gpu_memory, size_tracker);

    comparison_buffer = new int_comparison_buffer<Torus>(
        streams, COMPARISON_TYPE::NE, params, num_blocks, false,
        allocate_gpu_memory, size_tracker);

    init_lookup_tables(streams, num_blocks, allocate_gpu_memory, size_tracker);
    init_temporary_buffers(streams, num_blocks, allocate_gpu_memory,
                           size_tracker);

    sub_streams_1.create_on_same_gpus(active_streams);
    sub_streams_2.create_on_same_gpus(active_streams);
    sub_streams_3.create_on_same_gpus(active_streams);
    sub_streams_4.create_on_same_gpus(active_streams);
  }

  void create_indexes_for_overflow_sub(CudaStreams streams, uint32_t num_blocks,
                                       uint32_t group_size, bool use_seq,
                                       bool allocate_gpu_memory,
                                       uint64_t &size_tracker) {
    max_indexes_to_erase = num_blocks;

    first_indexes_for_overflow_sub =
        (Torus **)malloc(num_blocks * sizeof(Torus *));
    second_indexes_for_overflow_sub =
        (Torus **)malloc(num_blocks * sizeof(Torus *));
    scalars_for_overflow_sub = (Torus **)malloc(num_blocks * sizeof(Torus *));

    Torus *h_lut_indexes = (Torus *)malloc(num_blocks * sizeof(Torus));
    Torus *h_scalar = (Torus *)malloc(num_blocks * sizeof(Torus));

    // Extra indexes for the luts in first step
    for (int nb = 1; nb <= num_blocks; nb++) {
      first_indexes_for_overflow_sub[nb - 1] =
          (Torus *)cuda_malloc_with_size_tracking_async(
              nb * sizeof(Torus), streams.stream(0), streams.gpu_index(0),
              size_tracker, allocate_gpu_memory);
      for (int index = 0; index < nb; index++) {
        uint32_t grouping_index = index / group_size;
        bool is_in_first_grouping = (grouping_index == 0);
        uint32_t index_in_grouping = index % group_size;
        bool is_last_index = (index == (nb - 1));
        if (is_last_index) {
          if (nb == 1) {
            h_lut_indexes[index] = 2 * group_size;
          } else {
            h_lut_indexes[index] = 2;
          }
        } else if (is_in_first_grouping) {
          h_lut_indexes[index] = index_in_grouping;
        } else {
          h_lut_indexes[index] = index_in_grouping + group_size;
        }
      }
      cuda_memcpy_with_size_tracking_async_to_gpu(
          first_indexes_for_overflow_sub[nb - 1], h_lut_indexes,
          nb * sizeof(Torus), streams.stream(0), streams.gpu_index(0),
          allocate_gpu_memory);
    }
    // Extra indexes for the luts in second step
    for (int nb = 1; nb <= num_blocks; nb++) {
      second_indexes_for_overflow_sub[nb - 1] =
          (Torus *)cuda_malloc_with_size_tracking_async(
              nb * sizeof(Torus), streams.stream(0), streams.gpu_index(0),
              size_tracker, allocate_gpu_memory);
      scalars_for_overflow_sub[nb - 1] =
          (Torus *)cuda_malloc_with_size_tracking_async(
              nb * sizeof(Torus), streams.stream(0), streams.gpu_index(0),
              size_tracker, allocate_gpu_memory);

      for (int index = 0; index < nb; index++) {
        uint32_t grouping_index = index / group_size;
        bool is_in_first_grouping = (grouping_index == 0);
        uint32_t index_in_grouping = index % group_size;

        if (is_in_first_grouping) {
          h_lut_indexes[index] = index_in_grouping;
        } else if (index_in_grouping == (group_size - 1)) {
          if (use_seq) {
            int inner_index = (grouping_index - 1) % (group_size - 1);
            h_lut_indexes[index] = inner_index + 2 * group_size;
          } else {
            h_lut_indexes[index] = 2 * group_size;
          }
        } else {
          h_lut_indexes[index] = index_in_grouping + group_size;
        }

        bool may_have_its_padding_bit_set =
            !is_in_first_grouping && (index_in_grouping == group_size - 1);

        if (may_have_its_padding_bit_set) {
          if (use_seq) {
            h_scalar[index] = 1 << ((grouping_index - 1) % (group_size - 1));
          } else {
            h_scalar[index] = 1;
          }
        } else {
          h_scalar[index] = 0;
        }
      }
      cuda_memcpy_with_size_tracking_async_to_gpu(
          second_indexes_for_overflow_sub[nb - 1], h_lut_indexes,
          nb * sizeof(Torus), streams.stream(0), streams.gpu_index(0),
          allocate_gpu_memory);
      cuda_memcpy_with_size_tracking_async_to_gpu(
          scalars_for_overflow_sub[nb - 1], h_scalar, nb * sizeof(Torus),
          streams.stream(0), streams.gpu_index(0), allocate_gpu_memory);
    }
    free(h_lut_indexes);
    free(h_scalar);
  };

  void release(CudaStreams streams) {

    if (params.message_modulus == 4 && params.carry_modulus == 4 &&
        streams.count() >= 4) {
      div_rem_2_2_mem->release(streams);
      delete div_rem_2_2_mem;
      return;
    }

    uint32_t num_bits_in_message = 31 - __builtin_clz(params.message_modulus);

    // release and delete other operation memory objects
    shift_mem_1->release(streams);
    shift_mem_2->release(streams);
    overflow_sub_mem->release(streams);
    comparison_buffer->release(streams);
    delete shift_mem_1;
    delete shift_mem_2;
    delete overflow_sub_mem;
    delete comparison_buffer;

    // drop temporary buffers
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   remainder1, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   remainder2, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   numerator_block_stack, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   numerator_block_1, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_radix, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   interesting_remainder1,
                                   gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   interesting_remainder2,
                                   gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   interesting_divisor, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   divisor_ms_blocks, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   new_remainder, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   subtraction_overflowed,
                                   gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   did_not_overflow, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   overflow_sum, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   overflow_sum_radix, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_1, gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   at_least_one_upper_block_is_non_zero,
                                   gpu_memory_allocated);
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   cleaned_merged_interesting_remainder,
                                   gpu_memory_allocated);
    // release and delete lookup tables

    // masking_luts_1 and masking_luts_2
    for (int i = 0; i < params.message_modulus - 1; i++) {
      masking_luts_1[i]->release(streams);
      masking_luts_2[i]->release(streams);

      delete masking_luts_1[i];
      delete masking_luts_2[i];
    }
    delete[] masking_luts_1;
    delete[] masking_luts_2;

    // message_extract_lut_1 and message_extract_lut_2
    message_extract_lut_1->release(streams);
    message_extract_lut_2->release(streams);

    delete message_extract_lut_1;
    delete message_extract_lut_2;

    // zero_out_if_overflow_did_not_happen
    zero_out_if_overflow_did_not_happen[0]->release(streams);
    zero_out_if_overflow_did_not_happen[1]->release(streams);

    delete zero_out_if_overflow_did_not_happen[0];
    delete zero_out_if_overflow_did_not_happen[1];

    delete[] zero_out_if_overflow_did_not_happen;

    // zero_out_if_overflow_happened
    zero_out_if_overflow_happened[0]->release(streams);
    zero_out_if_overflow_happened[1]->release(streams);

    delete zero_out_if_overflow_happened[0];
    delete zero_out_if_overflow_happened[1];

    delete[] zero_out_if_overflow_happened;

    // merge_overflow_flags_luts
    for (int i = 0; i < num_bits_in_message; i++) {
      merge_overflow_flags_luts[i]->release(streams);

      delete merge_overflow_flags_luts[i];
    }
    delete[] merge_overflow_flags_luts;

    // release sub streams
    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
    sub_streams_1.release();
    sub_streams_2.release();
    sub_streams_3.release();
    sub_streams_4.release();

    // Delete temporary buffers
    delete remainder1;
    delete remainder2;
    delete numerator_block_stack;
    delete numerator_block_1;
    delete tmp_radix;
    delete interesting_remainder1;
    delete interesting_remainder2;
    delete interesting_divisor;
    delete divisor_ms_blocks;
    delete new_remainder;
    delete subtraction_overflowed;
    delete did_not_overflow;
    delete overflow_sum;
    delete overflow_sum_radix;
    delete tmp_1;
    delete at_least_one_upper_block_is_non_zero;
    delete cleaned_merged_interesting_remainder;

    for (int i = 0; i < max_indexes_to_erase; i++) {
      cuda_drop_with_size_tracking_async(
          first_indexes_for_overflow_sub[i], streams.stream(0),
          streams.gpu_index(0), gpu_memory_allocated);
      cuda_drop_with_size_tracking_async(
          second_indexes_for_overflow_sub[i], streams.stream(0),
          streams.gpu_index(0), gpu_memory_allocated);
      cuda_drop_with_size_tracking_async(
          scalars_for_overflow_sub[i], streams.stream(0), streams.gpu_index(0),
          gpu_memory_allocated);
    }
    free(first_indexes_for_overflow_sub);
    free(second_indexes_for_overflow_sub);
    free(scalars_for_overflow_sub);
  }
};

template <typename Torus> struct int_bitop_buffer {

  int_radix_params params;
  int_radix_lut<Torus> *lut;
  BITOP_TYPE op;
  bool gpu_memory_allocated;

  int_bitop_buffer(CudaStreams streams, BITOP_TYPE op, int_radix_params params,
                   uint32_t num_radix_blocks, bool allocate_gpu_memory,
                   uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;
    this->op = op;
    this->params = params;
    auto active_streams = streams.active_gpu_subset(num_radix_blocks);
    switch (op) {
    case BITAND:
    case BITOR:
    case BITXOR:
      lut = new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                     allocate_gpu_memory, size_tracker);
      {
        auto lut_bivariate_f = [op](Torus lhs, Torus rhs) -> Torus {
          if (op == BITOP_TYPE::BITAND) {
            // AND
            return lhs & rhs;
          } else if (op == BITOP_TYPE::BITOR) {
            // OR
            return lhs | rhs;
          } else {
            // XOR
            return lhs ^ rhs;
          }
        };

        generate_device_accumulator_bivariate<Torus>(
            streams.stream(0), streams.gpu_index(0), lut->get_lut(0, 0),
            lut->get_degree(0), lut->get_max_degree(0), params.glwe_dimension,
            params.polynomial_size, params.message_modulus,
            params.carry_modulus, lut_bivariate_f, gpu_memory_allocated);
        lut->broadcast_lut(active_streams);
      }
      break;
    default:
      // Scalar OP
      lut = new int_radix_lut<Torus>(streams, params, params.message_modulus,
                                     num_radix_blocks, allocate_gpu_memory,
                                     size_tracker);

      for (int i = 0; i < params.message_modulus; i++) {
        auto rhs = i;

        auto lut_univariate_scalar_f = [op, rhs](Torus x) -> Torus {
          if (op == BITOP_TYPE::SCALAR_BITAND) {
            // AND
            return x & rhs;
          } else if (op == BITOP_TYPE::SCALAR_BITOR) {
            // OR
            return x | rhs;
          } else {
            // XOR
            return x ^ rhs;
          }
        };
        generate_device_accumulator<Torus>(
            streams.stream(0), streams.gpu_index(0), lut->get_lut(0, i),
            lut->get_degree(i), lut->get_max_degree(i), params.glwe_dimension,
            params.polynomial_size, params.message_modulus,
            params.carry_modulus, lut_univariate_scalar_f,
            gpu_memory_allocated);
        lut->broadcast_lut(active_streams);
      }
    }
  }

  void release(CudaStreams streams) {
    lut->release(streams);
    delete lut;
  }
};

template <typename Torus> struct int_scalar_mul_buffer {
  int_radix_params params;
  int_logical_scalar_shift_buffer<Torus> *logical_scalar_shift_buffer;
  int_sum_ciphertexts_vec_memory<Torus> *sum_ciphertexts_vec_mem;
  CudaRadixCiphertextFFI *preshifted_buffer;
  CudaRadixCiphertextFFI *all_shifted_buffer;
  int_sc_prop_memory<Torus> *sc_prop_mem;
  bool anticipated_buffers_drop;
  bool gpu_memory_allocated;
  uint32_t num_ciphertext_bits;

  int_scalar_mul_buffer(CudaStreams streams, int_radix_params params,
                        uint32_t num_radix_blocks, uint32_t num_scalar_bits,
                        bool allocate_gpu_memory, bool anticipated_buffer_drop,
                        uint64_t &size_tracker) {
    gpu_memory_allocated = allocate_gpu_memory;
    this->params = params;
    this->anticipated_buffers_drop = anticipated_buffer_drop;

    uint32_t msg_bits = (uint32_t)std::log2(params.message_modulus);
    num_ciphertext_bits = msg_bits * num_scalar_bits;

    //// Contains all shifted values of lhs for shift in range (0..msg_bits)
    //// The idea is that with these we can create all other shift that are
    /// in / range (0..total_bits) for free (block rotation)
    preshifted_buffer = new CudaRadixCiphertextFFI;
    uint64_t anticipated_drop_mem = 0;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), preshifted_buffer,
        msg_bits * num_radix_blocks, params.big_lwe_dimension,
        anticipated_drop_mem, allocate_gpu_memory);

    all_shifted_buffer = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), all_shifted_buffer,
        num_ciphertext_bits * num_radix_blocks, params.big_lwe_dimension,
        size_tracker, allocate_gpu_memory);

    if (num_ciphertext_bits * num_radix_blocks >= num_radix_blocks + 2)
      logical_scalar_shift_buffer = new int_logical_scalar_shift_buffer<Torus>(
          streams, LEFT_SHIFT, params, num_radix_blocks, allocate_gpu_memory,
          all_shifted_buffer, anticipated_drop_mem);
    else
      logical_scalar_shift_buffer = new int_logical_scalar_shift_buffer<Torus>(
          streams, LEFT_SHIFT, params, num_radix_blocks, allocate_gpu_memory,
          anticipated_drop_mem);

    uint64_t last_step_mem = 0;
    if (num_ciphertext_bits > 0) {
      sum_ciphertexts_vec_mem = new int_sum_ciphertexts_vec_memory<Torus>(
          streams, params, num_radix_blocks, num_ciphertext_bits, true,
          allocate_gpu_memory, last_step_mem);
    }
    uint32_t requested_flag = outputFlag::FLAG_NONE;
    sc_prop_mem = new int_sc_prop_memory<Torus>(
        streams, params, num_radix_blocks, requested_flag, allocate_gpu_memory,
        last_step_mem);
    if (anticipated_buffer_drop) {
      size_tracker += std::max(anticipated_drop_mem, last_step_mem);
    } else {
      size_tracker += anticipated_drop_mem + last_step_mem;
    }
  }

  void release(CudaStreams streams) {
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   all_shifted_buffer, gpu_memory_allocated);
    if (num_ciphertext_bits > 0) {
      sum_ciphertexts_vec_mem->release(streams);
      delete sum_ciphertexts_vec_mem;
    }
    sc_prop_mem->release(streams);
    delete sc_prop_mem;
    delete all_shifted_buffer;
    release_buffers(streams);
  }

  void release_buffers(CudaStreams streams) {
    if (preshifted_buffer) {
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     preshifted_buffer, gpu_memory_allocated);
      delete preshifted_buffer;
      preshifted_buffer = nullptr;
    }

    if (logical_scalar_shift_buffer) {
      logical_scalar_shift_buffer->release(streams);
      delete logical_scalar_shift_buffer;
      logical_scalar_shift_buffer = nullptr;
    }
  }
};

template <typename Torus> struct int_abs_buffer {
  int_radix_params params;

  int_arithmetic_scalar_shift_buffer<Torus> *arithmetic_scalar_shift_mem;
  int_sc_prop_memory<Torus> *scp_mem;
  int_bitop_buffer<Torus> *bitxor_mem;

  CudaRadixCiphertextFFI *mask;
  bool allocate_gpu_memory;

  int_abs_buffer(CudaStreams streams, int_radix_params params,
                 uint32_t num_radix_blocks, bool allocate_gpu_memory,
                 uint64_t &size_tracker) {
    this->params = params;
    this->allocate_gpu_memory = allocate_gpu_memory;
    arithmetic_scalar_shift_mem = new int_arithmetic_scalar_shift_buffer<Torus>(
        streams, SHIFT_OR_ROTATE_TYPE::RIGHT_SHIFT, params, num_radix_blocks,
        allocate_gpu_memory, size_tracker);
    uint32_t requested_flag = outputFlag::FLAG_NONE;
    scp_mem = new int_sc_prop_memory<Torus>(streams, params, num_radix_blocks,
                                            requested_flag, allocate_gpu_memory,
                                            size_tracker);
    bitxor_mem = new int_bitop_buffer<Torus>(streams, BITOP_TYPE::BITXOR,
                                             params, num_radix_blocks,
                                             allocate_gpu_memory, size_tracker);

    mask = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), mask, num_radix_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
  }

  void release(CudaStreams streams) {
    arithmetic_scalar_shift_mem->release(streams);
    scp_mem->release(streams);
    bitxor_mem->release(streams);

    delete arithmetic_scalar_shift_mem;
    delete scp_mem;
    delete bitxor_mem;

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   mask, this->allocate_gpu_memory);
    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
    delete mask;
  }
};

template <typename Torus> struct int_div_rem_memory {
  int_radix_params params;
  CudaStreams active_streams;
  bool is_signed;
  // memory objects for other operations
  unsigned_int_div_rem_memory<Torus> *unsigned_mem;
  int_abs_buffer<Torus> *abs_mem_1;
  int_abs_buffer<Torus> *abs_mem_2;
  int_sc_prop_memory<Torus> *scp_mem_1;
  int_sc_prop_memory<Torus> *scp_mem_2;
  int_cmux_buffer<Torus> *cmux_quotient_mem;
  int_cmux_buffer<Torus> *cmux_remainder_mem;

  // lookup tables
  int_radix_lut<Torus> *compare_signed_bits_lut;

  // sub streams
  CudaStreams sub_streams_1;
  CudaStreams sub_streams_2;

  // temporary device buffers
  CudaRadixCiphertextFFI *positive_numerator;
  CudaRadixCiphertextFFI *positive_divisor;
  CudaRadixCiphertextFFI *sign_bits_are_different;
  CudaRadixCiphertextFFI *negated_quotient;
  CudaRadixCiphertextFFI *negated_remainder;
  bool gpu_memory_allocated;

  int_div_rem_memory(CudaStreams streams, int_radix_params params,
                     bool is_signed, uint32_t num_blocks,
                     bool allocate_gpu_memory, uint64_t &size_tracker) {

    gpu_memory_allocated = allocate_gpu_memory;
    this->active_streams = streams.active_gpu_subset(num_blocks);
    this->params = params;
    this->is_signed = is_signed;

    unsigned_mem = new unsigned_int_div_rem_memory<Torus>(
        streams, params, num_blocks, allocate_gpu_memory, size_tracker);

    if (is_signed) {
      Torus sign_bit_pos = 31 - __builtin_clz(params.message_modulus) - 1;

      // init memory objects for other integer operations
      abs_mem_1 = new int_abs_buffer<Torus>(streams, params, num_blocks,
                                            allocate_gpu_memory, size_tracker);
      abs_mem_2 = new int_abs_buffer<Torus>(streams, params, num_blocks,
                                            allocate_gpu_memory, size_tracker);
      uint32_t requested_flag = outputFlag::FLAG_NONE;
      scp_mem_1 = new int_sc_prop_memory<Torus>(
          streams, params, num_blocks, requested_flag, allocate_gpu_memory,
          size_tracker);
      scp_mem_2 = new int_sc_prop_memory<Torus>(
          streams, params, num_blocks, requested_flag, allocate_gpu_memory,
          size_tracker);

      std::function<uint64_t(uint64_t)> quotient_predicate_lut_f =
          [](uint64_t x) -> uint64_t { return x == 1; };
      std::function<uint64_t(uint64_t)> remainder_predicate_lut_f =
          [sign_bit_pos](uint64_t x) -> uint64_t {
        return (x >> sign_bit_pos) == 1;
      };

      cmux_quotient_mem = new int_cmux_buffer<Torus>(
          streams, quotient_predicate_lut_f, params, num_blocks,
          allocate_gpu_memory, size_tracker);
      cmux_remainder_mem = new int_cmux_buffer<Torus>(
          streams, remainder_predicate_lut_f, params, num_blocks,
          allocate_gpu_memory, size_tracker);
      // init temporary memory buffers
      positive_numerator = new CudaRadixCiphertextFFI;
      create_zero_radix_ciphertext_async<Torus>(
          streams.stream(0), streams.gpu_index(0), positive_numerator,
          num_blocks, params.big_lwe_dimension, size_tracker,
          allocate_gpu_memory);
      positive_divisor = new CudaRadixCiphertextFFI;
      create_zero_radix_ciphertext_async<Torus>(
          streams.stream(0), streams.gpu_index(0), positive_divisor, num_blocks,
          params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
      negated_quotient = new CudaRadixCiphertextFFI;
      create_zero_radix_ciphertext_async<Torus>(
          streams.stream(0), streams.gpu_index(0), negated_quotient, num_blocks,
          params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
      negated_remainder = new CudaRadixCiphertextFFI;
      create_zero_radix_ciphertext_async<Torus>(
          streams.stream(0), streams.gpu_index(0), negated_remainder,
          num_blocks, params.big_lwe_dimension, size_tracker,
          allocate_gpu_memory);

      // init boolean temporary buffers
      sign_bits_are_different = new CudaRadixCiphertextFFI;
      create_zero_radix_ciphertext_async<Torus>(
          streams.stream(0), streams.gpu_index(0), sign_bits_are_different, 1,
          params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

      // init sub streams
      sub_streams_1.create_on_same_gpus(streams);
      sub_streams_2.create_on_same_gpus(streams);

      // init lookup tables
      //  to extract and compare signed bits
      auto f_compare_extracted_signed_bits = [sign_bit_pos](Torus x,
                                                            Torus y) -> Torus {
        Torus x_sign_bit = (x >> sign_bit_pos) & 1;
        Torus y_sign_bit = (y >> sign_bit_pos) & 1;
        return (Torus)(x_sign_bit != y_sign_bit);
      };

      compare_signed_bits_lut = new int_radix_lut<Torus>(
          streams, params, 1, 1, allocate_gpu_memory, size_tracker);

      generate_device_accumulator_bivariate<Torus>(
          streams.stream(0), streams.gpu_index(0),
          compare_signed_bits_lut->get_lut(0, 0),
          compare_signed_bits_lut->get_degree(0),
          compare_signed_bits_lut->get_max_degree(0), params.glwe_dimension,
          params.polynomial_size, params.message_modulus, params.carry_modulus,
          f_compare_extracted_signed_bits, gpu_memory_allocated);
      auto active_gpu_count_cmp =
          streams.active_gpu_subset(1); // only 1 block needed
      compare_signed_bits_lut->broadcast_lut(active_gpu_count_cmp);
    }
  }

  void release(CudaStreams streams) {
    unsigned_mem->release(streams);
    delete unsigned_mem;

    if (is_signed) {
      // release objects for other integer operations
      abs_mem_1->release(streams);
      abs_mem_2->release(streams);
      scp_mem_1->release(streams);
      scp_mem_2->release(streams);
      cmux_quotient_mem->release(streams);
      cmux_remainder_mem->release(streams);

      delete abs_mem_1;
      delete abs_mem_2;
      delete scp_mem_1;
      delete scp_mem_2;
      delete cmux_quotient_mem;
      delete cmux_remainder_mem;

      // drop temporary buffers
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     positive_numerator, gpu_memory_allocated);
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     positive_divisor, gpu_memory_allocated);
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     sign_bits_are_different,
                                     gpu_memory_allocated);
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     negated_quotient, gpu_memory_allocated);
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     negated_remainder, gpu_memory_allocated);
      // release lookup tables
      compare_signed_bits_lut->release(streams);
      delete compare_signed_bits_lut;

      // release sub streams
      cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
      sub_streams_1.release();
      sub_streams_2.release();

      // delete temporary buffers
      delete positive_numerator;
      delete positive_divisor;
      delete sign_bits_are_different;
      delete negated_quotient;
      delete negated_remainder;
    }
  }
};

template <typename Torus> struct int_scalar_mul_high_buffer {
  int_radix_params params;
  bool allocate_gpu_memory;

  int_logical_scalar_shift_buffer<Torus> *logical_scalar_shift_mem;
  int_scalar_mul_buffer<Torus> *scalar_mul_mem;

  CudaRadixCiphertextFFI *tmp;

  int_scalar_mul_high_buffer(CudaStreams streams, const int_radix_params params,
                             uint32_t num_radix_blocks,
                             uint32_t num_scalar_bits,
                             const bool allocate_gpu_memory,
                             uint64_t &size_tracker) {

    this->params = params;
    this->allocate_gpu_memory = allocate_gpu_memory;

    this->logical_scalar_shift_mem = new int_logical_scalar_shift_buffer<Torus>(
        streams, RIGHT_SHIFT, params, 2 * num_radix_blocks, allocate_gpu_memory,
        size_tracker);

    this->scalar_mul_mem = new int_scalar_mul_buffer<Torus>(
        streams, params, 2 * num_radix_blocks, num_scalar_bits,
        allocate_gpu_memory, true, size_tracker);

    this->tmp = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp, 2 * num_radix_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
  }

  void release(CudaStreams streams) {

    logical_scalar_shift_mem->release(streams);
    delete logical_scalar_shift_mem;

    scalar_mul_mem->release(streams);
    delete scalar_mul_mem;
    scalar_mul_mem = nullptr;

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0), tmp,
                                   allocate_gpu_memory);
    delete tmp;
  }
};

template <typename Torus> struct int_sub_and_propagate {
  int_radix_params params;
  bool allocate_gpu_memory;

  CudaRadixCiphertextFFI *neg_rhs_array;

  int_sc_prop_memory<Torus> *sc_prop_mem;

  int_sub_and_propagate(CudaStreams streams, const int_radix_params params,
                        uint32_t num_radix_blocks, uint32_t requested_flag_in,
                        bool allocate_gpu_memory, uint64_t &size_tracker) {

    this->params = params;
    this->allocate_gpu_memory = allocate_gpu_memory;

    this->sc_prop_mem = new int_sc_prop_memory<Torus>(
        streams, params, num_radix_blocks, requested_flag_in,
        allocate_gpu_memory, size_tracker);

    this->neg_rhs_array = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), neg_rhs_array,
        num_radix_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);
  }

  void release(CudaStreams streams) {

    sc_prop_mem->release(streams);
    delete sc_prop_mem;

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   neg_rhs_array, allocate_gpu_memory);
    delete neg_rhs_array;
  }
};

template <typename Torus> struct int_extend_radix_with_sign_msb_buffer {

  int_radix_params params;
  bool allocate_gpu_memory;

  int_radix_lut<Torus> *lut = nullptr;

  CudaRadixCiphertextFFI *last_block = nullptr;
  CudaRadixCiphertextFFI *padding_block = nullptr;

  int_extend_radix_with_sign_msb_buffer(CudaStreams streams,
                                        const int_radix_params params,
                                        uint32_t num_radix_blocks,
                                        uint32_t num_additional_blocks,
                                        const bool allocate_gpu_memory,
                                        uint64_t &size_tracker) {

    this->params = params;
    this->allocate_gpu_memory = allocate_gpu_memory;

    if (num_additional_blocks != 0) {
      this->lut = new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                           allocate_gpu_memory, size_tracker);

      uint32_t bits_per_block = std::log2(params.message_modulus);
      uint32_t msg_modulus = params.message_modulus;

      generate_device_accumulator<Torus>(
          streams.stream(0), streams.gpu_index(0), lut->get_lut(0, 0),
          lut->get_degree(0), lut->get_max_degree(0), params.glwe_dimension,
          params.polynomial_size, params.message_modulus, params.carry_modulus,
          [msg_modulus, bits_per_block](Torus x) {
            const auto xm = x % msg_modulus;
            const auto sign_bit = (xm >> (bits_per_block - 1)) & 1;
            return (Torus)((msg_modulus - 1) * sign_bit);
          },
          allocate_gpu_memory);

      auto active_streams = streams.active_gpu_subset(num_radix_blocks);
      lut->broadcast_lut(active_streams);

      this->last_block = new CudaRadixCiphertextFFI;

      create_zero_radix_ciphertext_async<Torus>(
          streams.stream(0), streams.gpu_index(0), last_block, 1,
          params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

      this->padding_block = new CudaRadixCiphertextFFI;

      create_zero_radix_ciphertext_async<Torus>(
          streams.stream(0), streams.gpu_index(0), padding_block, 1,
          params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
    }
  }

  void release(CudaStreams streams) {

    if (lut != nullptr) {
      lut->release(streams);
      delete lut;
    }
    if (last_block != nullptr) {
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     last_block, allocate_gpu_memory);
      delete last_block;
    }
    if (padding_block != nullptr) {
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     padding_block, allocate_gpu_memory);
      delete padding_block;
    }
  }
};

template <typename Torus> struct int_unsigned_scalar_div_mem {
  int_radix_params params;
  bool allocate_gpu_memory;

  CudaRadixCiphertextFFI *tmp_ffi = nullptr;

  int_logical_scalar_shift_buffer<Torus> *logical_scalar_shift_mem = nullptr;
  int_scalar_mul_high_buffer<Torus> *scalar_mul_high_mem = nullptr;
  int_sc_prop_memory<Torus> *scp_mem = nullptr;
  int_sub_and_propagate<Torus> *sub_and_propagate_mem = nullptr;

  int_unsigned_scalar_div_mem(CudaStreams streams,
                              const int_radix_params params,
                              uint32_t num_radix_blocks,
                              const CudaScalarDivisorFFI *scalar_divisor_ffi,
                              const bool allocate_gpu_memory,
                              uint64_t &size_tracker) {

    this->params = params;
    this->allocate_gpu_memory = allocate_gpu_memory;

    if (!scalar_divisor_ffi->is_abs_divisor_one) {
      if (scalar_divisor_ffi->is_divisor_pow2) {

        logical_scalar_shift_mem = new int_logical_scalar_shift_buffer<Torus>(
            streams, RIGHT_SHIFT, params, num_radix_blocks, allocate_gpu_memory,
            size_tracker);

      } else if (scalar_divisor_ffi->divisor_has_more_bits_than_numerator) {

        tmp_ffi = new CudaRadixCiphertextFFI;
        create_zero_radix_ciphertext_async<Torus>(
            streams.stream(0), streams.gpu_index(0), tmp_ffi, num_radix_blocks,
            params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

      } else if (scalar_divisor_ffi
                     ->is_chosen_multiplier_geq_two_pow_numerator) {

        logical_scalar_shift_mem = new int_logical_scalar_shift_buffer<Torus>(
            streams, RIGHT_SHIFT, params, num_radix_blocks, allocate_gpu_memory,
            size_tracker);
        scalar_mul_high_mem = new int_scalar_mul_high_buffer<Torus>(
            streams, params, num_radix_blocks, scalar_divisor_ffi->active_bits,
            allocate_gpu_memory, size_tracker);
        scp_mem = new int_sc_prop_memory<Torus>(
            streams, params, num_radix_blocks, FLAG_NONE, allocate_gpu_memory,
            size_tracker);
        sub_and_propagate_mem = new int_sub_and_propagate<Torus>(
            streams, params, num_radix_blocks, FLAG_NONE, allocate_gpu_memory,
            size_tracker);
        tmp_ffi = new CudaRadixCiphertextFFI;
        create_zero_radix_ciphertext_async<Torus>(
            streams.stream(0), streams.gpu_index(0), tmp_ffi, num_radix_blocks,
            params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

      } else {

        logical_scalar_shift_mem = new int_logical_scalar_shift_buffer<Torus>(
            streams, RIGHT_SHIFT, params, num_radix_blocks, allocate_gpu_memory,
            size_tracker);
        scalar_mul_high_mem = new int_scalar_mul_high_buffer<Torus>(
            streams, params, num_radix_blocks, scalar_divisor_ffi->active_bits,
            allocate_gpu_memory, size_tracker);
      }
    }
  }

  void release(CudaStreams streams) {

    if (logical_scalar_shift_mem != nullptr) {
      logical_scalar_shift_mem->release(streams);
      delete logical_scalar_shift_mem;
    }
    if (scalar_mul_high_mem != nullptr) {
      scalar_mul_high_mem->release(streams);
      delete scalar_mul_high_mem;
    }
    if (scp_mem != nullptr) {
      scp_mem->release(streams);
      delete scp_mem;
    }
    if (sub_and_propagate_mem != nullptr) {
      sub_and_propagate_mem->release(streams);
      delete sub_and_propagate_mem;
    }
    if (tmp_ffi != nullptr) {
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     tmp_ffi, allocate_gpu_memory);
      delete tmp_ffi;
    }
  }
};

template <typename Torus> struct int_signed_scalar_mul_high_buffer {
  int_radix_params params;
  bool allocate_gpu_memory;

  int_logical_scalar_shift_buffer<Torus> *logical_scalar_shift_mem;
  int_scalar_mul_buffer<Torus> *scalar_mul_mem;
  int_extend_radix_with_sign_msb_buffer<Torus> *extend_radix_mem;

  CudaRadixCiphertextFFI *tmp;

  int_signed_scalar_mul_high_buffer(CudaStreams streams,
                                    const int_radix_params params,
                                    uint32_t num_radix_blocks,
                                    uint32_t num_scalar_bits,
                                    const bool allocate_gpu_memory,
                                    uint64_t &size_tracker) {

    this->params = params;
    this->allocate_gpu_memory = allocate_gpu_memory;

    this->logical_scalar_shift_mem = new int_logical_scalar_shift_buffer<Torus>(
        streams, RIGHT_SHIFT, params, 2 * num_radix_blocks, allocate_gpu_memory,
        size_tracker);

    this->scalar_mul_mem = new int_scalar_mul_buffer<Torus>(
        streams, params, 2 * num_radix_blocks, num_scalar_bits,
        allocate_gpu_memory, true, size_tracker);

    this->tmp = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp, 2 * num_radix_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

    this->extend_radix_mem = new int_extend_radix_with_sign_msb_buffer<Torus>(
        streams, params, num_radix_blocks, num_radix_blocks,
        allocate_gpu_memory, size_tracker);
  }

  void release(CudaStreams streams) {

    logical_scalar_shift_mem->release(streams);
    delete logical_scalar_shift_mem;

    scalar_mul_mem->release(streams);
    delete scalar_mul_mem;

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0), tmp,
                                   allocate_gpu_memory);
    delete tmp;

    extend_radix_mem->release(streams);
    delete extend_radix_mem;
  }
};

template <typename Torus> struct int_signed_scalar_div_mem {
  int_radix_params params;
  bool allocate_gpu_memory;

  CudaRadixCiphertextFFI *tmp_ffi = nullptr;
  CudaRadixCiphertextFFI *xsign_ffi = nullptr;

  int_arithmetic_scalar_shift_buffer<Torus> *arithmetic_scalar_shift_mem =
      nullptr;
  int_logical_scalar_shift_buffer<Torus> *logical_scalar_shift_mem = nullptr;
  int_signed_scalar_mul_high_buffer<Torus> *scalar_mul_high_mem = nullptr;
  int_sc_prop_memory<Torus> *scp_mem = nullptr;
  int_sub_and_propagate<Torus> *sub_and_propagate_mem = nullptr;

  int_signed_scalar_div_mem(CudaStreams streams, const int_radix_params params,
                            uint32_t num_radix_blocks,
                            const CudaScalarDivisorFFI *scalar_divisor_ffi,
                            const bool allocate_gpu_memory,
                            uint64_t &size_tracker) {

    this->params = params;
    this->allocate_gpu_memory = allocate_gpu_memory;

    if (!scalar_divisor_ffi->chosen_multiplier_has_more_bits_than_numerator) {

      if (scalar_divisor_ffi->is_abs_divisor_one &&
          scalar_divisor_ffi->is_divisor_negative) {
        tmp_ffi = new CudaRadixCiphertextFFI;

        create_zero_radix_ciphertext_async<Torus>(
            streams.stream(0), streams.gpu_index(0), tmp_ffi, num_radix_blocks,
            params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

      } else if (!scalar_divisor_ffi->is_abs_divisor_one) {

        tmp_ffi = new CudaRadixCiphertextFFI;
        create_zero_radix_ciphertext_async<Torus>(
            streams.stream(0), streams.gpu_index(0), tmp_ffi, num_radix_blocks,
            params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

        arithmetic_scalar_shift_mem =
            new int_arithmetic_scalar_shift_buffer<Torus>(
                streams, RIGHT_SHIFT, params, num_radix_blocks,
                allocate_gpu_memory, size_tracker);

        if (scalar_divisor_ffi->is_divisor_pow2) {

          logical_scalar_shift_mem = new int_logical_scalar_shift_buffer<Torus>(
              streams, RIGHT_SHIFT, params, num_radix_blocks,
              allocate_gpu_memory, size_tracker);
          scp_mem = new int_sc_prop_memory<Torus>(
              streams, params, num_radix_blocks, FLAG_NONE, allocate_gpu_memory,
              size_tracker);

        } else {

          xsign_ffi = new CudaRadixCiphertextFFI;
          create_zero_radix_ciphertext_async<Torus>(
              streams.stream(0), streams.gpu_index(0), xsign_ffi,
              num_radix_blocks, params.big_lwe_dimension, size_tracker,
              allocate_gpu_memory);

          scalar_mul_high_mem = new int_signed_scalar_mul_high_buffer<Torus>(
              streams, params, num_radix_blocks,
              scalar_divisor_ffi->active_bits, allocate_gpu_memory,
              size_tracker);

          sub_and_propagate_mem = new int_sub_and_propagate<Torus>(
              streams, params, num_radix_blocks, FLAG_NONE, allocate_gpu_memory,
              size_tracker);

          if (scalar_divisor_ffi->is_chosen_multiplier_geq_two_pow_numerator) {
            scp_mem = new int_sc_prop_memory<Torus>(
                streams, params, num_radix_blocks, FLAG_NONE,
                allocate_gpu_memory, size_tracker);
          }
        }
      }
    }
  }

  void release(CudaStreams streams) {

    if (arithmetic_scalar_shift_mem != nullptr) {
      arithmetic_scalar_shift_mem->release(streams);
      delete arithmetic_scalar_shift_mem;
    }
    if (logical_scalar_shift_mem != nullptr) {
      logical_scalar_shift_mem->release(streams);
      delete logical_scalar_shift_mem;
    }
    if (scalar_mul_high_mem != nullptr) {
      scalar_mul_high_mem->release(streams);
      delete scalar_mul_high_mem;
    }
    if (scp_mem != nullptr) {
      scp_mem->release(streams);
      delete scp_mem;
    }
    if (sub_and_propagate_mem != nullptr) {
      sub_and_propagate_mem->release(streams);
      delete sub_and_propagate_mem;
    }
    if (tmp_ffi != nullptr) {
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     tmp_ffi, allocate_gpu_memory);
      delete tmp_ffi;
    }
    if (xsign_ffi != nullptr) {
      release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                     xsign_ffi, allocate_gpu_memory);
      delete xsign_ffi;
    }
  }
};

template <typename Torus> struct int_unsigned_scalar_div_rem_buffer {
  int_radix_params params;
  bool allocate_gpu_memory;

  CudaRadixCiphertextFFI *numerator_ct;

  int_unsigned_scalar_div_mem<Torus> *unsigned_div_mem;
  int_bitop_buffer<Torus> *bitop_mem = nullptr;
  int_scalar_mul_buffer<Torus> *scalar_mul_mem = nullptr;
  int_sub_and_propagate<Torus> *sub_and_propagate_mem = nullptr;

  int_unsigned_scalar_div_rem_buffer(
      CudaStreams streams, const int_radix_params params,
      uint32_t num_radix_blocks, const CudaScalarDivisorFFI *scalar_divisor_ffi,
      uint32_t const active_bits_divisor, const bool allocate_gpu_memory,
      uint64_t &size_tracker) {

    this->params = params;
    this->allocate_gpu_memory = allocate_gpu_memory;

    this->numerator_ct = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), numerator_ct, num_radix_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

    this->unsigned_div_mem = new int_unsigned_scalar_div_mem<Torus>(
        streams, params, num_radix_blocks, scalar_divisor_ffi,
        allocate_gpu_memory, size_tracker);

    if (scalar_divisor_ffi->is_divisor_pow2) {
      this->bitop_mem = new int_bitop_buffer<Torus>(
          streams, SCALAR_BITAND, params, num_radix_blocks, allocate_gpu_memory,
          size_tracker);
    } else {
      if (!scalar_divisor_ffi->is_divisor_zero &&
          !scalar_divisor_ffi->is_abs_divisor_one && num_radix_blocks != 0) {
        this->scalar_mul_mem = new int_scalar_mul_buffer<Torus>(
            streams, params, num_radix_blocks, active_bits_divisor,
            allocate_gpu_memory, true, size_tracker);
      }
      this->sub_and_propagate_mem = new int_sub_and_propagate<Torus>(
          streams, params, num_radix_blocks, FLAG_NONE, allocate_gpu_memory,
          size_tracker);
    }
  }

  void release(CudaStreams streams) {

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   numerator_ct, allocate_gpu_memory);
    delete numerator_ct;

    unsigned_div_mem->release(streams);
    delete unsigned_div_mem;

    if (bitop_mem != nullptr) {
      bitop_mem->release(streams);
      delete bitop_mem;
    }
    if (scalar_mul_mem != nullptr) {
      scalar_mul_mem->release(streams);
      delete scalar_mul_mem;
    }
    if (sub_and_propagate_mem != nullptr) {
      sub_and_propagate_mem->release(streams);
      delete sub_and_propagate_mem;
    }
  }
};

template <typename Torus> struct int_signed_scalar_div_rem_buffer {
  int_radix_params params;
  bool allocate_gpu_memory;

  CudaRadixCiphertextFFI *numerator_ct;

  int_signed_scalar_div_mem<Torus> *signed_div_mem;
  int_logical_scalar_shift_buffer<Torus> *logical_scalar_shift_mem = nullptr;
  int_scalar_mul_buffer<Torus> *scalar_mul_mem = nullptr;
  int_sub_and_propagate<Torus> *sub_and_propagate_mem;
  int_sc_prop_memory<Torus> *scp_mem;

  int_signed_scalar_div_rem_buffer(
      CudaStreams streams, const int_radix_params params,
      uint32_t num_radix_blocks, const CudaScalarDivisorFFI *scalar_divisor_ffi,
      uint32_t const active_bits_divisor, const bool allocate_gpu_memory,
      uint64_t &size_tracker) {

    this->params = params;
    this->allocate_gpu_memory = allocate_gpu_memory;

    this->numerator_ct = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), numerator_ct, num_radix_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

    this->signed_div_mem = new int_signed_scalar_div_mem<Torus>(
        streams, params, num_radix_blocks, scalar_divisor_ffi,
        allocate_gpu_memory, size_tracker);

    this->scp_mem = new int_sc_prop_memory<Torus>(
        streams, params, num_radix_blocks, FLAG_NONE, allocate_gpu_memory,
        size_tracker);

    bool is_divisor_one = scalar_divisor_ffi->is_abs_divisor_one &&
                          !scalar_divisor_ffi->is_divisor_negative;

    if (!scalar_divisor_ffi->is_divisor_negative &&
        scalar_divisor_ffi->is_divisor_pow2) {
      this->logical_scalar_shift_mem =
          new int_logical_scalar_shift_buffer<Torus>(
              streams, LEFT_SHIFT, params, num_radix_blocks,
              allocate_gpu_memory, size_tracker);

    } else if (!scalar_divisor_ffi->is_divisor_zero && !is_divisor_one &&
               num_radix_blocks != 0) {
      this->scalar_mul_mem = new int_scalar_mul_buffer<Torus>(
          streams, params, num_radix_blocks, active_bits_divisor,
          allocate_gpu_memory, true, size_tracker);
    }

    this->sub_and_propagate_mem = new int_sub_and_propagate<Torus>(
        streams, params, num_radix_blocks, FLAG_NONE, allocate_gpu_memory,
        size_tracker);
  }

  void release(CudaStreams streams) {

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   numerator_ct, allocate_gpu_memory);
    delete numerator_ct;
    numerator_ct = nullptr;

    signed_div_mem->release(streams);
    delete signed_div_mem;
    signed_div_mem = nullptr;

    scp_mem->release(streams);
    delete scp_mem;
    scp_mem = nullptr;

    if (logical_scalar_shift_mem != nullptr) {
      logical_scalar_shift_mem->release(streams);
      delete logical_scalar_shift_mem;
      logical_scalar_shift_mem = nullptr;
    }
    if (scalar_mul_mem != nullptr) {
      scalar_mul_mem->release(streams);
      delete scalar_mul_mem;
      scalar_mul_mem = nullptr;
    }
    sub_and_propagate_mem->release(streams);
    delete sub_and_propagate_mem;
    sub_and_propagate_mem = nullptr;
  }
};

template <typename Torus> struct int_prepare_count_of_consecutive_bits_buffer {
  int_radix_params params;
  bool allocate_gpu_memory;

  int_radix_lut<Torus> *univ_lut_mem;
  int_radix_lut<Torus> *biv_lut_mem;

  Direction direction;
  BitValue bit_value;

  CudaRadixCiphertextFFI *tmp_ct;

  int_prepare_count_of_consecutive_bits_buffer(
      CudaStreams streams, const int_radix_params params,
      uint32_t num_radix_blocks, Direction direction, BitValue bit_value,
      const bool allocate_gpu_memory, uint64_t &size_tracker) {
    this->params = params;
    this->allocate_gpu_memory = allocate_gpu_memory;
    this->direction = direction;
    this->bit_value = bit_value;
    auto active_streams = streams.active_gpu_subset(num_radix_blocks);
    this->univ_lut_mem =
        new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                 allocate_gpu_memory, size_tracker);
    this->biv_lut_mem =
        new int_radix_lut<Torus>(streams, params, 1, num_radix_blocks,
                                 allocate_gpu_memory, size_tracker);

    const uint32_t num_bits = std::log2(this->params.message_modulus);

    auto generate_uni_lut_lambda = [this, num_bits](Torus x) -> Torus {
      x %= this->params.message_modulus;
      uint64_t count = 0;

      if (this->direction == Trailing) {
        for (uint32_t i = 0; i < num_bits; ++i) {
          if (((x >> i) & 1) != this->bit_value) {
            break;
          }
          count++;
        }
      } else {
        for (int32_t i = num_bits - 1; i >= 0; --i) {
          if (((x >> i) & 1) != this->bit_value) {
            break;
          }
          count++;
        }
      }
      return count;
    };

    generate_device_accumulator<Torus>(
        streams.stream(0), streams.gpu_index(0), univ_lut_mem->get_lut(0, 0),
        univ_lut_mem->get_degree(0), univ_lut_mem->get_max_degree(0),
        params.glwe_dimension, params.polynomial_size, params.message_modulus,
        params.carry_modulus, generate_uni_lut_lambda, allocate_gpu_memory);

    univ_lut_mem->broadcast_lut(active_streams);

    auto generate_bi_lut_lambda =
        [num_bits](Torus block_num_bit_count,
                   Torus more_significant_block_bit_count) -> Torus {
      if (more_significant_block_bit_count == num_bits) {
        return block_num_bit_count;
      }
      return 0;
    };

    generate_device_accumulator_bivariate<Torus>(
        streams.stream(0), streams.gpu_index(0), biv_lut_mem->get_lut(0, 0),
        biv_lut_mem->get_degree(0), biv_lut_mem->get_max_degree(0),
        params.glwe_dimension, params.polynomial_size, params.message_modulus,
        params.carry_modulus, generate_bi_lut_lambda, allocate_gpu_memory);

    biv_lut_mem->broadcast_lut(active_streams);

    this->tmp_ct = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), tmp_ct, num_radix_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);
  }

  void release(CudaStreams streams) {
    univ_lut_mem->release(streams);
    delete univ_lut_mem;

    biv_lut_mem->release(streams);
    delete biv_lut_mem;

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   tmp_ct, allocate_gpu_memory);
    delete tmp_ct;
  }
};

template <typename Torus> struct int_count_of_consecutive_bits_buffer {
  int_radix_params params;
  bool allocate_gpu_memory;
  uint32_t counter_num_blocks;

  int_prepare_count_of_consecutive_bits_buffer<Torus> *prepare_mem = nullptr;
  CudaRadixCiphertextFFI *ct_prepared = nullptr;

  int_sum_ciphertexts_vec_memory<Torus> *sum_mem = nullptr;
  int_sc_prop_memory<Torus> *propagate_mem = nullptr;
  CudaRadixCiphertextFFI *cts = nullptr;

  int_count_of_consecutive_bits_buffer(CudaStreams streams,
                                       const int_radix_params params,
                                       uint32_t num_radix_blocks,
                                       uint32_t counter_num_blocks,
                                       Direction direction, BitValue bit_value,
                                       const bool allocate_gpu_memory,
                                       uint64_t &size_tracker) {

    this->params = params;
    this->allocate_gpu_memory = allocate_gpu_memory;
    this->counter_num_blocks = counter_num_blocks;

    this->ct_prepared = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), ct_prepared, num_radix_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

    this->prepare_mem = new int_prepare_count_of_consecutive_bits_buffer<Torus>(
        streams, params, num_radix_blocks, direction, bit_value,
        allocate_gpu_memory, size_tracker);

    this->cts = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), cts,
        counter_num_blocks * num_radix_blocks, params.big_lwe_dimension,
        size_tracker, allocate_gpu_memory);

    this->sum_mem = new int_sum_ciphertexts_vec_memory<Torus>(
        streams, params, counter_num_blocks, num_radix_blocks, true,
        allocate_gpu_memory, size_tracker);

    this->propagate_mem = new int_sc_prop_memory<Torus>(
        streams, params, counter_num_blocks, FLAG_NONE, allocate_gpu_memory,
        size_tracker);
  }

  void release(CudaStreams streams) {

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   ct_prepared, allocate_gpu_memory);
    delete ct_prepared;
    ct_prepared = nullptr;

    prepare_mem->release(streams);
    delete prepare_mem;
    prepare_mem = nullptr;

    sum_mem->release(streams);
    delete sum_mem;
    sum_mem = nullptr;

    propagate_mem->release(streams);
    delete propagate_mem;
    propagate_mem = nullptr;

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0), cts,
                                   allocate_gpu_memory);
    delete cts;
    cts = nullptr;
  }
};

template <typename Torus> struct int_grouped_oprf_memory {
  int_radix_params params;
  bool allocate_gpu_memory;

  int_radix_lut<Torus> *luts;
  CudaRadixCiphertextFFI *plaintext_corrections;
  Torus *h_lut_indexes;

  // with message_bits_per_block == ilog2(msg_modulus) from crypto params
  int_grouped_oprf_memory(CudaStreams streams, int_radix_params params,
                          uint32_t num_blocks_to_process,
                          uint32_t message_bits_per_block,
                          uint64_t total_random_bits, bool allocate_gpu_memory,
                          uint64_t &size_tracker) {

    uint32_t calculated_active_blocks =
        total_random_bits == 0
            ? 0
            : (total_random_bits + message_bits_per_block - 1) /
                  message_bits_per_block;
    if (num_blocks_to_process != calculated_active_blocks) {
      PANIC(
          "num_blocks_to_process should be equal to calculated_active_blocks");
    }

    this->params = params;
    this->allocate_gpu_memory = allocate_gpu_memory;

    this->luts = new int_radix_lut<Torus>(
        streams, params, message_bits_per_block, num_blocks_to_process,
        allocate_gpu_memory, size_tracker);

    this->plaintext_corrections = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), this->plaintext_corrections,
        num_blocks_to_process, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    uint64_t message_modulus_log2 = (uint64_t)std::log2(params.message_modulus);
    if (message_modulus_log2 != message_bits_per_block) {
      PANIC("message_modulus_log2 should be equal to message_bits_per_block");
    }
    uint64_t carry_modulus_log2 = (uint64_t)std::log2(params.carry_modulus);
    uint64_t full_bits_count = 1 + carry_modulus_log2 + message_modulus_log2;
    uint64_t delta = 1ULL << (64 - full_bits_count);
    size_t lwe_size = params.big_lwe_dimension + 1;

    // Pre-generate all possible LUTs.
    //
    for (uint32_t random_bit = 1; random_bit <= message_bits_per_block;
         ++random_bit) {
      uint64_t p = 1ULL << random_bit;
      uint64_t poly_delta =
          2 * static_cast<uint64_t>(params.polynomial_size) / p;

      if (2 * static_cast<uint64_t>(params.polynomial_size) < p) {
        PANIC("2 * static_cast<uint64_t>(params.polynomial_size) should be "
              "smaller than p");
      }

      auto lut_f = [poly_delta, delta](uint32_t x) -> Torus {
        return (2 * (x / poly_delta) + 1) * delta / 2;
      };

      uint64_t degree = 0;
      uint32_t lut_index = random_bit - 1;
      generate_device_accumulator_no_encoding<Torus>(
          streams.stream(0), streams.gpu_index(0), luts->get_lut(0, lut_index),
          degree, params.message_modulus, params.carry_modulus,
          params.glwe_dimension, params.polynomial_size, lut_f,
          allocate_gpu_memory);
      // In  OPRF the degree is hard set to p - 1 instead of the LUT degree
      degree = p - 1;
      *luts->get_degree(lut_index) = degree;
    }

    // For each block, this loop determines the exact number of bits to generate
    // (handling both bounded and unbounded cases), which pre-computed LUT to
    // use, and the final plaintext correction to add.
    //
    Torus *h_corrections =
        (Torus *)calloc(num_blocks_to_process * lwe_size, sizeof(Torus));
    this->h_lut_indexes = (Torus *)calloc(num_blocks_to_process, sizeof(Torus));

    uint64_t bits_processed = 0;
    for (uint32_t i = 0; i < num_blocks_to_process; ++i) {

      if (total_random_bits <= bits_processed) {
        PANIC("total_random_bits should be greater than bits_processed");
      }
      uint64_t bits_remaining = total_random_bits - bits_processed;
      uint32_t bits_for_this_block =
          std::min((uint64_t)message_bits_per_block, bits_remaining);

      uint64_t p = 1ULL << bits_for_this_block;
      Torus plaintext_to_add = (p - 1) * delta / 2;

      h_corrections[i * lwe_size + params.big_lwe_dimension] = plaintext_to_add;
      if (bits_for_this_block < 1) {
        PANIC("bits_for_this_block should be greater than 1");
      }
      this->h_lut_indexes[i] = bits_for_this_block - 1;

      bits_processed += bits_for_this_block;
    }

    // h_corrections contains num_blocks lwes of dimension big_lwe_dim
    // of which num_blocks_to_process lwes have a body that his set
    // to a correction and all others to 0.
    // All lwes in h_corrections have a mask equal to 0.
    // Copy the prepared plaintext corrections to the GPU.
    cuda_memcpy_with_size_tracking_async_to_gpu(
        this->plaintext_corrections->ptr, h_corrections,
        num_blocks_to_process * lwe_size * sizeof(Torus), streams.stream(0),
        streams.gpu_index(0), allocate_gpu_memory);

    // Copy the prepared LUT indexes to the GPU 0, before broadcast to all other
    // GPUs.
    cuda_memcpy_with_size_tracking_async_to_gpu(
        luts->get_lut_indexes(0, 0), this->h_lut_indexes,
        num_blocks_to_process * sizeof(Torus), streams.stream(0),
        streams.gpu_index(0), allocate_gpu_memory);
    auto active_streams = streams.active_gpu_subset(num_blocks_to_process);
    luts->broadcast_lut(active_streams);

    free(h_corrections);
  }

  void release(CudaStreams streams) {
    this->luts->release(streams);
    delete this->luts;
    this->luts = nullptr;

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   this->plaintext_corrections,
                                   this->allocate_gpu_memory);
    delete this->plaintext_corrections;
    this->plaintext_corrections = nullptr;

    free(this->h_lut_indexes);
    this->h_lut_indexes = nullptr;
  }
};

template <typename Torus> struct int_ilog2_buffer {
  int_radix_params params;
  bool allocate_gpu_memory;
  uint32_t input_num_blocks;
  uint32_t counter_num_blocks;
  uint32_t num_bits_in_ciphertext;

  int_prepare_count_of_consecutive_bits_buffer<Torus> *prepare_mem;
  int_sum_ciphertexts_vec_memory<Torus> *sum_mem;
  int_fullprop_buffer<Torus> *final_propagate_mem;

  CudaRadixCiphertextFFI *ct_in_buffer;
  CudaRadixCiphertextFFI *sum_input_cts;
  CudaRadixCiphertextFFI *sum_output_not_propagated;
  CudaRadixCiphertextFFI *message_blocks_not;
  CudaRadixCiphertextFFI *carry_blocks_not;
  CudaRadixCiphertextFFI *rotated_carry_blocks;

  int_radix_lut<Torus> *lut_message_not;
  int_radix_lut<Torus> *lut_carry_not;

  int_ilog2_buffer(CudaStreams streams, const int_radix_params params,
                   uint32_t input_num_blocks, uint32_t counter_num_blocks,
                   uint32_t num_bits_in_ciphertext,
                   const bool allocate_gpu_memory, uint64_t &size_tracker) {

    this->params = params;
    this->allocate_gpu_memory = allocate_gpu_memory;
    this->input_num_blocks = input_num_blocks;
    this->counter_num_blocks = counter_num_blocks;
    this->num_bits_in_ciphertext = num_bits_in_ciphertext;

    this->ct_in_buffer = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), this->ct_in_buffer,
        input_num_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    this->prepare_mem = new int_prepare_count_of_consecutive_bits_buffer<Torus>(
        streams, params, input_num_blocks, Leading, Zero, allocate_gpu_memory,
        size_tracker);

    uint32_t sum_input_total_blocks =
        (input_num_blocks + 1) * counter_num_blocks;
    this->sum_input_cts = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), this->sum_input_cts,
        sum_input_total_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    this->sum_mem = new int_sum_ciphertexts_vec_memory<Torus>(
        streams, params, counter_num_blocks, input_num_blocks + 1, false,
        allocate_gpu_memory, size_tracker);

    this->sum_output_not_propagated = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0),
        this->sum_output_not_propagated, counter_num_blocks,
        params.big_lwe_dimension, size_tracker, allocate_gpu_memory);

    this->lut_message_not =
        new int_radix_lut<Torus>(streams, params, 1, counter_num_blocks,
                                 allocate_gpu_memory, size_tracker);
    std::function<Torus(Torus)> lut_message_lambda =
        [this](uint64_t x) -> uint64_t {
      uint64_t message = x % this->params.message_modulus;
      return (~message) % this->params.message_modulus;
    };
    generate_device_accumulator(streams.stream(0), streams.gpu_index(0),
                                this->lut_message_not->get_lut(0, 0),
                                this->lut_message_not->get_degree(0),
                                this->lut_message_not->get_max_degree(0),
                                params.glwe_dimension, params.polynomial_size,
                                params.message_modulus, params.carry_modulus,
                                lut_message_lambda, allocate_gpu_memory);
    auto active_streams = streams.active_gpu_subset(counter_num_blocks);
    lut_message_not->broadcast_lut(active_streams);

    this->lut_carry_not =
        new int_radix_lut<Torus>(streams, params, 1, counter_num_blocks,
                                 allocate_gpu_memory, size_tracker);
    std::function<Torus(Torus)> lut_carry_lambda =
        [this](uint64_t x) -> uint64_t {
      uint64_t carry = x / this->params.message_modulus;
      return (~carry) % this->params.message_modulus;
    };
    generate_device_accumulator(
        streams.stream(0), streams.gpu_index(0),
        this->lut_carry_not->get_lut(0, 0), this->lut_carry_not->get_degree(0),
        this->lut_carry_not->get_max_degree(0), params.glwe_dimension,
        params.polynomial_size, params.message_modulus, params.carry_modulus,
        lut_carry_lambda, allocate_gpu_memory);
    lut_carry_not->broadcast_lut(active_streams);

    this->message_blocks_not = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), this->message_blocks_not,
        counter_num_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    this->carry_blocks_not = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), this->carry_blocks_not,
        counter_num_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    this->rotated_carry_blocks = new CudaRadixCiphertextFFI;
    create_zero_radix_ciphertext_async<Torus>(
        streams.stream(0), streams.gpu_index(0), this->rotated_carry_blocks,
        counter_num_blocks, params.big_lwe_dimension, size_tracker,
        allocate_gpu_memory);

    this->final_propagate_mem = new int_fullprop_buffer<Torus>(
        streams, params, allocate_gpu_memory, size_tracker);
  }

  void release(CudaStreams streams) {
    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   this->ct_in_buffer, allocate_gpu_memory);
    delete this->ct_in_buffer;
    this->ct_in_buffer = nullptr;

    this->prepare_mem->release(streams);
    delete this->prepare_mem;
    this->prepare_mem = nullptr;

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   this->sum_input_cts, allocate_gpu_memory);
    delete this->sum_input_cts;
    this->sum_input_cts = nullptr;

    this->sum_mem->release(streams);
    delete this->sum_mem;
    this->sum_mem = nullptr;

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   this->sum_output_not_propagated,
                                   allocate_gpu_memory);
    delete this->sum_output_not_propagated;
    this->sum_output_not_propagated = nullptr;

    this->lut_message_not->release(streams);
    delete this->lut_message_not;
    this->lut_message_not = nullptr;

    this->lut_carry_not->release(streams);
    delete this->lut_carry_not;
    this->lut_carry_not = nullptr;

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   this->message_blocks_not,
                                   allocate_gpu_memory);
    delete this->message_blocks_not;
    this->message_blocks_not = nullptr;

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   this->carry_blocks_not, allocate_gpu_memory);
    delete this->carry_blocks_not;
    this->carry_blocks_not = nullptr;

    release_radix_ciphertext_async(streams.stream(0), streams.gpu_index(0),
                                   this->rotated_carry_blocks,
                                   allocate_gpu_memory);
    delete this->rotated_carry_blocks;
    this->rotated_carry_blocks = nullptr;

    this->final_propagate_mem->release(streams);
    delete this->final_propagate_mem;
    this->final_propagate_mem = nullptr;
  }
};

void update_degrees_after_bitand(uint64_t *output_degrees,
                                 uint64_t *lwe_array_1_degrees,
                                 uint64_t *lwe_array_2_degrees,
                                 uint32_t num_radix_blocks);
void update_degrees_after_bitor(uint64_t *output_degrees,
                                uint64_t *lwe_array_1_degrees,
                                uint64_t *lwe_array_2_degrees,
                                uint32_t num_radix_blocks);
void update_degrees_after_bitxor(uint64_t *output_degrees,
                                 uint64_t *lwe_array_1_degrees,
                                 uint64_t *lwe_array_2_degrees,
                                 uint32_t num_radix_blocks);
void update_degrees_after_scalar_bitand(uint64_t *output_degrees,
                                        uint64_t const *clear_degrees,
                                        uint64_t const *input_degrees,
                                        uint32_t num_clear_blocks);
void update_degrees_after_scalar_bitor(uint64_t *output_degrees,
                                       uint64_t const *clear_degrees,
                                       uint64_t const *input_degrees,
                                       uint32_t num_clear_blocks);
void update_degrees_after_scalar_bitxor(uint64_t *output_degrees,
                                        uint64_t const *clear_degrees,
                                        uint64_t const *input_degrees,
                                        uint32_t num_clear_blocks);
std::pair<bool, bool> get_invert_flags(COMPARISON_TYPE compare);
void reverseArray(uint64_t arr[], size_t n);
#endif // CUDA_INTEGER_UTILITIES_H
