feat(math): introduce MixedRadixEvaluationDomain

tachyon/math/polynomial/domains/mixed_radix/BUILD.bazel

@@ -0,0 +1,28 @@
+load("//bazel:tachyon_cc.bzl", "tachyon_cc_library", "tachyon_cc_test")
+
+package(default_visibility = ["//visibility:public"])
+
+tachyon_cc_library(
+    name = "mixed_radix_evaluation_domain",
+    hdrs = ["mixed_radix_evaluation_domain.h"],
+    deps = [
+        "//tachyon/base:logging",
+        "//tachyon/base/strings:string_number_conversions",
+        "//tachyon/math/base/gmp:gmp_util",
+        "//tachyon/math/finite_fields:prime_field_base",
+        "//tachyon/math/polynomial/domains:evaluation_domain",
+        "//tachyon/math/polynomial/domains:evaluation_domain_util",
+    ],
+)
+
+tachyon_cc_test(
+    name = "mixed_radix_evaluation_domain_unittest",
+    size = "small",
+    srcs = ["mixed_radix_evaluation_domain_unittest.cc"],
+    deps = [
+        ":mixed_radix_evaluation_domain",
+        "//tachyon/math/elliptic_curves/bn/bn384_small_two_adicity:fq",
+        "//tachyon/math/polynomial/domains:evaluation_domain",
+        "//tachyon/math/polynomial/polynomials/univariate:univariate_polynomial",
+    ],
+)

tachyon/math/polynomial/domains/mixed_radix/mixed_radix_evaluation_domain.h

