ntt.hpp

Go to the documentation of this file.

// Copyright (C) 2020-2021 Intel Corporation
// SPDX-License-Identifier: Apache-2.0

#pragma once

#include <stdint.h>

#include <memory>
#include <vector>

#include "hexl/util/aligned-allocator.hpp"
#include "hexl/util/allocator.hpp"

namespace intel {
namespace hexl {

class NTT {
 public:
  template <class Adaptee, class... Args>
  struct AllocatorAdapter
      : public AllocatorInterface<AllocatorAdapter<Adaptee, Args...>> {
    explicit AllocatorAdapter(Adaptee&& _a, Args&&... args);
    AllocatorAdapter(const Adaptee& _a, Args&... args);

    // interface implementation
    void* allocate_impl(size_t bytes_count);
    void deallocate_impl(void* p, size_t n);

   private:
    Adaptee alloc;
  };

  NTT() = default;

  ~NTT() = default;

  NTT(uint64_t degree, uint64_t q,
      std::shared_ptr<AllocatorBase> alloc_ptr = {});

  template <class Allocator, class... AllocatorArgs>
  NTT(uint64_t degree, uint64_t q, Allocator&& a, AllocatorArgs&&... args)
      : NTT(degree, q,
            std::static_pointer_cast<AllocatorBase>(
                std::make_shared<AllocatorAdapter<Allocator, AllocatorArgs...>>(
                    std::move(a), std::forward<AllocatorArgs>(args)...))) {}

  NTT(uint64_t degree, uint64_t q, uint64_t root_of_unity,
      std::shared_ptr<AllocatorBase> alloc_ptr = {});

  template <class Allocator, class... AllocatorArgs>
  NTT(uint64_t degree, uint64_t q, uint64_t root_of_unity, Allocator&& a,
      AllocatorArgs&&... args)
      : NTT(degree, q, root_of_unity,
            std::static_pointer_cast<AllocatorBase>(
                std::make_shared<AllocatorAdapter<Allocator, AllocatorArgs...>>(
                    std::move(a), std::forward<AllocatorArgs>(args)...))) {}

  static bool CheckArguments(uint64_t degree, uint64_t modulus);

  void ComputeForward(uint64_t* result, const uint64_t* operand,
                      uint64_t input_mod_factor, uint64_t output_mod_factor);

  void ComputeInverse(uint64_t* result, const uint64_t* operand,
                      uint64_t input_mod_factor, uint64_t output_mod_factor);

  uint64_t GetMinimalRootOfUnity() const { return m_w; }

  uint64_t GetDegree() const { return m_degree; }

  uint64_t GetModulus() const { return m_q; }

  const AlignedVector64<uint64_t>& GetRootOfUnityPowers() const {
    return m_root_of_unity_powers;
  }

  uint64_t GetRootOfUnityPower(size_t i) { return GetRootOfUnityPowers()[i]; }

  const AlignedVector64<uint64_t>& GetPrecon32RootOfUnityPowers() const {
    return m_precon32_root_of_unity_powers;
  }

  const AlignedVector64<uint64_t>& GetPrecon64RootOfUnityPowers() const {
    return m_precon64_root_of_unity_powers;
  }

  const AlignedVector64<uint64_t>& GetAVX512RootOfUnityPowers() const {
    return m_avx512_root_of_unity_powers;
  }

  const AlignedVector64<uint64_t>& GetAVX512Precon32RootOfUnityPowers() const {
    return m_avx512_precon32_root_of_unity_powers;
  }

  const AlignedVector64<uint64_t>& GetAVX512Precon52RootOfUnityPowers() const {
    return m_avx512_precon52_root_of_unity_powers;
  }

  const AlignedVector64<uint64_t>& GetAVX512Precon64RootOfUnityPowers() const {
    return m_avx512_precon64_root_of_unity_powers;
  }

  const AlignedVector64<uint64_t>& GetInvRootOfUnityPowers() const {
    return m_inv_root_of_unity_powers;
  }

  uint64_t GetInvRootOfUnityPower(size_t i) {
    return GetInvRootOfUnityPowers()[i];
  }

  // powers for the modulus and root of unity.
  const AlignedVector64<uint64_t>& GetPrecon32InvRootOfUnityPowers() const {
    return m_precon32_inv_root_of_unity_powers;
  }

  // powers for the modulus and root of unity.
  const AlignedVector64<uint64_t>& GetPrecon52InvRootOfUnityPowers() const {
    return m_precon52_inv_root_of_unity_powers;
  }

  // powers for the modulus and root of unity.
  const AlignedVector64<uint64_t>& GetPrecon64InvRootOfUnityPowers() const {
    return m_precon64_inv_root_of_unity_powers;
  }

  static size_t MaxDegreeBits() { return 20; }

  static size_t MaxModulusBits() { return 62; }

  static const size_t s_default_shift_bits{64};

  static const size_t s_ifma_shift_bits{52};

  static const size_t s_max_fwd_32_modulus{1ULL << (32 - 2)};

  static const size_t s_max_inv_32_modulus{1ULL << (32 - 1)};

  static const size_t s_max_fwd_ifma_modulus{1ULL << (s_ifma_shift_bits - 2)};

  static const size_t s_max_inv_ifma_modulus{1ULL << (s_ifma_shift_bits - 1)};

 private:
  void ComputeRootOfUnityPowers();

  uint64_t m_degree;  // N: size of NTT transform, should be power of 2
  uint64_t m_q;       // prime modulus. Must satisfy q == 1 mod 2n

  uint64_t m_degree_bits;  // log_2(m_degree)

  uint64_t m_w_inv;  // Inverse of minimal root of unity
  uint64_t m_w;      // A 2N'th root of unity

  std::shared_ptr<AllocatorBase> m_alloc;

  AlignedAllocator<uint64_t, 64> m_aligned_alloc;

  // powers of the minimal root of unity
  AlignedVector64<uint64_t> m_root_of_unity_powers;
  // vector of floor(W * 2**32 / m_q), with W the root of unity powers
  AlignedVector64<uint64_t> m_precon32_root_of_unity_powers;
  // vector of floor(W * 2**64 / m_q), with W the root of unity powers
  AlignedVector64<uint64_t> m_precon64_root_of_unity_powers;

  // powers of the minimal root of unity adjusted for use in AVX512
  // implementations
  AlignedVector64<uint64_t> m_avx512_root_of_unity_powers;
  // vector of floor(W * 2**32 / m_q), with W the AVX512 root of unity powers
  AlignedVector64<uint64_t> m_avx512_precon32_root_of_unity_powers;
  // vector of floor(W * 2**52 / m_q), with W the AVX512 root of unity powers
  AlignedVector64<uint64_t> m_avx512_precon52_root_of_unity_powers;
  // vector of floor(W * 2**64 / m_q), with W the AVX512 root of unity powers
  AlignedVector64<uint64_t> m_avx512_precon64_root_of_unity_powers;

  // vector of floor(W * 2**32 / m_q), with W the inverse root of unity powers
  AlignedVector64<uint64_t> m_precon32_inv_root_of_unity_powers;
  // vector of floor(W * 2**52 / m_q), with W the inverse root of unity powers
  AlignedVector64<uint64_t> m_precon52_inv_root_of_unity_powers;
  // vector of floor(W * 2**64 / m_q), with W the inverse root of unity powers
  AlignedVector64<uint64_t> m_precon64_inv_root_of_unity_powers;

  AlignedVector64<uint64_t> m_inv_root_of_unity_powers;
};

}  // namespace hexl
}  // namespace intel