test(math): add radix2_dit_parallel testing

Values were gathered by running [these tests](https://github.com/Plonky3/Plonky3/blob/33b94a8ebaf8fbaace3def6e792cba0dd97b3c42/dft/src/testing.rs) on specific cases.
kroma-network · Jun 18, 2024 · a660cf2 · a660cf2
1 parent 72b5056
commit a660cf2
Show file tree

Hide file tree

Showing 2 changed files with 112 additions and 0 deletions.
diff --git 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
@@ -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 <typename F>
+void PrintTo(const tachyon::math::RowMajorMatrix<F> &m, std::ostream *os) {
+  size_t rows = static_cast<size_t>(m.rows());
+  size_t cols = static_cast<size_t>(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<F> small_mat(1, 3);
+  small_mat << F(0), F(1), F(2);
+  Radix2DitParallel<F>::DFTBatch(small_mat);
+  RowMajorMatrix<BabyBear> kSmallDFTExpected(1, 3);
+  kSmallDFTExpected << F(0), F(1), F(2);
+  EXPECT_EQ(kSmallDFTExpected, small_mat);
+  // Larger case
+  RowMajorMatrix<F> 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<F>::DFTBatch(large_mat);
+  RowMajorMatrix<BabyBear> 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<F> small_mat(1, 3);
+  small_mat << F(0), F(1), F(2);
+  F shift = F(31);
+  Radix2DitParallel<F>::CosetLDEBatch(small_mat, 1, shift);
+  RowMajorMatrix<BabyBear> 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<F> 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<F>::CosetLDEBatch(large_mat, 1, shift);
+  // clang-format off
+  RowMajorMatrix<BabyBear> 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