concrete/backends/concrete-cuda/implementation/test/test_circuit_bootstrap.cpp

#include "../include/circuit_bootstrap.h"
#include "../include/device.h"
#include "concrete-cpu.h"
#include "utils.h"
#include "gtest/gtest.h"
#include <cstdint>
#include <stdio.h>
#include <stdlib.h>

const unsigned REPETITIONS = 2;
const unsigned SAMPLES = 50;

typedef struct {
  int lwe_dimension;
  int glwe_dimension;
  int polynomial_size;
  double lwe_modular_variance;
  double glwe_modular_variance;
  int pbs_base_log;
  int pbs_level;
  int pksk_base_log;
  int pksk_level;
  int cbs_base_log;
  int cbs_level;
} CircuitBootstrapTestParams;

class CircuitBootstrapTestPrimitives_u64
    : public ::testing::TestWithParam<CircuitBootstrapTestParams> {
protected:
  int lwe_dimension;
  int glwe_dimension;
  int polynomial_size;
  double lwe_modular_variance;
  double glwe_modular_variance;
  int pbs_base_log;
  int pbs_level;
  int pksk_base_log;
  int pksk_level;
  int cbs_base_log;
  int cbs_level;
  int number_of_bits_of_message_including_padding;
  int ggsw_size;
  uint64_t delta;
  int delta_log;
  Csprng *csprng;
  cudaStream_t *stream;
  int gpu_index = 0;
  uint64_t *lwe_sk_in_array;
  uint64_t *lwe_sk_out_array;
  uint64_t *lwe_in_ct;
  uint64_t *ggsw_out_ct;
  uint64_t *plaintexts;
  double *d_fourier_bsk_array;
  uint64_t *d_pksk_array;
  uint64_t *d_lwe_in_ct;
  uint64_t *d_ggsw_out_ct;
  uint64_t *d_lut_vector_indexes;
  int8_t *cbs_buffer;

public:
  // Test arithmetic functions
  void SetUp() {
    stream = cuda_create_stream(0);

    // TestParams
    lwe_dimension = (int)GetParam().lwe_dimension;
    glwe_dimension = (int)GetParam().glwe_dimension;
    polynomial_size = (int)GetParam().polynomial_size;
    lwe_modular_variance = (double)GetParam().lwe_modular_variance;
    glwe_modular_variance = (double)GetParam().glwe_modular_variance;
    pbs_base_log = (int)GetParam().pbs_base_log;
    pbs_level = (int)GetParam().pbs_level;
    pksk_base_log = (int)GetParam().pksk_base_log;
    pksk_level = (int)GetParam().pksk_level;
    cbs_base_log = (int)GetParam().cbs_base_log;
    cbs_level = (int)GetParam().cbs_level;
    // We generate binary messages
    number_of_bits_of_message_including_padding = 2;
    delta_log = 60;
    delta = (uint64_t)(1) << delta_log;
    ggsw_size = cbs_level * (glwe_dimension + 1) * (glwe_dimension + 1) *
                polynomial_size;

    // Create a Csprng
    csprng =
        (Csprng *)aligned_alloc(CONCRETE_CSPRNG_ALIGN, CONCRETE_CSPRNG_SIZE);
    uint8_t seed[16] = {(uint8_t)0};
    concrete_cpu_construct_concrete_csprng(
        csprng, Uint128{.little_endian_bytes = {*seed}});

    // Generate the keys
    generate_lwe_secret_keys(&lwe_sk_in_array, lwe_dimension, csprng,
                             REPETITIONS);
    generate_lwe_secret_keys(&lwe_sk_out_array,
                             glwe_dimension * polynomial_size, csprng,
                             REPETITIONS);
    generate_lwe_bootstrap_keys(
        stream, gpu_index, &d_fourier_bsk_array, lwe_sk_in_array,
        lwe_sk_out_array, lwe_dimension, glwe_dimension, polynomial_size,
        pbs_level, pbs_base_log, csprng, glwe_modular_variance, REPETITIONS);
    generate_lwe_private_functional_keyswitch_key_lists(
        stream, gpu_index, &d_pksk_array, lwe_sk_out_array, lwe_sk_out_array,
        glwe_dimension * polynomial_size, glwe_dimension, polynomial_size,
        pksk_level, pksk_base_log, csprng, lwe_modular_variance, REPETITIONS);
    plaintexts =
        generate_plaintexts(number_of_bits_of_message_including_padding, delta,
                            1, REPETITIONS, SAMPLES);

    d_ggsw_out_ct = (uint64_t *)cuda_malloc_async(ggsw_size * sizeof(uint64_t),
                                                  stream, gpu_index);

    d_lwe_in_ct = (uint64_t *)cuda_malloc_async(
        (lwe_dimension + 1) * sizeof(uint64_t), stream, gpu_index);

    lwe_in_ct = (uint64_t *)malloc((lwe_dimension + 1) * sizeof(uint64_t));
    ggsw_out_ct = (uint64_t *)malloc(ggsw_size * sizeof(uint64_t));
    // Execute cbs scratch
    scratch_cuda_circuit_bootstrap_64(
        stream, gpu_index, &cbs_buffer, glwe_dimension, lwe_dimension,
        polynomial_size, cbs_level, 1, cuda_get_max_shared_memory(gpu_index),
        true);
    // Build LUT vector indexes
    uint64_t *h_lut_vector_indexes =
        (uint64_t *)malloc(cbs_level * sizeof(uint64_t));
    for (int index = 0; index < cbs_level; index++) {
      h_lut_vector_indexes[index] = 0; // index % cbs_level;
    }
    d_lut_vector_indexes = (uint64_t *)cuda_malloc_async(
        cbs_level * sizeof(uint64_t), stream, gpu_index);
    cuda_memcpy_async_to_gpu(d_lut_vector_indexes, h_lut_vector_indexes,
                             cbs_level * sizeof(uint64_t), stream, gpu_index);
    free(h_lut_vector_indexes);
  }

  void TearDown() {
    void *v_stream = (void *)stream;

    cuda_synchronize_stream(v_stream);
    concrete_cpu_destroy_concrete_csprng(csprng);
    free(csprng);
    free(lwe_sk_in_array);
    free(lwe_sk_out_array);
    free(plaintexts);
    free(lwe_in_ct);
    free(ggsw_out_ct);
    cleanup_cuda_circuit_bootstrap(stream, gpu_index, &cbs_buffer);
    cuda_drop_async(d_fourier_bsk_array, stream, gpu_index);
    cuda_drop_async(d_pksk_array, stream, gpu_index);
    cuda_drop_async(d_lwe_in_ct, stream, gpu_index);
    cuda_drop_async(d_ggsw_out_ct, stream, gpu_index);
    cuda_drop_async(d_lut_vector_indexes, stream, gpu_index);
    cuda_destroy_stream(stream, gpu_index);
  }
};

