feat(math): introduce GeneralEvaluationDomain

tachyon/math/polynomial/domains/BUILD.bazel

@@ -1,4 +1,4 @@
-load("//bazel:tachyon_cc.bzl", "tachyon_cc_library")
+load("//bazel:tachyon_cc.bzl", "tachyon_cc_library", "tachyon_cc_test")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -20,3 +20,28 @@ tachyon_cc_library(
         "//tachyon/base:compiler_specific",
     ],
 )
+
+tachyon_cc_library(
+    name = "general_evaluation_domain",
+    hdrs = ["general_evaluation_domain.h"],
+    deps = [
+        ":evaluation_domain",
+        "//tachyon/base:logging",
+        "//tachyon/math/polynomial/domains/mixed_radix:mixed_radix_evaluation_domain",
+        "//tachyon/math/polynomial/domains/radix2:radix2_evaluation_domain",
+        "//tachyon/math/polynomial/polynomials/univariate:univariate_polynomial",
+    ],
+)
+
+tachyon_cc_test(
+    name = "general_evaluation_domain_unittest",
+    size = "small",
+    srcs = ["general_evaluation_domain_unittest.cc"],
+    deps = [
+        ":general_evaluation_domain",
+        "//tachyon/math/elliptic_curves/bls/bls12_381:fr",
+        "//tachyon/math/elliptic_curves/bls/bls12_381:g1",
+        "//tachyon/math/elliptic_curves/bn/bn384_small_two_adicity:fr",
+        "//tachyon/math/polynomial/polynomials/univariate:univariate_polynomial",
+    ],
+)

tachyon/math/polynomial/domains/general_evaluation_domain.h

@@ -0,0 +1,142 @@
+// 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 |GeneralEvaluationDomain| for performing various kinds
+// of polynomial arithmetic on top of a FFT-friendly finite field.
+//
+// It is a wrapper around specific implementations of |EvaluationDomain| that
+// automatically chooses the most efficient implementation depending on the
+// number of coefficients and the two-adicity of the prime.
+
+#ifndef TACHYON_MATH_POLYNOMIAL_DOMAINS_GENERAL_EVALUATION_DOMAIN_H_
+#define TACHYON_MATH_POLYNOMIAL_DOMAINS_GENERAL_EVALUATION_DOMAIN_H_
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include "tachyon/base/logging.h"
+#include "tachyon/math/polynomial/domains/evaluation_domain.h"
+#include "tachyon/math/polynomial/domains/mixed_radix/mixed_radix_evaluation_domain.h"
+#include "tachyon/math/polynomial/domains/radix2/radix2_evaluation_domain.h"
+
+namespace tachyon::math {
+
+// Defines a domain over which finite field (I)FFTs can be performed.
+// Generally tries to build a radix-2 domain and falls back to a mixed-radix
+// domain if the radix-2 multiplicative subgroup is too small.
+// pub enum GeneralEvaluationDomain<F: FftField> {
+//     /// Radix-2 domain
+//     Radix2(Radix2EvaluationDomain<F>),
+//     /// Mixed-radix domain
+//     MixedRadix(MixedRadixEvaluationDomain<F>),
+// }
+template <typename F>
+class GeneralEvaluationDomain
+    : public EvaluationDomain<GeneralEvaluationDomain<F>> {
+ public:
+  using DomainType =
+      std::variant<Radix2EvaluationDomain<F>, MixedRadixEvaluationDomain<F>>;
+  constexpr GeneralEvaluationDomain() = default;
+  // Construct a domain that is large enough for evaluations of a polynomial
+  // having |num_coeffs| coefficients.
+  //
+  // If the field specifies a small subgroup for a mixed-radix FFT and the
+  // radix-2 FFT cannot be constructed, this method tries constructing a
+  // mixed-radix FFT instead.
+  constexpr explicit GeneralEvaluationDomain(size_t num_coeffs) {
+    uint64_t size = 0;
+    if (Radix2EvaluationDomain<F>::ComputeSizeOfDomain(num_coeffs, &size)) {
+      domain_ = Radix2EvaluationDomain<F>(num_coeffs);
+    } else if (MixedRadixEvaluationDomain<F>::ComputeSizeOfDomain(num_coeffs,
+                                                                  &size)) {
+      domain_ = MixedRadixEvaluationDomain<F>(num_coeffs);
+    } else {
+      LOG(ERROR) << "Cannot construct a domain for the given prime field and "
+                    "the number of coefficients.";
+    }
+  }
+
+  constexpr GeneralEvaluationDomain<F> GetCoset(F offset) {
+    GeneralEvaluationDomain<F> coset;
+    coset.domain_ = std::visit(
+        [offset](auto&& domain) -> DomainType {
+          return domain.GetCoset(offset);
+        },
+        domain_);
+    return coset;
+  }
+
+  constexpr bool ComputeSizeOfDomain(size_t num_coeffs, uint64_t* size) {
+    return std::visit(
+        [num_coeffs, size](auto&& domain) {
+          return domain.ComputeSizeOfDomain(num_coeffs, size);
+        },
+        domain_);
+  }
+
+  constexpr uint64_t GetSize() const {
+    return std::visit([](auto&& domain) { return domain.GetSize(); }, domain_);
+  }
+
+  constexpr uint32_t GetLogSizeOfGroup() const {
+    return std::visit([](auto&& domain) { return domain.GetLogSizeOfGroup(); },
+                      domain_);
+  }
+
+  constexpr F GetSizeAsFieldElement() const {
+    return std::visit(
+        [](auto&& domain) { return domain.GetSizeAsFieldElement(); }, domain_);
+  }
+
+  constexpr F GetGroupGen() const {
+    return std::visit([](auto&& domain) { return domain.GetGroupGen(); },
+                      domain_);
+  }
+
+  constexpr F GetGroupGenInv() const {
+    return std::visit([](auto&& domain) { return domain.GetGroupGenInv(); },
+                      domain_);
+  }
+
+  constexpr F GetCosetOffset() const {
+    return std::visit([](auto&& domain) { return domain.GetCosetOffset(); },
+                      domain_);
+  }
+
+  constexpr F GetCosetOffsetInv() const {
+    return std::visit([](auto&& domain) { return domain.GetCosetOffsetInv(); },
+                      domain_);
+  }
+
+  constexpr F GetCosetOffsetPowSize() const {
+    return std::visit(
+        [](auto&& domain) { return domain.GetCosetOffsetPowSize(); }, domain_);
+  }
+
+  template <typename DensePoly>
+  constexpr void FFTInPlace(DensePoly& poly) {
+    return std::visit([&poly](auto&& domain) { domain.FFTInPlace(poly); },
+                      domain_);
+  }
+
+  template <typename DensePoly>
+  constexpr void IFFTInPlace(DensePoly& eval) {
+    return std::visit([&eval](auto&& domain) { domain.IFFTInPlace(eval); },
+                      domain_);
+  }
+
+ private:
+  DomainType domain_;
+};
+
+template <typename F>
+class EvaluationDomainTraits<GeneralEvaluationDomain<F>> {
+ public:
+  using Field = F;
+};
+
+}  // namespace tachyon::math
+
+#endif  // TACHYON_MATH_POLYNOMIAL_DOMAINS_GENERAL_EVALUATION_DOMAIN_H_

