diff --git a/tachyon/math/polynomial/domains/BUILD.bazel b/tachyon/math/polynomial/domains/BUILD.bazel index 4ba4addefa..08ed415a9d 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_unittests", + 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/evaluation_domain.h b/tachyon/math/polynomial/domains/evaluation_domain.h index 327c4c136b..a3be646601 100644 --- a/tachyon/math/polynomial/domains/evaluation_domain.h +++ b/tachyon/math/polynomial/domains/evaluation_domain.h @@ -45,23 +45,23 @@ class EvaluationDomain { return poly; } - // Multiply the i-th element of |poly| with |g|^i. + // Multiply the i-th element of |poly| with |f|^i. template - constexpr void DistributePowers(DensePoly& poly, const Field& g) { - DistributePowersAndMulByConst(poly, g, Field::One()); + constexpr void DistributePowers(DensePoly& poly, const Field& f) { + DistributePowersAndMulByConst(poly, f, Field::One()); } - // Multiply the i-th element of |poly| with |c|*|g|^i. + // Multiply the i-th element of |poly| with |c|*|f|^i. // TODO(TomTaehoonKim): Parallelize this. (See // https://github.com/arkworks-rs/algebra/blob/4152c41769ae0178fc110bfd15cc699673a2ce4b/poly/src/domain/mod.rs#L130) template - constexpr void DistributePowersAndMulByConst(DensePoly& poly, const Field& g, + constexpr void DistributePowersAndMulByConst(DensePoly& poly, const Field& f, const Field& c) { - // Invariant: |pow| = |c|*|g|^i at the i-th iteration of the loop + // Invariant: |pow| = |c|*|f|^i at the i-th iteration of the loop Field pow = c; for (Field& coeff : poly.coefficients_.coefficients_) { coeff *= pow; - pow *= g; + pow *= f; } } 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..fb9cde2d01 --- /dev/null +++ b/tachyon/math/polynomial/domains/general_evaluation_domain.h @@ -0,0 +1,144 @@ +// 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 + +#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 dbbad2b427..d5b23b4731 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 89b317cb4a..38f7b3bed9 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;