TEST_P(CircuitBootstrapTestPrimitives_u64, circuit_bootstrap) {
  void *v_stream = (void *)stream;
  for (uint r = 0; r < REPETITIONS; r++) {
    for (uint s = 0; s < SAMPLES; s++) {
      uint64_t plaintext = plaintexts[r * SAMPLES + s];
      int bsk_size = (glwe_dimension + 1) * (glwe_dimension + 1) * pbs_level *
                     polynomial_size * (lwe_dimension + 1);
      double *d_fourier_bsk = d_fourier_bsk_array + (ptrdiff_t)(bsk_size * r);
      int pksk_list_size = pksk_level * (glwe_dimension + 1) * polynomial_size *
                           (glwe_dimension * polynomial_size + 1) *
                           (glwe_dimension + 1);
      uint64_t *d_pksk_list = d_pksk_array + (ptrdiff_t)(pksk_list_size * r);
      uint64_t *lwe_in_sk = lwe_sk_in_array + (ptrdiff_t)(lwe_dimension * r);
      uint64_t *lwe_sk_out =
          lwe_sk_out_array + (ptrdiff_t)(r * glwe_dimension * polynomial_size);
      concrete_cpu_encrypt_lwe_ciphertext_u64(
          lwe_in_sk, lwe_in_ct, plaintext, lwe_dimension, lwe_modular_variance,
          csprng, &CONCRETE_CSPRNG_VTABLE);
      cuda_synchronize_stream(v_stream);
      cuda_memcpy_async_to_gpu(d_lwe_in_ct, lwe_in_ct,
                               (lwe_dimension + 1) * sizeof(uint64_t), stream,
                               gpu_index);

      // Execute circuit bootstrap
      cuda_circuit_bootstrap_64(
          stream, gpu_index, (void *)d_ggsw_out_ct, (void *)d_lwe_in_ct,
          (void *)d_fourier_bsk, (void *)d_pksk_list,
          (void *)d_lut_vector_indexes, cbs_buffer, delta_log, polynomial_size,
          glwe_dimension, lwe_dimension, pbs_level, pbs_base_log, pksk_level,
          pksk_base_log, cbs_level, cbs_base_log, 1,
          cuda_get_max_shared_memory(gpu_index));

      // Copy result back
      cuda_memcpy_async_to_cpu(ggsw_out_ct, d_ggsw_out_ct,
                               ggsw_size * sizeof(uint64_t), stream, gpu_index);
      cuda_synchronize_stream(v_stream);

      uint64_t *decrypted =
          (uint64_t *)malloc(polynomial_size * (glwe_dimension + 1) *
                             cbs_level * sizeof(uint64_t));

      uint64_t multiplying_factor = -(plaintext >> delta_log);
      for (int l = 1; l < cbs_level + 1; l++) {
        for (int j = 0; j < glwe_dimension; j++) {
          uint64_t *res = decrypted + (ptrdiff_t)((l - 1) * polynomial_size *
                                                      (glwe_dimension + 1) +
                                                  j * polynomial_size);
          uint64_t *glwe_ct_out =
              ggsw_out_ct +
              (ptrdiff_t)((l - 1) * polynomial_size * (glwe_dimension + 1) *
                              (glwe_dimension + 1) +
                          j * polynomial_size * (glwe_dimension + 1));
          concrete_cpu_decrypt_glwe_ciphertext_u64(
              lwe_sk_out, res, glwe_ct_out, glwe_dimension, polynomial_size);

          for (int k = 0; k < polynomial_size; k++) {
            uint64_t expected_decryption =
                lwe_sk_out[j * polynomial_size + k] * multiplying_factor;
            expected_decryption >>= (64 - cbs_base_log * l);
            uint64_t decoded_plaintext =
                closest_representable(res[k], l, cbs_base_log) >>
                (64 - cbs_base_log * l);
            EXPECT_EQ(expected_decryption, decoded_plaintext);
          }
        }
      }
      // Check last glwe on last level
      uint64_t *res =
          decrypted +
          (ptrdiff_t)((cbs_level - 1) * polynomial_size * (glwe_dimension + 1) +
                      glwe_dimension * polynomial_size);
      uint64_t *glwe_ct_out =
          ggsw_out_ct +
          (ptrdiff_t)((cbs_level - 1) * polynomial_size * (glwe_dimension + 1) *
                          (glwe_dimension + 1) +
                      glwe_dimension * polynomial_size * (glwe_dimension + 1));
      concrete_cpu_decrypt_glwe_ciphertext_u64(lwe_sk_out, res, glwe_ct_out,
                                               glwe_dimension, polynomial_size);

      for (int k = 0; k < polynomial_size; k++) {
        uint64_t expected_decryption = (k == 0) ? plaintext / delta : 0;
        uint64_t decoded_plaintext =
            closest_representable(res[k], cbs_level, cbs_base_log) >>
            (64 - cbs_base_log * cbs_level);
        EXPECT_EQ(expected_decryption, decoded_plaintext);
      }
      free(decrypted);
    }
  }
}

