From bcd48d25e81c34edf8c53e8e3a33347fb197b0e4 Mon Sep 17 00:00:00 2001 From: TomTaehoonKim Date: Wed, 4 Oct 2023 01:43:18 +0900 Subject: [PATCH] feat(math): introduce `GeneralEvaluationDomain` --- tachyon/math/polynomial/domains/BUILD.bazel | 27 +++- .../domains/general_evaluation_domain.h | 142 ++++++++++++++++++ .../general_evaluation_domain_unittest.cc | 103 +++++++++++++ .../univariate/dense_coefficients.h | 5 + .../univariate/univariate_polynomial.h | 4 + 5 files changed, 280 insertions(+), 1 deletion(-) create mode 100644 tachyon/math/polynomial/domains/general_evaluation_domain.h create mode 100644 tachyon/math/polynomial/domains/general_evaluation_domain_unittest.cc diff --git a/tachyon/math/polynomial/domains/BUILD.bazel b/tachyon/math/polynomial/domains/BUILD.bazel index 4ba4addefa..a4a2db4fdc 100644 --- a/tachyon/math/polynomial/domains/BUILD.bazel +++ b/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", + ], +) diff --git a/tachyon/math/polynomial/domains/general_evaluation_domain.h b/tachyon/math/polynomial/domains/general_evaluation_domain.h new file mode 100644 index 0000000000..524bfa3cbd --- /dev/null +++ b/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 +#include + +#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 { +// /// Radix-2 domain +// Radix2(Radix2EvaluationDomain), +// /// Mixed-radix domain +// MixedRadix(MixedRadixEvaluationDomain), +// } +template +class GeneralEvaluationDomain + : public EvaluationDomain> { + public: + using DomainType = + std::variant, MixedRadixEvaluationDomain>; + 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::ComputeSizeOfDomain(num_coeffs, &size)) { + domain_ = Radix2EvaluationDomain(num_coeffs); + } else if (MixedRadixEvaluationDomain::ComputeSizeOfDomain(num_coeffs, + &size)) { + domain_ = MixedRadixEvaluationDomain(num_coeffs); + } else { + LOG(ERROR) << "Cannot construct a domain for the given prime field and " + "the number of coefficients."; + } + } + + constexpr GeneralEvaluationDomain GetCoset(F offset) { + GeneralEvaluationDomain 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 + constexpr void FFTInPlace(DensePoly& poly) { + return std::visit([&poly](auto&& domain) { domain.FFTInPlace(poly); }, + domain_); + } + + template + constexpr void IFFTInPlace(DensePoly& eval) { + return std::visit([&eval](auto&& domain) { domain.IFFTInPlace(eval); }, + domain_); + } + + private: + DomainType domain_; +}; + +template +class EvaluationDomainTraits> { + public: + using Field = F; +}; + +} // namespace tachyon::math + +#endif // TACHYON_MATH_POLYNOMIAL_DOMAINS_GENERAL_EVALUATION_DOMAIN_H_ diff --git a/tachyon/math/polynomial/domains/general_evaluation_domain_unittest.cc b/tachyon/math/polynomial/domains/general_evaluation_domain_unittest.cc new file mode 100644 index 0000000000..734f805863 --- /dev/null +++ b/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 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 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 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 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 domain(degree); + GeneralEvaluationDomain coset_domain = domain.GetCoset( + bls12_381::Fr::FromMontgomery(bls12_381::Fr::Config::kSubgroupGenerator)); + + DenseUnivariatePolynomial v = + DenseUnivariatePolynomial::Random(degree); + DenseUnivariatePolynomial 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 diff --git a/tachyon/math/polynomial/polynomials/univariate/dense_coefficients.h b/tachyon/math/polynomial/polynomials/univariate/dense_coefficients.h index cefc1d79f9..0d6dd175fa 100644 --- a/tachyon/math/polynomial/polynomials/univariate/dense_coefficients.h +++ b/tachyon/math/polynomial/polynomials/univariate/dense_coefficients.h @@ -21,6 +21,9 @@ class SparseCoefficients; template class EvaluationDomain; +template +class GeneralEvaluationDomain; + template class Radix2EvaluationDomain; @@ -132,8 +135,10 @@ class DenseCoefficients { DenseCoefficients>; friend class internal::UnivariatePolynomialOp< SparseCoefficients>; + friend class EvaluationDomain>; friend class EvaluationDomain>; friend class EvaluationDomain>; + friend class GeneralEvaluationDomain; friend class Radix2EvaluationDomain; friend class MixedRadixEvaluationDomain; diff --git a/tachyon/math/polynomial/polynomials/univariate/univariate_polynomial.h b/tachyon/math/polynomial/polynomials/univariate/univariate_polynomial.h index 191fb8f31a..1deda212e4 100644 --- a/tachyon/math/polynomial/polynomials/univariate/univariate_polynomial.h +++ b/tachyon/math/polynomial/polynomials/univariate/univariate_polynomial.h @@ -156,8 +156,12 @@ class UnivariatePolynomial private: friend class internal::UnivariatePolynomialOp; + // TODO(TomTaehoonKim): Enable friend class declaration over other types of + // fields than |Coefficients::Field|. + friend class EvaluationDomain>; friend class EvaluationDomain>; friend class EvaluationDomain>; + friend class GeneralEvaluationDomain; friend class Radix2EvaluationDomain; friend class MixedRadixEvaluationDomain;