@@ -0,0 +1,356 @@
+// Copyright 2022 arkworks contributors
+// Use of this source code is governed by a MIT/Apache-2.0 style license that
+// can be found in the LICENSE-MIT.arkworks and the LICENCE-APACHE.arkworks
+// file.
+
+// This header contains a |MixedRadixEvaluationDomain| for performing various
+// kinds of polynomial arithmetic on top of fields that are FFT-friendly but do
+// not have high-enough two-adicity to perform the FFT efficiently, i.e. the
+// multiplicative subgroup G generated by |F::Config::kTwoAdicRootOfUnity| is
+// not large enough. |MixedRadixEvaluationDomain| resolves this issue by using a
+// larger subgroup obtained by combining G with another subgroup of size
+// |F::Config::kSmallSubgroupBase|^|F::Config::kSmallSubgroupAdicity|, to obtain
+// a subgroup generated by |F::Config::kLargeSubgroupRootOfUnity|.
+
+#ifndef TACHYON_MATH_POLYNOMIAL_DOMAINS_MIXED_RADIX_MIXED_RADIX_EVALUATION_DOMAIN_H_
+#define TACHYON_MATH_POLYNOMIAL_DOMAINS_MIXED_RADIX_MIXED_RADIX_EVALUATION_DOMAIN_H_
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include "tachyon/base/logging.h"
+#include "tachyon/base/strings/string_number_conversions.h"
+#include "tachyon/math/base/gmp/gmp_util.h"
+#include "tachyon/math/finite_fields/prime_field_util.h"
+#include "tachyon/math/polynomial/domains/evaluation_domain.h"
+#include "tachyon/math/polynomial/domains/evaluation_domain_util.h"
+
+namespace tachyon::math {
+
+// Defines a domain over which finite field (I)FFTs can be performed. Works only
+// for fields that have a multiplicative subgroup of size that is a power-of-2
+// and another small subgroup over a different base defined.
+template <typename F>
+class MixedRadixEvaluationDomain
+    : public EvaluationDomain<MixedRadixEvaluationDomain<F>> {
+ public:
+  using Field = F;
+
+  constexpr MixedRadixEvaluationDomain() = default;
+  // Construct a domain that is large enough for evaluations of a polynomial
+  // having |num_coeffs| coefficients.
+  constexpr explicit MixedRadixEvaluationDomain(size_t num_coeffs) {
+    CHECK(F::Config::kHasLargeSubgroupRootOfUnity);
+
+    // Compute the best size of our evaluation domain.
+    size_ = BestMixedDomainSize(num_coeffs);
+
+    // Compute the size of our evaluation domain
+    uint64_t q = static_cast<uint64_t>(Field::Config::kSmallSubgroupBase);
+    uint32_t q_adicity =
+        ComputeAdicity(q, gmp::FromDecString(base::NumberToString(size_)));
+    uint64_t q_part = static_cast<uint64_t>(std::pow(q, q_adicity));
+
+    uint32_t two_adicity =
+        ComputeAdicity(2, gmp::FromDecString(base::NumberToString(size_)));
+    log_size_of_group_ = two_adicity;
+    uint64_t two_part = static_cast<uint64_t>(std::pow(2, two_adicity));
+    CHECK(size_ == q_part * two_part);
+    size_as_field_element_ = Field::FromDecString(base::NumberToString(size_));
+    size_inv_ = size_as_field_element_.Inverse();
+
+    // Compute the generator for the multiplicative subgroup.
+    // It should be the |num_coeffs| root of unity.
+    Field::GetRootOfUnity(size_, &group_gen_);
+    // Check that it is indeed the requested root of unity.
+    DCHECK(group_gen_.Pow(BigInt<1>(size_)) == Field::One());
+    group_gen_inv_ = group_gen_.Inverse();
+    offset_ = Field::One();
+    offset_inv_ = Field::One();
+    offset_pow_size = Field::One();
+  }
+
+  constexpr MixedRadixEvaluationDomain<Field> GetCoset(Field offset) {
+    MixedRadixEvaluationDomain<Field> coset;
+    coset.size_ = size_;
+    coset.log_size_of_group_ = log_size_of_group_;
+    coset.size_as_field_element_ = size_as_field_element_;
+    coset.size_inv_ = size_inv_;
+    coset.group_gen_ = group_gen_;
+    coset.group_gen_inv_ = group_gen_inv_;
+    coset.offset_ = offset;
+    coset.offset_inv_ = offset.Inverse();
+    coset.offset_pow_size = offset.Pow(BigInt<1>(size_));
+    return coset;
+  }
+
+  constexpr static bool ComputeSizeOfDomain(uint64_t num_coeffs,
+                                            uint64_t* size) {
+    if (F::Config::kHasLargeSubgroupRootOfUnity) {
+      // Compute the best size of our evaluation domain.
+      *size = BestMixedDomainSize(num_coeffs);
+
+      uint64_t q = static_cast<uint64_t>(Field::Config::kSmallSubgroupBase);
+      uint32_t q_adicity =
+          ComputeAdicity(q, gmp::FromDecString(base::NumberToString(*size)));
+      uint64_t q_part = static_cast<uint64_t>(std::pow(q, q_adicity));
+
+      uint32_t two_adicity =
+          ComputeAdicity(2, gmp::FromDecString(base::NumberToString(*size)));
+      uint64_t two_part = static_cast<uint64_t>(std::pow(2, two_adicity));
+
+      if (*size != q_part * two_part) {
+        LOG(ERROR) << "Cannot compute a mixed radix domain of size " << *size
+                   << " for a polynomial of degree " << num_coeffs
+                   << " (best size is " << q_part * two_part << ")";
+        return false;
+      }
+      return true;
+    } else {
+      LOG(ERROR) << "Small subgroup base must be defined for the field to use "
+                 << "a mixed radix domain.";
+      return false;
+    }
+  }
+
+  constexpr uint64_t GetSize() const { return size_; }
+
+  constexpr uint32_t GetLogSizeOfGroup() const { return log_size_of_group_; }
+
+  constexpr Field GetSizeAsFieldElement() const {
+    return size_as_field_element_;
+  }
+
+  constexpr Field GetGroupGen() const { return group_gen_; }
+
+  constexpr Field GetGroupGenInv() const { return group_gen_inv_; }
+
+  constexpr Field GetCosetOffset() const { return offset_; }
+
+  constexpr Field GetCosetOffsetInv() const { return offset_inv_; }
+
+  constexpr Field GetCosetOffsetPowSize() const { return offset_pow_size; }
+
+  template <typename DensePoly>
+  constexpr void FFTInPlace(DensePoly& poly) {
+    if (!offset_.IsOne()) {
+      EvaluationDomain<MixedRadixEvaluationDomain<Field>>::DistributePowers(
+          poly, offset_);
+    }
+    poly.coefficients_.coefficients_.resize(size_, Field::Zero());
+    SerialMixedRadixFFT(poly, group_gen_, log_size_of_group_);
+  }
+
+  template <typename DensePoly>
+  constexpr void IFFTInPlace(DensePoly& eval) {
+    eval.coefficients_.coefficients_.resize(size_, Field::Zero());
+    SerialMixedRadixFFT(eval, group_gen_inv_, log_size_of_group_);
+    if (offset_.IsOne()) {
+      for (Field& coeff : eval.coefficients_.coefficients_) {
+        coeff *= size_inv_;
+      }
+    } else {
+      EvaluationDomain<MixedRadixEvaluationDomain<Field>>::
+          DistributePowersAndMulByConst(eval, offset_inv_, size_inv_);
+    }
+  }
+
+ private:
+  friend class EvaluationDomain<MixedRadixEvaluationDomain<Field>>;
+
+  constexpr size_t MixedRadixFFTPermute(uint32_t two_adicity,
+                                        uint32_t q_adicity, size_t q, size_t n,
+                                        size_t i) {
+    // This is the permutation obtained by splitting into 2 groups |two_adicity|
+    // times and then |q| groups |q_adicity| many times. It can be efficiently
+    // described as follows i = 2^0 b_0 + 2^1 b_1 + ... + 2^{|two_adicity| - 1}
+    // b_{|two_adicity| - 1} + 2^|two_adicity| ( x_0 + |q|^1 x_1 + .. +
+    // |q|^{|q_adicity|-1} x_{|q_adicity|-1}). We want to return
+    // j = b_0 (|n|/2) + b_1 (|n|/ 2^2) + ... + b_{|two_adicity|-1}
+    // (|n|/2^|two_adicity|) + x_0 (|n| / 2^|two_adicity| / |q|) + .. +
+    // x_{|q_adicity|-1} (|n| / 2^|two_adicity| / |q|^|q_adicity|).
+    size_t res = 0;
+    size_t shift = n;
+    for (size_t j = 0; j < two_adicity; ++j) {
+      shift /= 2;
+      res += (i % 2) * shift;
+      i /= 2;
+    }
+    for (size_t j = 0; j < q_adicity; ++j) {
+      shift /= q;
+      res += (i % q) * shift;
+      i /= q;
+    }
+    return res;
+  }
+
+  constexpr static uint64_t BestMixedDomainSize(uint64_t min_size) {
+    uint64_t best = UINT64_MAX;
+    uint32_t small_subgroup_base = Field::Config::kSmallSubgroupBase;
+    uint32_t small_subgroup_adicity = Field::Config::kSmallSubgroupAdicity;
+    for (size_t i = 0; i <= small_subgroup_adicity; ++i) {
+      uint64_t r = static_cast<uint64_t>(std::pow(small_subgroup_base, i));
+      uint32_t two_adicity = 0;
+      while (r < min_size) {
+        r *= 2;
+        two_adicity += 1;
+      }
+      if (two_adicity <= Field::Config::kTwoAdicity) {
+        best = std::min(best, r);
+      }
+    }
+    return best;
+  }
+
+  // TODO(TomTaehoonKim): Implement ParallelFFT and BestFFT. (See
+  // https://github.com/arkworks-rs/algebra/blob/master/poly/src/domain/utils.rs#L105)
+  template <typename DensePoly>
+  void SerialMixedRadixFFT(DensePoly& a, Field omega, uint32_t two_adicity) {
+    // Conceptually, this FFT first splits into 2 sub-arrays |two_adicity| many
+    // times, and then splits into q sub-arrays |q_adicity| many times.
+
+    size_t n = a.coefficients_.coefficients_.size();
+    uint32_t q = Field::Config::kSmallSubgroupBase;
+    uint64_t q_u64 = static_cast<uint64_t>(Field::Config::kSmallSubgroupBase);
+    uint64_t n_u64 = static_cast<uint64_t>(n);
+
+    uint32_t q_adicity =
+        ComputeAdicity(q_u64, gmp::FromDecString(base::NumberToString(n_u64)));
+    uint64_t q_part = static_cast<uint64_t>(std::pow(q_u64, q_adicity));
+    uint64_t two_part =
+        static_cast<uint64_t>(std::pow(uint64_t(2), two_adicity));
+
+    CHECK(n_u64 == q_part * two_part);
+
+    size_t m = 1;
+    if (q_adicity > 0) {
+      // If we're using the other radix, we have to do two things differently
+      // than in the radix 2 case. 1. Applying the index permutation is a bit
+      // more complicated. It isn't an involution (like it is in the radix 2
+      // case) so we need to remember which elements we've moved as we go along
+      // and can't use the trick of just swapping when processing the first
+      // element of a 2-cycle.
+      //
+      // 2. We need to do |q_adicity| many merge passes, each of which is a bit
+      // more complicated than the specialized q=2 case.
+
+      // Applying the permutation
+      std::vector<bool> seen(n, false);
+      for (size_t k = 0; k < n; ++k) {
+        size_t i = k;
+        Field& a_i = a.coefficients_.coefficients_[i];
+        while (!seen[i]) {
+          size_t dest = MixedRadixFFTPermute(two_adicity, q_adicity, q, n, i);
+
+          Field a_dest = a.coefficients_.coefficients_[dest];
+          a.coefficients_.coefficients_[dest] = a_i;
+
+          seen[i] = true;
+
+          a_i = a_dest;
+          i = dest;
+        }
+      }
+
+      Field omega_q = omega.Pow(BigInt<1>(n / q));
+      std::vector<Field> qth_roots(q, Field::One());
+      for (size_t i = 1; i < q; ++i) {
+        qth_roots[i] = qth_roots[i - 1] * omega_q;
+      }
+
+      std::vector<Field> terms(q - 1, Field::Zero());
+
+      // Doing the q_adicity passes.
+      for (size_t _i = 0; _i < q_adicity; ++_i) {
+        Field w_m = omega.Pow(BigInt<1>(n / (q * m)));
+        size_t k = 0;
+        while (k < n) {
+          Field w_j = Field::One();  // w_j is omega_m ^ j
+          for (size_t j = 0; j < m; ++j) {
+            Field base_term = a.coefficients_.coefficients_[k + j];
+            Field w_j_i = w_j;
+            for (size_t i = 1; i < q; ++i) {
+              terms[i - 1] = a.coefficients_.coefficients_[k + j + i * m];
+              terms[i - 1] *= w_j_i;
+              w_j_i *= w_j;
+            }
+
+            for (size_t i = 0; i < q; ++i) {
+              a.coefficients_.coefficients_[k + j + i * m] = base_term;
+              for (size_t l = 1; l < q; ++l) {
+                Field tmp = terms[l - 1];
+                tmp *= qth_roots[(i * l) % q];
+                a.coefficients_.coefficients_[k + j + i * m] += tmp;
+              }
+            }
+
+            w_j *= w_m;
+          }
+
+          k += q * m;
+        }
+        m *= q;
+      }
+    } else {
+      // Swapping in place (from Storer's book)
+      for (uint64_t k = 0; k < n; ++k) {
+        uint64_t rk = BitRev(k, two_adicity);
+        if (k < rk) {
+          std::swap(a.coefficients_.coefficients_[k],
+                    a.coefficients_.coefficients_[rk]);
+        }
+      }
+    }
+
+    for (size_t _i = 0; _i < two_adicity; ++_i) {
+      // w_m is 2^|two_adicity|-th root of unity now
+      Field w_m = omega.Pow(BigInt<1>(n / (2 * m)));
+
+      size_t k = 0;
+      while (k < n) {
+        Field w = Field::One();
+        for (size_t j = 0; j < m; ++j) {
+          Field t = a.coefficients_.coefficients_[(k + m) + j];
+          t *= w;
+          a.coefficients_.coefficients_[(k + m) + j] =
+              a.coefficients_.coefficients_[k + j];
+          a.coefficients_.coefficients_[(k + m) + j] -= t;
+          a.coefficients_.coefficients_[k + j] += t;
+          w *= w_m;
+        }
+        k += 2 * m;
+      }
+      m *= 2;
+    }
+  }
+
+  // The size of the domain.
+  uint64_t size_;
+  // log2(|size_|).
+  uint32_t log_size_of_group_;
+  // Size of the domain as a field element.
+  Field size_as_field_element_;
+  // Inverse of the size in the field.
+  Field size_inv_;
+  // A generator of the subgroup.
+  Field group_gen_;
+  // Inverse of the generator of the subgroup.
+  Field group_gen_inv_;
+  // Offset that specifies the coset.
+  Field offset_;
+  // Inverse of the offset that specifies the coset.
+  Field offset_inv_;
+  // Constant coefficient for the vanishing polynomial.
+  // Equals |offset_|^|size_|.
+  Field offset_pow_size;
+};
+
+template <typename F>
+class EvaluationDomainTraits<MixedRadixEvaluationDomain<F>> {
+ public:
+  using Field = F;
+};
+
+}  // namespace tachyon::math
+
+#endif  // TACHYON_MATH_POLYNOMIAL_DOMAINS_MIXED_RADIX_MIXED_RADIX_EVALUATION_DOMAIN_H_