test(math): test FFTBatch and CosetLDEBatch

kroma-network · Jun 27, 2024 · 4cc0ac3 · 4cc0ac3
1 parent ac54e72
commit 4cc0ac3
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 0 deletions.
diff --git a/tachyon/math/polynomials/univariate/BUILD.bazel b/tachyon/math/polynomials/univariate/BUILD.bazel
@@ -140,6 +140,7 @@ tachyon_cc_unittest(
     name = "univariate_unittests",
     srcs = [
         "lagrange_interpolation_unittest.cc",
+        "radix2_evaluation_domain_unittest.cc",
         "univariate_dense_polynomial_unittest.cc",
         "univariate_evaluation_domain_unittest.cc",
         "univariate_evaluations_unittest.cc",
@@ -148,6 +149,7 @@ tachyon_cc_unittest(
     deps = [
         ":lagrange_interpolation",
         ":mixed_radix_evaluation_domain",
+        ":naive_batch_fft",
         ":radix2_evaluation_domain",
         ":univariate_polynomial",
         "//tachyon/base:optional",
@@ -158,6 +160,8 @@ 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:packed_baby_bear",
+        "//tachyon/math/finite_fields/koala_bear:packed_koala_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_evaluation_domain_unittest.cc b/tachyon/math/polynomials/univariate/radix2_evaluation_domain_unittest.cc
@@ -0,0 +1,55 @@
+#include "tachyon/math/polynomials/univariate/radix2_evaluation_domain.h"
+
+#include <memory>
+
+#include "gtest/gtest.h"
+
+#include "tachyon/math/finite_fields/baby_bear/packed_baby_bear.h"
+#include "tachyon/math/finite_fields/koala_bear/packed_koala_bear.h"
+#include "tachyon/math/finite_fields/test/finite_field_test.h"
+#include "tachyon/math/matrix/matrix_types.h"
+#include "tachyon/math/polynomials/univariate/naive_batch_fft.h"
+
+namespace tachyon::math {
+
+namespace {
+
+template <typename F>
+class Radix2EvaluationDomainTest
+    : public FiniteFieldTest<
+          typename PackedPrimeFieldTraits<F>::PackedPrimeField> {};
+
+}  // namespace
+
+using PrimeFieldTypes = testing::Types<BabyBear, KoalaBear>;
+TYPED_TEST_SUITE(Radix2EvaluationDomainTest, PrimeFieldTypes);
+
+TYPED_TEST(Radix2EvaluationDomainTest, FFTBatch) {
+  using F = TypeParam;
+  for (size_t log_r = 0; log_r < 5; ++log_r) {
+    RowMajorMatrix<F> expected = RowMajorMatrix<F>::Random(1 << log_r, 3);
+    RowMajorMatrix<F> result = expected;
+    NaiveBatchFFT<F> naive;
+    naive.FFTBatch(expected);
+    std::unique_ptr<Radix2EvaluationDomain<F>> domain =
+        Radix2EvaluationDomain<F>::Create(1 << log_r);
+    domain->FFTBatch(result);
+    EXPECT_EQ(expected, result);
+  }
+}
+
+TYPED_TEST(Radix2EvaluationDomainTest, CosetLDEBatch) {
+  using F = TypeParam;
+  for (size_t log_r = 0; log_r < 5; ++log_r) {
+    RowMajorMatrix<F> expected = RowMajorMatrix<F>::Random(1 << log_r, 3);
+    RowMajorMatrix<F> result = expected;
+    NaiveBatchFFT<F> naive;
+    naive.CosetLDEBatch(expected, 1);
+    std::unique_ptr<Radix2EvaluationDomain<F>> domain =
+        Radix2EvaluationDomain<F>::Create(1 << log_r);
+    domain->CosetLDEBatch(result, 1);
+    EXPECT_EQ(expected, result);
+  }
+}
+
+}  // namespace tachyon::math