tachyon/math/polynomial/domains/general_evaluation_domain_unittest.cc

@@ -0,0 +1,103 @@
+// 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.
+
+#include "tachyon/math/polynomial/domains/general_evaluation_domain.h"
+
+#include "gtest/gtest.h"
+
+#include "tachyon/math/elliptic_curves/bls/bls12_381/fr.h"
+#include "tachyon/math/elliptic_curves/bls/bls12_381/g1.h"
+#include "tachyon/math/elliptic_curves/bn/bn384_small_two_adicity/fr.h"
+#include "tachyon/math/polynomial/polynomials/univariate/univariate_polynomial.h"
+
+namespace tachyon::math {
+
+namespace {
+
+class GeneralEvaluationDomainTest : public testing::Test {
+ public:
+  static void SetUpTestSuite() {
+    bls12_381::Fr::Init();
+    bn384_small_two_adicity::Fr::Init();
+  }
+
+  GeneralEvaluationDomainTest() = default;
+  GeneralEvaluationDomainTest(const GeneralEvaluationDomainTest&) = delete;
+  GeneralEvaluationDomainTest& operator=(const GeneralEvaluationDomainTest&) =
+      delete;
+  ~GeneralEvaluationDomainTest() override = default;
+};
+
+}  // namespace
+
+TEST_F(GeneralEvaluationDomainTest, VanishingPolynomialEvaluation) {
+  for (size_t coeffs = 1; coeffs < 10; ++coeffs) {
+    GeneralEvaluationDomain<bls12_381::Fr> domain(coeffs);
+    auto z = domain.GetVanishingPolynomial();
+    for (size_t _i = 0; _i < 100; ++_i) {
+      bls12_381::Fr point = bls12_381::Fr::Random();
+      EXPECT_EQ(z.Evaluate(point), domain.EvaluateVanishingPolynomial(point));
+    }
+  }
+
+  for (size_t coeffs = 15; coeffs < 17; ++coeffs) {
+    GeneralEvaluationDomain<bn384_small_two_adicity::Fr> domain(coeffs);
+    auto z = domain.GetVanishingPolynomial();
+    for (size_t _i = 0; _i < 100; ++_i) {
+      bn384_small_two_adicity::Fr point = bn384_small_two_adicity::Fr::Random();
+      EXPECT_EQ(z.Evaluate(point), domain.EvaluateVanishingPolynomial(point));
+    }
+  }
+}
+
+TEST_F(GeneralEvaluationDomainTest, VanishingPolynomialVanishesOnDomain) {
+  // NOTE(TomTaehoonKim): In arkworks, the test is done for 0..1000, but it
+  // takes ~90s with size = "medium" in Tachyon, so we only test for 0..10.
+  for (size_t coeffs = 0; coeffs < 10; ++coeffs) {
+    GeneralEvaluationDomain<bls12_381::Fr> domain(coeffs);
+    auto z = domain.GetVanishingPolynomial();
+    for (bls12_381::Fr& element : domain.GetElements()) {
+      EXPECT_TRUE(z.Evaluate(element).IsZero());
+    }
+  }
+}
+
+TEST_F(GeneralEvaluationDomainTest, SizeOfElements) {
+  for (size_t coeffs = 1; coeffs < 10; ++coeffs) {
+    size_t size = 1 << coeffs;
+    GeneralEvaluationDomain<bls12_381::Fr> domain(size);
+    EXPECT_EQ(domain.GetSize(), domain.GetElements().size());
+  }
+}
+
+TEST_F(GeneralEvaluationDomainTest, FFTComposition) {
+  const size_t log_degree = 5;
+  const size_t degree = (1 << log_degree) - 1;
+  GeneralEvaluationDomain<bls12_381::Fr> domain(degree);
+  GeneralEvaluationDomain<bls12_381::Fr> coset_domain = domain.GetCoset(
+      bls12_381::Fr::FromMontgomery(bls12_381::Fr::Config::kSubgroupGenerator));
+
+  DenseUnivariatePolynomial<bls12_381::Fr, degree> v =
+      DenseUnivariatePolynomial<bls12_381::Fr, degree>::Random(degree);
+  DenseUnivariatePolynomial<bls12_381::Fr, degree> v2 = v;
+
+  domain.IFFTInPlace(v2);
+  domain.FFTInPlace(v2);
+  EXPECT_EQ(v.ToString(), v2.ToString());
+
+  domain.FFTInPlace(v2);
+  domain.IFFTInPlace(v2);
+  EXPECT_EQ(v.ToString(), v2.ToString());
+
+  coset_domain.IFFTInPlace(v2);
+  coset_domain.FFTInPlace(v2);
+  EXPECT_EQ(v.ToString(), v2.ToString());
+
+  coset_domain.FFTInPlace(v2);
+  coset_domain.IFFTInPlace(v2);
+  EXPECT_EQ(v.ToString(), v2.ToString());
+}
+
+}  // namespace tachyon::math

