text: // Copyright (C) 2020-2021 Intel Corporation
  // SPDX-License-Identifier: Apache-2.0
  
  #include <benchmark/benchmark.h>
  
  #include <vector>
  
  #include "hexl/logging/logging.hpp"
  #include "hexl/ntt/ntt.hpp"
  #include "hexl/number-theory/number-theory.hpp"
  #include "hexl/util/aligned-allocator.hpp"
  #include "ntt/fwd-ntt-avx512.hpp"
  #include "ntt/inv-ntt-avx512.hpp"
  #include "ntt/ntt-internal.hpp"
  
  namespace intel {
  namespace hexl {
  
  // Forward transforms
  
  //=================================================================
  
  // state[0] is the degree
  static void BM_FwdNTTNative(benchmark::State& state) {  //  NOLINT
    size_t ntt_size = state.range(0);
    size_t modulus = GeneratePrimes(1, 45, ntt_size)[0];
  
    AlignedVector64<uint64_t> input(ntt_size, 1);
    NTT ntt(ntt_size, modulus);
  
    for (auto _ : state) {
      ForwardTransformToBitReverse64(
          input.data(), ntt_size, modulus, ntt.GetRootOfUnityPowers().data(),
          ntt.GetPrecon64RootOfUnityPowers().data(), 2, 1);
    }
  }
  
  BENCHMARK(BM_FwdNTTNative)
      ->Unit(benchmark::kMicrosecond)
      ->Args({1024})
      ->Args({4096})
      ->Args({16384});
  //=================================================================
  
  #ifdef HEXL_HAS_AVX512IFMA
  // state[0] is the degree
  static void BM_FwdNTT_AVX512IFMA(benchmark::State& state) {  //  NOLINT
    size_t ntt_size = state.range(0);
    size_t modulus_bits = 49;
    size_t modulus = GeneratePrimes(1, modulus_bits, ntt_size)[0];
  
    AlignedVector64<uint64_t> input(ntt_size, 1);
    NTT ntt(ntt_size, modulus);
  
    const AlignedVector64<uint64_t> root_of_unity =
        ntt.GetAVX512RootOfUnityPowers();
    const AlignedVector64<uint64_t> precon_root_of_unity =
        ntt.GetAVX512Precon52RootOfUnityPowers();
  
    for (auto _ : state) {
      ForwardTransformToBitReverseAVX512<NTT::s_ifma_shift_bits>(
          input.data(), ntt_size, modulus, root_of_unity.data(),
          precon_root_of_unity.data(), 2, 1);
    }
  }
  
  BENCHMARK(BM_FwdNTT_AVX512IFMA)
      ->Unit(benchmark::kMicrosecond)
      ->Args({1024})
      ->Args({4096})
      ->Args({16384});
  
  //=================================================================
  
  // state[0] is the degree
  static void BM_FwdNTT_AVX512IFMALazy(benchmark::State& state) {  //  NOLINT
    size_t ntt_size = state.range(0);
    size_t modulus_bits = 49;
    size_t modulus = GeneratePrimes(1, modulus_bits, ntt_size)[0];
  
    AlignedVector64<uint64_t> input(ntt_size, 1);
    NTT ntt(ntt_size, modulus);
  
    const AlignedVector64<uint64_t> root_of_unity =
        ntt.GetAVX512RootOfUnityPowers();
    const AlignedVector64<uint64_t> precon_root_of_unity =
        ntt.GetAVX512Precon52RootOfUnityPowers();
  
    for (auto _ : state) {
      ForwardTransformToBitReverseAVX512<NTT::s_ifma_shift_bits>(
          input.data(), ntt_size, modulus, root_of_unity.data(),
          precon_root_of_unity.data(), 4, 4);
    }
  }
  
  BENCHMARK(BM_FwdNTT_AVX512IFMALazy)
      ->Unit(benchmark::kMicrosecond)
      ->Args({1024})
      ->Args({4096})
      ->Args({16384});
  
  //=================================================================
  
  #endif
  
  //=================================================================
  
  #ifdef HEXL_HAS_AVX512DQ
  // state[0] is the degree
  // state[1] is the output modulus factor
  static void BM_FwdNTT_AVX512DQ_32(benchmark::State& state) {  //  NOLINT
    size_t ntt_size = state.range(0);
    uint64_t output_mod_factor = state.range(1);
    size_t modulus_bits = 29;
    size_t modulus = GeneratePrimes(1, modulus_bits, ntt_size)[0];
  
    AlignedVector64<uint64_t> input(ntt_size, 1);
    NTT ntt(ntt_size, modulus);
  
    const AlignedVector64<uint64_t> root_of_unity =
        ntt.GetAVX512RootOfUnityPowers();
    const AlignedVector64<uint64_t> precon_root_of_unity =
        ntt.GetAVX512Precon32RootOfUnityPowers();
    for (auto _ : state) {
      ForwardTransformToBitReverseAVX512<32>(
          input.data(), ntt_size, modulus, root_of_unity.data(),
          precon_root_of_unity.data(), 4, output_mod_factor);
    }
  }
  
  BENCHMARK(BM_FwdNTT_AVX512DQ_32)
      ->Unit(benchmark::kMicrosecond)
      ->Args({1024, 1})
      ->Args({1024, 4})
      ->Args({4096, 1})
      ->Args({4096, 4})
      ->Args({16384, 1})
      ->Args({16384, 4});
  
  // state[0] is the degree
  // state[1] is the output modulus factor
  static void BM_FwdNTT_AVX512DQ_64(benchmark::State& state) {  //  NOLINT
    size_t ntt_size = state.range(0);
    uint64_t output_mod_factor = state.range(1);
    size_t modulus_bits = 55;
    size_t modulus = GeneratePrimes(1, modulus_bits, ntt_size)[0];
  
    AlignedVector64<uint64_t> input(ntt_size, 1);
    NTT ntt(ntt_size, modulus);
  
    const AlignedVector64<uint64_t> root_of_unity =
        ntt.GetAVX512RootOfUnityPowers();
    const AlignedVector64<uint64_t> precon_root_of_unity =
        ntt.GetAVX512Precon64RootOfUnityPowers();
    for (auto _ : state) {
      ForwardTransformToBitReverseAVX512<64>(
          input.data(), ntt_size, modulus, root_of_unity.data(),
          precon_root_of_unity.data(), 4, output_mod_factor);
    }
  }
  
  BENCHMARK(BM_FwdNTT_AVX512DQ_64)
      ->Unit(benchmark::kMicrosecond)
      ->Args({1024, 1})
      ->Args({1024, 4})
      ->Args({4096, 1})
      ->Args({4096, 4})
      ->Args({16384, 1})
      ->Args({16384, 4});
  
  #endif
  
  //=================================================================
  
  // state[0] is the degree
  static void BM_FwdNTTInPlace(benchmark::State& state) {  //  NOLINT
    size_t ntt_size = state.range(0);
    size_t modulus = GeneratePrimes(1, 61, ntt_size)[0];
  
    AlignedVector64<uint64_t> input(ntt_size, 1);
    NTT ntt(ntt_size, modulus);
  
    for (auto _ : state) {
      ntt.ComputeForward(input.data(), input.data(), 1, 1);
    }
  }
  
  BENCHMARK(BM_FwdNTTInPlace)
      ->Unit(benchmark::kMicrosecond)
      ->Args({1024})
      ->Args({4096})
      ->Args({16384});
  
  //=================================================================
  
  // state[0] is the degree
  static void BM_FwdNTTCopy(benchmark::State& state) {  //  NOLINT
    size_t ntt_size = state.range(0);
    size_t modulus = GeneratePrimes(1, 61, ntt_size)[0];
  
    AlignedVector64<uint64_t> input(ntt_size, 1);
    AlignedVector64<uint64_t> output(ntt_size, 1);
    NTT ntt(ntt_size, modulus);
  
    for (auto _ : state) {
      ntt.ComputeForward(input.data(), output.data(), 1, 1);
    }
  }
  
  BENCHMARK(BM_FwdNTTCopy)
      ->Unit(benchmark::kMicrosecond)
      ->Args({1024})
      ->Args({4096})
      ->Args({16384});
  
  //=================================================================
  
  // Inverse transforms
  
  // state[0] is the degree
  static void BM_InvNTTNative(benchmark::State& state) {  //  NOLINT
    size_t ntt_size = state.range(0);
    size_t modulus = GeneratePrimes(1, 45, ntt_size)[0];
  
    AlignedVector64<uint64_t> input(ntt_size, 1);
    NTT ntt(ntt_size, modulus);
  
    const AlignedVector64<uint64_t> root_of_unity = ntt.GetInvRootOfUnityPowers();
    const AlignedVector64<uint64_t> precon_root_of_unity =
        ntt.GetPrecon64InvRootOfUnityPowers();
    for (auto _ : state) {
      InverseTransformFromBitReverse64(input.data(), ntt_size, modulus,
                                       root_of_unity.data(),
                                       precon_root_of_unity.data(), 1, 1);
    }
  }
  
  BENCHMARK(BM_InvNTTNative)
      ->Unit(benchmark::kMicrosecond)
      ->Args({1024})
      ->Args({4096})
      ->Args({16384});
  
  //=================================================================
  
  #ifdef HEXL_HAS_AVX512IFMA
  // state[0] is the degree
  static void BM_InvNTT_AVX512IFMA(benchmark::State& state) {  //  NOLINT
    size_t ntt_size = state.range(0);
    size_t modulus = GeneratePrimes(1, 49, ntt_size)[0];
  
    AlignedVector64<uint64_t> input(ntt_size, 1);
    NTT ntt(ntt_size, modulus);
  
    const AlignedVector64<uint64_t> root_of_unity = ntt.GetInvRootOfUnityPowers();
    const AlignedVector64<uint64_t> precon_root_of_unity =
        ntt.GetPrecon52InvRootOfUnityPowers();
    for (auto _ : state) {
      InverseTransformFromBitReverseAVX512<NTT::s_ifma_shift_bits>(
          input.data(), ntt_size, modulus, root_of_unity.data(),
          precon_root_of_unity.data(), 1, 1);
    }
  }
  
  BENCHMARK(BM_InvNTT_AVX512IFMA)
      ->Unit(benchmark::kMicrosecond)
      ->Args({1024})
      ->Args({4096})
      ->Args({16384});
  
  //=================================================================
  
  // state[0] is the degree
  static void BM_InvNTT_AVX512IFMALazy(benchmark::State& state) {  //  NOLINT
    size_t ntt_size = state.range(0);
    size_t modulus = GeneratePrimes(1, 49, ntt_size)[0];
  
    AlignedVector64<uint64_t> input(ntt_size, 1);
    NTT ntt(ntt_size, modulus);
  
    const AlignedVector64<uint64_t> root_of_unity = ntt.GetInvRootOfUnityPowers();
    const AlignedVector64<uint64_t> precon_root_of_unity =
        ntt.GetPrecon52InvRootOfUnityPowers();
    for (auto _ : state) {
      InverseTransformFromBitReverseAVX512<NTT::s_ifma_shift_bits>(
          input.data(), ntt_size, modulus, root_of_unity.data(),
          precon_root_of_unity.data(), 2, 2);
    }
  }
  
  BENCHMARK(BM_InvNTT_AVX512IFMALazy)
      ->Unit(benchmark::kMicrosecond)
      ->Args({1024})
      ->Args({4096})
      ->Args({16384});
  
  #endif
  
  //=================================================================
  
  #ifdef HEXL_HAS_AVX512DQ
  // state[0] is the degree
  static void BM_InvNTT_AVX512DQ_32(benchmark::State& state) {  //  NOLINT
    size_t ntt_size = state.range(0);
    uint64_t output_mod_factor = state.range(1);
    size_t modulus = GeneratePrimes(1, 30, ntt_size)[0];
  
    AlignedVector64<uint64_t> input(ntt_size, 1);
    NTT ntt(ntt_size, modulus);
  
    const AlignedVector64<uint64_t> root_of_unity = ntt.GetInvRootOfUnityPowers();
    const AlignedVector64<uint64_t> precon_root_of_unity =
        ntt.GetPrecon32InvRootOfUnityPowers();
  
    for (auto _ : state) {
      InverseTransformFromBitReverseAVX512<32>(
          input.data(), ntt_size, modulus, root_of_unity.data(),
          precon_root_of_unity.data(), output_mod_factor, output_mod_factor);
    }
  }
  
  BENCHMARK(BM_InvNTT_AVX512DQ_32)
      ->Unit(benchmark::kMicrosecond)
      ->Args({1024, 1})
      ->Args({1024, 2})
      ->Args({4096, 1})
      ->Args({4096, 2})
      ->Args({16384, 1})
      ->Args({16384, 2});
  
  static void BM_InvNTT_AVX512DQ_64(benchmark::State& state) {  //  NOLINT
    size_t ntt_size = state.range(0);
    uint64_t output_mod_factor = state.range(1);
    size_t modulus = GeneratePrimes(1, 62, ntt_size)[0];
  
    AlignedVector64<uint64_t> input(ntt_size, 1);
    NTT ntt(ntt_size, modulus);
  
    const AlignedVector64<uint64_t> root_of_unity = ntt.GetInvRootOfUnityPowers();
    const AlignedVector64<uint64_t> precon_root_of_unity =
        ntt.GetPrecon64InvRootOfUnityPowers();
  
    for (auto _ : state) {
      InverseTransformFromBitReverseAVX512<NTT::s_default_shift_bits>(
          input.data(), ntt_size, modulus, root_of_unity.data(),
          precon_root_of_unity.data(), output_mod_factor, output_mod_factor);
    }
  }
  
  BENCHMARK(BM_InvNTT_AVX512DQ_64)
      ->Unit(benchmark::kMicrosecond)
      ->Args({1024, 1})
      ->Args({1024, 2})
      ->Args({4096, 1})
      ->Args({4096, 2})
      ->Args({16384, 1})
      ->Args({16384, 2});
  #endif
  
  //=================================================================
  
  }  // namespace hexl
  }  // namespace intel
  