diff --git a/tachyon/math/polynomials/univariate/BUILD.bazel b/tachyon/math/polynomials/univariate/BUILD.bazel index ddc5c4abde..bd7d71e30e 100644 --- a/tachyon/math/polynomials/univariate/BUILD.bazel +++ b/tachyon/math/polynomials/univariate/BUILD.bazel @@ -130,6 +130,7 @@ tachyon_cc_unittest( name = "univariate_unittests", srcs = [ "lagrange_interpolation_unittest.cc", + "radix2_dit_parallel_unittest.cc", "univariate_dense_polynomial_unittest.cc", "univariate_evaluation_domain_unittest.cc", "univariate_evaluations_unittest.cc", @@ -138,6 +139,7 @@ tachyon_cc_unittest( deps = [ ":lagrange_interpolation", ":mixed_radix_evaluation_domain", + ":radix2_dit_parallel", ":radix2_evaluation_domain", ":univariate_polynomial", "//tachyon/base:optional", @@ -148,6 +150,7 @@ tachyon_cc_unittest( "//tachyon/math/elliptic_curves/bls12/bls12_381:fr", "//tachyon/math/elliptic_curves/bn/bn254:fr", "//tachyon/math/elliptic_curves/bn/bn384_small_two_adicity:fq", + "//tachyon/math/finite_fields/baby_bear", "//tachyon/math/finite_fields/test:finite_field_test", "//tachyon/math/finite_fields/test:gf7", "@com_google_absl//absl/hash:hash_testing", diff --git a/tachyon/math/polynomials/univariate/radix2_dit_parallel_unittest.cc b/tachyon/math/polynomials/univariate/radix2_dit_parallel_unittest.cc new file mode 100644 index 0000000000..c123e02242 --- /dev/null +++ b/tachyon/math/polynomials/univariate/radix2_dit_parallel_unittest.cc @@ -0,0 +1,109 @@ +#include "tachyon/math/polynomials/univariate/radix2_dit_parallel.h" + +#include "gtest/gtest.h" + +#include "tachyon/math/finite_fields/baby_bear/baby_bear.h" +#include "tachyon/math/matrix/matrix_types.h" + +namespace Eigen { + +template +void PrintTo(const tachyon::math::RowMajorMatrix &m, std::ostream *os) { + size_t rows = static_cast(m.rows()); + size_t cols = static_cast(m.cols()); + for (size_t row = 0; row < rows; ++row) { + for (size_t col = 0; col < cols; ++col) { + *os << m(row, col).ToString() << ","; + } + *os << std::endl; + } +} + +} // namespace Eigen + +namespace tachyon::math { + +// Values are pulled from results in Plonky3 +TEST(Radix2DitParallelTest, DFTBatch) { + using F = BabyBear; + // Small case + RowMajorMatrix small_mat(1, 3); + small_mat << F(0), F(1), F(2); + Radix2DitParallel::DFTBatch(small_mat); + RowMajorMatrix kSmallDFTExpected(1, 3); + kSmallDFTExpected << F(0), F(1), F(2); + EXPECT_EQ(kSmallDFTExpected, small_mat); + // Larger case + RowMajorMatrix large_mat(1 << 3, 3); + // clang-format off + large_mat << F(0), F(1), F(2), + F(3), F(4), F(5), + F(6), F(7), F(8), + F(9), F(10), F(11), + F(12), F(13), F(14), + F(15), F(16), F(17), + F(18), F(19), F(20), + F(21), F(22), F(23); + // clang-format on + Radix2DitParallel::DFTBatch(large_mat); + RowMajorMatrix kLargeDFTExpected(1 << 3, 3); + // clang-format off + kLargeDFTExpected << F(84), F(92), F(100), + F(1901923198), F(1901923198), F(1901923198), + F(1405070963), F(1405070963), F(1405070963), + F(1105047169), F(1105047169), F(1105047169), + F(2013265909), F(2013265909), F(2013265909), + F(908218728), F(908218728), F(908218728), + F(608194934), F(608194934), F(608194934), + F(111342699), F(111342699), F(111342699); + // clang-format on + EXPECT_EQ(kLargeDFTExpected, large_mat); +} + +TEST(Radix2DitParallelTest, CosetLDEBatch) { + using F = BabyBear; + // Small case + RowMajorMatrix small_mat(1, 3); + small_mat << F(0), F(1), F(2); + F shift = F(31); + Radix2DitParallel::CosetLDEBatch(small_mat, 1, shift); + RowMajorMatrix kSmallCosetLDEExpected(1 << 1, 3); + kSmallCosetLDEExpected << F(0), F(1), F(2), F(0), F(1), F(2); + EXPECT_EQ(kSmallCosetLDEExpected, small_mat); + + // Larger case + RowMajorMatrix large_mat(1 << 3, 3); + // clang-format off + large_mat << F(0), F(1), F(2), + F(3), F(4), F(5), + F(6), F(7), F(8), + F(9), F(10), F(11), + F(12), F(13), F(14), + F(15), F(16), F(17), + F(18), F(19), F(20), + F(21), F(22), F(23); + // clang-format on + shift = F(31); + Radix2DitParallel::CosetLDEBatch(large_mat, 1, shift); + // clang-format off + RowMajorMatrix kLargeCosetLDEExpected(1 << 4, 3); + kLargeCosetLDEExpected << F(1754655111), F(1754655112), F(1754655113), + F(635057597), F(635057598), F(635057599), + F(1554846415), F(1554846416), F(1554846417), + F(1945388118), F(1945388119), F(1945388120), + F(1795728790), F(1795728791), F(1795728792), + F(299008665), F(299008666), F(299008667), + F(1099398705), F(1099398706), F(1099398707), + F(131814472), F(131814473), F(131814474), + F(527296628), F(527296629), F(527296630), + F(1553002944), F(1553002945), F(1553002946), + F(904503641), F(904503642), F(904503643), + F(54374936), F(54374937), F(54374938), + F(1956576150), F(1956576151), F(1956576152), + F(1666750140), F(1666750141), F(1666750142), + F(473324249), F(473324250), F(473324251), + F(1767666896), F(1767666897), F(1767666898); + // clang-format on + EXPECT_EQ(kLargeCosetLDEExpected, large_mat); +} +} // namespace tachyon::math