tachyon/math/polynomial/polynomials/univariate/dense_coefficients.h

@@ -21,6 +21,9 @@ class SparseCoefficients;
 template <typename Derived>
 class EvaluationDomain;
 
+template <typename F>
+class GeneralEvaluationDomain;
+
 template <typename F>
 class Radix2EvaluationDomain;
 
@@ -132,8 +135,10 @@ class DenseCoefficients {
       DenseCoefficients<F, MaxDegree>>;
   friend class internal::UnivariatePolynomialOp<
       SparseCoefficients<F, MaxDegree>>;
+  friend class EvaluationDomain<GeneralEvaluationDomain<F>>;
   friend class EvaluationDomain<Radix2EvaluationDomain<F>>;
   friend class EvaluationDomain<MixedRadixEvaluationDomain<F>>;
+  friend class GeneralEvaluationDomain<F>;
   friend class Radix2EvaluationDomain<F>;
   friend class MixedRadixEvaluationDomain<F>;
 

tachyon/math/polynomial/polynomials/univariate/univariate_polynomial.h

@@ -156,8 +156,12 @@ class UnivariatePolynomial
 
  private:
   friend class internal::UnivariatePolynomialOp<Coefficients>;
+  // TODO(TomTaehoonKim): Enable friend class declaration over other types of
+  // fields than |Coefficients::Field|.
+  friend class EvaluationDomain<GeneralEvaluationDomain<Field>>;
   friend class EvaluationDomain<Radix2EvaluationDomain<Field>>;
   friend class EvaluationDomain<MixedRadixEvaluationDomain<Field>>;
+  friend class GeneralEvaluationDomain<Field>;
   friend class Radix2EvaluationDomain<Field>;
   friend class MixedRadixEvaluationDomain<Field>;