// Defines for which parameters set the PBS will be tested.
// It executes each test for all pairs on phis X qs (Cartesian product)
::testing::internal::ParamGenerator<CircuitBootstrapTestParams> cbs_params_u64 =
    ::testing::Values(
        // n, k, N, lwe_variance, glwe_variance, pbs_base_log, pbs_level,
        // pksk_base_log, pksk_level, cbs_base_log, cbs_level
        (CircuitBootstrapTestParams){10, 2, 512, 7.52316384526264e-37,
                                     7.52316384526264e-37, 11, 2, 15, 2, 10,
                                     1});

std::string
printParamName(::testing::TestParamInfo<CircuitBootstrapTestParams> p) {
  CircuitBootstrapTestParams params = p.param;

  return "n_" + std::to_string(params.lwe_dimension) + "_k_" +
         std::to_string(params.glwe_dimension) + "_N_" +
         std::to_string(params.polynomial_size) + "_pbs_base_log_" +
         std::to_string(params.pbs_base_log) + "_pbs_level_" +
         std::to_string(params.pbs_level) + "_pksk_base_log_" +
         std::to_string(params.pksk_base_log) + "_pksk_level_" +
         std::to_string(params.pksk_level) + "_cbs_base_log_" +
         std::to_string(params.cbs_base_log) + "_cbs_level_" +
         std::to_string(params.cbs_level);
}

INSTANTIATE_TEST_CASE_P(CircuitBootstrapInstantiation,
                        CircuitBootstrapTestPrimitives_u64, cbs_params_u64,
                        printParamName);