From b6852aa767a0c9d66068d28b44f4b564a5aa23ee Mon Sep 17 00:00:00 2001
From: ashjeong <ashjeong@umich.edu>
Date: Fri, 21 Jun 2024 10:24:04 +0900
Subject: [PATCH 01/12] chore(math): remove unneeded dependency

---
 tachyon/math/matrix/matrix_utils_unittest.cc | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tachyon/math/matrix/matrix_utils_unittest.cc b/tachyon/math/matrix/matrix_utils_unittest.cc
index ee60e8603..76a7520e3 100644
--- a/tachyon/math/matrix/matrix_utils_unittest.cc
+++ b/tachyon/math/matrix/matrix_utils_unittest.cc
@@ -1,7 +1,6 @@
 #include "tachyon/math/matrix/matrix_utils.h"
 
 #include "tachyon/base/strings/string_util.h"
-#include "tachyon/build/build_config.h"
 #include "tachyon/math/finite_fields/baby_bear/packed_baby_bear.h"
 #include "tachyon/math/finite_fields/test/finite_field_test.h"
 #include "tachyon/math/finite_fields/test/gf7.h"

From 0a5b1be636b9471f9fcc71b4a3dfbfbd03df3e47 Mon Sep 17 00:00:00 2001
From: ashjeong <ashjeong@umich.edu>
Date: Tue, 18 Jun 2024 14:29:14 +0900
Subject: [PATCH 02/12] chore: change "panic" to "crash" in comments

---
 .../polynomials/univariate/univariate_evaluation_domain.h     | 2 +-
 tachyon/zk/plonk/constraint_system/constraint_system.h        | 2 +-
 tachyon/zk/plonk/layout/assignment.h                          | 4 ++--
 tachyon/zk/plonk/vanishing/vanishing_utils.h                  | 2 +-
 4 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tachyon/math/polynomials/univariate/univariate_evaluation_domain.h b/tachyon/math/polynomials/univariate/univariate_evaluation_domain.h
index d3df05b0e..f8396f07d 100644
--- a/tachyon/math/polynomials/univariate/univariate_evaluation_domain.h
+++ b/tachyon/math/polynomials/univariate/univariate_evaluation_domain.h
@@ -308,7 +308,7 @@ class UnivariateEvaluationDomain : public EvaluationDomain<F, MaxDegree> {
   // Return the filter polynomial of |*this| with respect to |subdomain|.
   // Assumes that |subdomain| is contained within |*this|.
   //
-  // Panics if |subdomain| is not contained within |*this|.
+  // Crashes if |subdomain| is not contained within |*this|.
   constexpr DensePoly GetFilterPolynomial(
       const UnivariateEvaluationDomain& subdomain) const {
     SparsePoly domain_vanishing_poly =
diff --git a/tachyon/zk/plonk/constraint_system/constraint_system.h b/tachyon/zk/plonk/constraint_system/constraint_system.h
index 4a0c03aa7..52cf433b0 100644
--- a/tachyon/zk/plonk/constraint_system/constraint_system.h
+++ b/tachyon/zk/plonk/constraint_system/constraint_system.h
@@ -403,7 +403,7 @@ class ConstraintSystem {
   // Creates a new gate.
   //
   // A gate is required to contain polynomial constraints. This method will
-  // panic if |constrain| returns an empty iterator.
+  // crash if |constrain| returns an empty iterator.
   void CreateGate(std::string_view name, ConstrainCallback constrain) {
     VirtualCells cells(this);
     std::vector<Constraint<F>> constraints = std::move(constrain).Run(cells);
diff --git a/tachyon/zk/plonk/layout/assignment.h b/tachyon/zk/plonk/layout/assignment.h
index ebeb20e07..fb60e5adf 100644
--- a/tachyon/zk/plonk/layout/assignment.h
+++ b/tachyon/zk/plonk/layout/assignment.h
@@ -29,7 +29,7 @@ class Assignment {
 
   // Creates a new region and enters into it.
   //
-  // Panics if we are currently in a region (if |ExitRegion()| was not called).
+  // Crashes if we are currently in a region (if |ExitRegion()| was not called).
   //
   // Not intended for downstream consumption; use |Layouter::AssignRegion()|
   // instead.
@@ -43,7 +43,7 @@ class Assignment {
 
   // Exits the current region.
   //
-  // Panics if we are not currently in a region (if |EnterRegion()| was not
+  // Crashes if we are not currently in a region (if |EnterRegion()| was not
   // called).
   //
   // Not intended for downstream consumption; use |Layouter::AssignRegion()|
diff --git a/tachyon/zk/plonk/vanishing/vanishing_utils.h b/tachyon/zk/plonk/vanishing/vanishing_utils.h
index c8993e57d..8cebf470a 100644
--- a/tachyon/zk/plonk/vanishing/vanishing_utils.h
+++ b/tachyon/zk/plonk/vanishing/vanishing_utils.h
@@ -111,7 +111,7 @@ void DistributePowersZeta(ExtendedPoly& poly, bool into_coset) {
 // This takes us from the extended evaluation domain and gets us the quotient
 // polynomial coefficients.
 //
-// This function will panic if the provided vector is not the correct length.
+// This function will crash if the provided vector is not the correct length.
 template <typename F, typename ExtendedPoly, typename ExtendedEvals,
           typename ExtendedDomain>
 ExtendedPoly ExtendedToCoeff(ExtendedEvals&& evals,

From ad8747aeee852832fa0aaab1da843b1861a8a862 Mon Sep 17 00:00:00 2001
From: ashjeong <ashjeong@umich.edu>
Date: Mon, 24 Jun 2024 10:42:44 +0900
Subject: [PATCH 03/12] refac(math): change `PackRowHorizontally` to work with
 `Radix2` fft batch butterflies

---
 tachyon/math/matrix/matrix_utils.h           | 27 +++++-----
 tachyon/math/matrix/matrix_utils_unittest.cc | 53 ++++++++++++--------
 2 files changed, 46 insertions(+), 34 deletions(-)

diff --git a/tachyon/math/matrix/matrix_utils.h b/tachyon/math/matrix/matrix_utils.h
index d7682316e..81c0906af 100644
--- a/tachyon/math/matrix/matrix_utils.h
+++ b/tachyon/math/matrix/matrix_utils.h
@@ -45,22 +45,25 @@ MakeCirculant(const Eigen::MatrixBase<ArgType>& arg) {
                                  CirculantFunctor<ArgType>(arg.derived()));
 }
 
+// NOTE(ashjeong): Important! |matrix| should carry the same amount of rows as
+// the parent matrix it is a block from. |PackRowHorizontally| currently only
+// supports row-major matrices.
 template <typename PackedPrimeField, typename Derived, typename PrimeField>
-std::vector<PackedPrimeField> PackRowHorizontally(
-    const Eigen::MatrixBase<Derived>& matrix, size_t row,
-    std::vector<PrimeField>& remaining_values) {
+std::vector<PackedPrimeField*> PackRowHorizontally(
+    Eigen::Block<Derived>& matrix, size_t row,
+    std::vector<PrimeField*>& remaining_values) {
+  static_assert(Derived::Options & Eigen::RowMajorBit);
   size_t num_packed = matrix.cols() / PackedPrimeField::N;
   size_t remaining_start_idx = num_packed * PackedPrimeField::N;
-  remaining_values =
-      base::CreateVector(matrix.cols() - remaining_start_idx,
-                         [row, remaining_start_idx, &matrix](size_t col) {
-                           return matrix(row, remaining_start_idx + col);
-                         });
-
+  remaining_values = base::CreateVector(
+      matrix.cols() - remaining_start_idx,
+      [row, remaining_start_idx, &matrix](size_t col) {
+        return reinterpret_cast<PrimeField*>(
+            matrix.data() + row * matrix.cols() + remaining_start_idx + col);
+      });
   return base::CreateVector(num_packed, [row, &matrix](size_t col) {
-    return PackedPrimeField::From([row, col, &matrix](size_t i) {
-      return matrix(row, PackedPrimeField::N * col + i);
-    });
+    return reinterpret_cast<PackedPrimeField*>(
+        matrix.data() + row * matrix.cols() + PackedPrimeField::N * col);
   });
 }
 
diff --git a/tachyon/math/matrix/matrix_utils_unittest.cc b/tachyon/math/matrix/matrix_utils_unittest.cc
index 76a7520e3..99dde96eb 100644
--- a/tachyon/math/matrix/matrix_utils_unittest.cc
+++ b/tachyon/math/matrix/matrix_utils_unittest.cc
@@ -26,30 +26,39 @@ TEST_F(MatrixPackingTest, PackRowHorizontally) {
   constexpr size_t N = PackedBabyBear::N;
   constexpr size_t R = 3;
 
-  Matrix<BabyBear> matrix = Matrix<BabyBear>::Random(2 * N, 2 * N);
-  std::vector<BabyBear> remaining_values;
-  std::vector<PackedBabyBear> packed_values =
-      PackRowHorizontally<PackedBabyBear>(matrix, R, remaining_values);
-  ASSERT_TRUE(remaining_values.empty());
-  ASSERT_EQ(packed_values.size(), 2);
-  for (size_t i = 0; i < packed_values.size(); ++i) {
-    for (size_t j = 0; j < N; ++j) {
-      EXPECT_EQ(packed_values[i][j], matrix(R, i * N + j));
+  {
+    RowMajorMatrix<BabyBear> matrix =
+        RowMajorMatrix<BabyBear>::Random(2 * N, 2 * N);
+    Eigen::Block<RowMajorMatrix<BabyBear>> mat =
+        matrix.block(0, 0, matrix.rows(), matrix.cols());
+    std::vector<BabyBear*> remaining_values;
+    std::vector<PackedBabyBear*> packed_values =
+        PackRowHorizontally<PackedBabyBear>(mat, R, remaining_values);
+    ASSERT_TRUE(remaining_values.empty());
+    ASSERT_EQ(packed_values.size(), 2);
+    for (size_t i = 0; i < packed_values.size(); ++i) {
+      for (size_t j = 0; j < N; ++j) {
+        EXPECT_EQ((*packed_values[i])[j], matrix(R, i * N + j));
+      }
     }
   }
-
-  matrix = Matrix<BabyBear>::Random(2 * N - 1, 2 * N - 1);
-  remaining_values.clear();
-  packed_values =
-      PackRowHorizontally<PackedBabyBear>(matrix, R, remaining_values);
-  ASSERT_EQ(remaining_values.size(), N - 1);
-  ASSERT_EQ(packed_values.size(), 1);
-  for (size_t i = 0; i < remaining_values.size(); ++i) {
-    EXPECT_EQ(remaining_values[i], matrix(R, packed_values.size() * N + i));
-  }
-  for (size_t i = 0; i < packed_values.size(); ++i) {
-    for (size_t j = 0; j < N; ++j) {
-      EXPECT_EQ(packed_values[i][j], matrix(R, i * N + j));
+  {
+    RowMajorMatrix<BabyBear> matrix =
+        RowMajorMatrix<BabyBear>::Random(2 * N - 1, 2 * N - 1);
+    Eigen::Block<RowMajorMatrix<BabyBear>> mat =
+        matrix.block(0, 0, matrix.rows(), matrix.cols());
+    std::vector<BabyBear*> remaining_values;
+    std::vector<PackedBabyBear*> packed_values =
+        PackRowHorizontally<PackedBabyBear>(mat, R, remaining_values);
+    ASSERT_EQ(remaining_values.size(), N - 1);
+    ASSERT_EQ(packed_values.size(), 1);
+    for (size_t i = 0; i < remaining_values.size(); ++i) {
+      EXPECT_EQ(*remaining_values[i], matrix(R, packed_values.size() * N + i));
+    }
+    for (size_t i = 0; i < packed_values.size(); ++i) {
+      for (size_t j = 0; j < N; ++j) {
+        EXPECT_EQ((*packed_values[i])[j], matrix(R, i * N + j));
+      }
     }
   }
 }

From be2eb1741f95ce1f5f8cb6c810c9217c505922d0 Mon Sep 17 00:00:00 2001
From: ashjeong <ashjeong@umich.edu>
Date: Mon, 24 Jun 2024 10:39:53 +0900
Subject: [PATCH 04/12] feat(math): add matrix manipulation functions

---
 tachyon/math/matrix/BUILD.bazel    |  2 ++
 tachyon/math/matrix/matrix_utils.h | 47 ++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/tachyon/math/matrix/BUILD.bazel b/tachyon/math/matrix/BUILD.bazel
index db7bae3f3..5c4ac12fa 100644
--- a/tachyon/math/matrix/BUILD.bazel
+++ b/tachyon/math/matrix/BUILD.bazel
@@ -26,6 +26,8 @@ tachyon_cc_library(
     name = "matrix_utils",
     hdrs = ["matrix_utils.h"],
     deps = [
+        "//tachyon/base:bits",
+        "//tachyon/base:openmp_util",
         "//tachyon/base/containers:container_util",
         "//tachyon/math/finite_fields:packed_prime_field_traits_forward",
         "@eigen_archive//:eigen3",
diff --git a/tachyon/math/matrix/matrix_utils.h b/tachyon/math/matrix/matrix_utils.h
index 81c0906af..6b8a2b551 100644
--- a/tachyon/math/matrix/matrix_utils.h
+++ b/tachyon/math/matrix/matrix_utils.h
@@ -1,11 +1,14 @@
 #ifndef TACHYON_MATH_MATRIX_MATRIX_UTILS_H_
 #define TACHYON_MATH_MATRIX_MATRIX_UTILS_H_
 
+#include <utility>
 #include <vector>
 
 #include "third_party/eigen3/Eigen/Core"
 
+#include "tachyon/base/bits.h"
 #include "tachyon/base/containers/container_util.h"
+#include "tachyon/base/openmp_util.h"
 #include "tachyon/math/finite_fields/packed_prime_field_traits_forward.h"
 
 namespace tachyon::math {
@@ -77,6 +80,50 @@ std::vector<PackedPrimeField> PackRowVertically(
   });
 }
 
+// Expands a |Eigen::MatrixBase|'s rows from |rows| to |rows|^(|added_bits|),
+// moving values from row |i| to row |i|^(|added_bits|). All new entries are set
+// to |F::Zero()|.
+template <typename Derived>
+void ExpandInPlaceWithZeroPad(Eigen::MatrixBase<Derived>& mat,
+                              size_t added_bits) {
+  if (added_bits == 0) {
+    return;
+  }
+
+  Eigen::Index original_rows = mat.rows();
+  Eigen::Index new_rows = mat.rows() << added_bits;
+  Eigen::Index cols = mat.cols();
+
+  Derived padded = Derived::Zero(new_rows, cols);
+
+  OPENMP_PARALLEL_FOR(Eigen::Index row = 0; row < original_rows; ++row) {
+    Eigen::Index padded_row_index = row << added_bits;
+    // TODO(ashjeong): Check if moved properly
+    padded.row(padded_row_index) = std::move(mat.row(row));
+  }
+  mat = std::move(padded);
+}
+
+// Swaps rows of a |Eigen::MatrixBase| such that each row is changed to the row
+// accessed with the reversed bits of the current index. Crashes if the number
+// of rows is not a power of two.
+template <typename Derived>
+void ReverseMatrixIndexBits(Eigen::MatrixBase<Derived>& mat) {
+  size_t rows = static_cast<size_t>(mat.rows());
+  if (rows == 0) {
+    return;
+  }
+  CHECK(base::bits::IsPowerOfTwo(rows));
+  size_t log_n = base::bits::Log2Ceiling(rows);
+
+  OPENMP_PARALLEL_FOR(size_t row = 1; row < rows; ++row) {
+    size_t ridx = base::bits::BitRev(row) >> (sizeof(size_t) * 8 - log_n);
+    if (row < ridx) {
+      mat.row(row).swap(mat.row(ridx));
+    }
+  }
+}
+
 }  // namespace tachyon::math
 
 #endif  // TACHYON_MATH_MATRIX_MATRIX_UTILS_H_

From da095d06e415ac2d4a2b911951c64c1c28f0da5b Mon Sep 17 00:00:00 2001
From: ashjeong <ashjeong@umich.edu>
Date: Thu, 20 Jun 2024 19:51:40 +0900
Subject: [PATCH 05/12] refac(math): change `BabyBear` to `fft_prime_field`

See [here](https://github.com/Plonky3/Plonky3/blob/d9ef3902abcc6a34b75f6820da168460aee26763/baby-bear/src/baby_bear.rs#L63)
---
 tachyon/math/finite_fields/baby_bear/BUILD.bazel | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tachyon/math/finite_fields/baby_bear/BUILD.bazel b/tachyon/math/finite_fields/baby_bear/BUILD.bazel
index 93ae5338d..346d4be69 100644
--- a/tachyon/math/finite_fields/baby_bear/BUILD.bazel
+++ b/tachyon/math/finite_fields/baby_bear/BUILD.bazel
@@ -1,17 +1,24 @@
+load("@bazel_skylib//rules:common_settings.bzl", "string_flag")
 load("//bazel:tachyon.bzl", "if_aarch64", "if_has_avx512", "if_x86_64")
 load("//bazel:tachyon_cc.bzl", "tachyon_avx512_defines", "tachyon_cc_library")
 load("//tachyon/math/finite_fields/generator/ext_prime_field_generator:build_defs.bzl", "generate_fp4s")
-load("//tachyon/math/finite_fields/generator/prime_field_generator:build_defs.bzl", "generate_prime_fields")
+load("//tachyon/math/finite_fields/generator/prime_field_generator:build_defs.bzl", "SUBGROUP_GENERATOR", "generate_fft_prime_fields")
 
 package(default_visibility = ["//visibility:public"])
 
-generate_prime_fields(
+string_flag(
+    name = SUBGROUP_GENERATOR,
+    build_setting_default = "31",
+)
+
+generate_fft_prime_fields(
     name = "baby_bear",
     class_name = "BabyBear",
     # 2³¹ - 2²⁷ + 1
     # Hex: 0x78000001
     modulus = "2013265921",
     namespace = "tachyon::math",
+    subgroup_generator = ":" + SUBGROUP_GENERATOR,
     use_montgomery = True,
 )
 

From 7d48e57bc63dd35f11b2487f3cbfa1462b0c7a6b Mon Sep 17 00:00:00 2001
From: ashjeong <ashjeong@umich.edu>
Date: Thu, 20 Jun 2024 19:51:18 +0900
Subject: [PATCH 06/12] refac(math): change `KoalaBear` to `fft_prime_field`

See [here](https://github.com/Plonky3/Plonky3/blob/d9ef3902abcc6a34b75f6820da168460aee26763/koala-bear/src/koala_bear.rs#L63)
---
 tachyon/math/finite_fields/koala_bear/BUILD.bazel | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tachyon/math/finite_fields/koala_bear/BUILD.bazel b/tachyon/math/finite_fields/koala_bear/BUILD.bazel
index 38238dbef..180a397eb 100644
--- a/tachyon/math/finite_fields/koala_bear/BUILD.bazel
+++ b/tachyon/math/finite_fields/koala_bear/BUILD.bazel
@@ -1,17 +1,24 @@
+load("@bazel_skylib//rules:common_settings.bzl", "string_flag")
 load("//bazel:tachyon.bzl", "if_aarch64", "if_has_avx512", "if_x86_64")
 load("//bazel:tachyon_cc.bzl", "tachyon_avx512_defines", "tachyon_cc_library")
 load("//tachyon/math/finite_fields/generator/ext_prime_field_generator:build_defs.bzl", "generate_fp2s", "generate_fp4s")
-load("//tachyon/math/finite_fields/generator/prime_field_generator:build_defs.bzl", "generate_prime_fields")
+load("//tachyon/math/finite_fields/generator/prime_field_generator:build_defs.bzl", "SUBGROUP_GENERATOR", "generate_fft_prime_fields")
 
 package(default_visibility = ["//visibility:public"])
 
-generate_prime_fields(
+string_flag(
+    name = SUBGROUP_GENERATOR,
+    build_setting_default = "3",
+)
+
+generate_fft_prime_fields(
     name = "koala_bear",
     class_name = "KoalaBear",
     # 2³¹ - 2²⁴ + 1
     # Hex: 0x7f000001
     modulus = "2130706433",
     namespace = "tachyon::math",
+    subgroup_generator = ":" + SUBGROUP_GENERATOR,
     use_montgomery = True,
 )
 

From 8fae137da5fd9f46a32c15e229e91f61081d75ae Mon Sep 17 00:00:00 2001
From: ashjeong <ashjeong@umich.edu>
Date: Thu, 20 Jun 2024 22:30:44 +0900
Subject: [PATCH 07/12] refac(math): add ability to obtain `PackedPrimeField`
 type from `PrimeField`

For `BabyBear`, `KoalaBear`, `Mersenne31`
---
 tachyon/math/finite_fields/BUILD.bazel                     | 1 +
 tachyon/math/finite_fields/baby_bear/packed_baby_bear.h    | 5 +++++
 tachyon/math/finite_fields/koala_bear/packed_koala_bear.h  | 5 +++++
 tachyon/math/finite_fields/mersenne31/packed_mersenne31.h  | 5 +++++
 .../math/finite_fields/packed_prime_field_traits_forward.h | 2 +-
 tachyon/math/finite_fields/prime_field_base.h              | 7 +++++++
 6 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/tachyon/math/finite_fields/BUILD.bazel b/tachyon/math/finite_fields/BUILD.bazel
index 0d60514ff..7b0ea3156 100644
--- a/tachyon/math/finite_fields/BUILD.bazel
+++ b/tachyon/math/finite_fields/BUILD.bazel
@@ -147,6 +147,7 @@ tachyon_cc_library(
     deps = [
         ":finite_field",
         ":legendre_symbol",
+        ":packed_prime_field_traits_forward",
         ":prime_field_util",
         "//tachyon/base:bits",
         "//tachyon/base/json",
diff --git a/tachyon/math/finite_fields/baby_bear/packed_baby_bear.h b/tachyon/math/finite_fields/baby_bear/packed_baby_bear.h
index 877210c9d..d153d6bc3 100644
--- a/tachyon/math/finite_fields/baby_bear/packed_baby_bear.h
+++ b/tachyon/math/finite_fields/baby_bear/packed_baby_bear.h
@@ -37,6 +37,11 @@ struct FiniteFieldTraits<PackedBabyBear> {
   using Config = BabyBear::Config;
 };
 
+template <>
+struct PackedPrimeFieldTraits<BabyBear> {
+  using PackedPrimeField = PackedBabyBear;
+};
+
 }  // namespace tachyon::math
 
 namespace Eigen {
diff --git a/tachyon/math/finite_fields/koala_bear/packed_koala_bear.h b/tachyon/math/finite_fields/koala_bear/packed_koala_bear.h
index 25d0c1b83..13b64d729 100644
--- a/tachyon/math/finite_fields/koala_bear/packed_koala_bear.h
+++ b/tachyon/math/finite_fields/koala_bear/packed_koala_bear.h
@@ -37,6 +37,11 @@ struct FiniteFieldTraits<PackedKoalaBear> {
   using Config = KoalaBear::Config;
 };
 
+template <>
+struct PackedPrimeFieldTraits<KoalaBear> {
+  using PackedPrimeField = PackedKoalaBear;
+};
+
 }  // namespace tachyon::math
 
 namespace Eigen {
diff --git a/tachyon/math/finite_fields/mersenne31/packed_mersenne31.h b/tachyon/math/finite_fields/mersenne31/packed_mersenne31.h
index b2921ea00..1deddeaf9 100644
--- a/tachyon/math/finite_fields/mersenne31/packed_mersenne31.h
+++ b/tachyon/math/finite_fields/mersenne31/packed_mersenne31.h
@@ -37,6 +37,11 @@ struct FiniteFieldTraits<PackedMersenne31> {
   using Config = Mersenne31::Config;
 };
 
+template <>
+struct PackedPrimeFieldTraits<Mersenne31> {
+  using PackedPrimeField = PackedMersenne31;
+};
+
 }  // namespace tachyon::math
 
 namespace Eigen {
diff --git a/tachyon/math/finite_fields/packed_prime_field_traits_forward.h b/tachyon/math/finite_fields/packed_prime_field_traits_forward.h
index 087fa3f25..fb351f4f4 100644
--- a/tachyon/math/finite_fields/packed_prime_field_traits_forward.h
+++ b/tachyon/math/finite_fields/packed_prime_field_traits_forward.h
@@ -3,7 +3,7 @@
 
 namespace tachyon::math {
 
-template <typename T>
+template <typename T, typename SFINAE = void>
 struct PackedPrimeFieldTraits;
 
 }  // namespace tachyon::math
diff --git a/tachyon/math/finite_fields/prime_field_base.h b/tachyon/math/finite_fields/prime_field_base.h
index b6f0d8ecd..b6619c81f 100644
--- a/tachyon/math/finite_fields/prime_field_base.h
+++ b/tachyon/math/finite_fields/prime_field_base.h
@@ -22,6 +22,7 @@
 #include "tachyon/math/base/gmp/gmp_util.h"
 #include "tachyon/math/finite_fields/finite_field.h"
 #include "tachyon/math/finite_fields/legendre_symbol.h"
+#include "tachyon/math/finite_fields/packed_prime_field_traits_forward.h"
 #include "tachyon/math/finite_fields/prime_field_util.h"
 
 namespace tachyon {
@@ -160,6 +161,12 @@ H AbslHashValue(H h, const F& prime_field) {
   return h;
 }
 
+template <typename T>
+struct PackedPrimeFieldTraits<
+    T, std::enable_if_t<std::is_base_of_v<math::PrimeFieldBase<T>, T>>> {
+  using PackedPrimeField = T;
+};
+
 }  // namespace math
 
 namespace base {

From 8cc6a994f53bb6251107030d25ca7623f1d6e9fe Mon Sep 17 00:00:00 2001
From: ashjeong <ashjeong@umich.edu>
Date: Thu, 27 Jun 2024 10:00:45 +0900
Subject: [PATCH 08/12] refac(math): templatize `ButterflyFnOutIn`

---
 .../polynomials/univariate/mixed_radix_evaluation_domain.h   | 2 +-
 .../math/polynomials/univariate/radix2_evaluation_domain.h   | 3 ++-
 .../polynomials/univariate/univariate_evaluation_domain.h    | 5 +++--
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/tachyon/math/polynomials/univariate/mixed_radix_evaluation_domain.h b/tachyon/math/polynomials/univariate/mixed_radix_evaluation_domain.h
index 37adcabd6..6dff55f60 100644
--- a/tachyon/math/polynomials/univariate/mixed_radix_evaluation_domain.h
+++ b/tachyon/math/polynomials/univariate/mixed_radix_evaluation_domain.h
@@ -270,7 +270,7 @@ class MixedRadixEvaluationDomain
       for (size_t k = 0; k < n; k += 2 * m) {
         F w = F::One();
         for (size_t j = 0; j < m; ++j) {
-          UnivariateEvaluationDomain<F, MaxDegree>::ButterflyFnOutIn(
+          UnivariateEvaluationDomain<F, MaxDegree>::template ButterflyFnOutIn(
               a.at(k + j), a.at((k + m) + j), w);
           w *= w_m;
         }
diff --git a/tachyon/math/polynomials/univariate/radix2_evaluation_domain.h b/tachyon/math/polynomials/univariate/radix2_evaluation_domain.h
index 54cf0cce3..e0ba38e70 100644
--- a/tachyon/math/polynomials/univariate/radix2_evaluation_domain.h
+++ b/tachyon/math/polynomials/univariate/radix2_evaluation_domain.h
@@ -188,7 +188,8 @@ class Radix2EvaluationDomain : public UnivariateEvaluationDomain<F, MaxDegree> {
       fn = UnivariateEvaluationDomain<F, MaxDegree>::ButterflyFnInOut;
     } else {
       static_assert(Order == FFTOrder::kOutIn);
-      fn = UnivariateEvaluationDomain<F, MaxDegree>::ButterflyFnOutIn;
+      fn = UnivariateEvaluationDomain<F,
+                                      MaxDegree>::template ButterflyFnOutIn<F>;
     }
 
     // Each butterfly cluster uses 2 * |gap| positions.
diff --git a/tachyon/math/polynomials/univariate/univariate_evaluation_domain.h b/tachyon/math/polynomials/univariate/univariate_evaluation_domain.h
index f8396f07d..8f789477c 100644
--- a/tachyon/math/polynomials/univariate/univariate_evaluation_domain.h
+++ b/tachyon/math/polynomials/univariate/univariate_evaluation_domain.h
@@ -458,10 +458,11 @@ class UnivariateEvaluationDomain : public EvaluationDomain<F, MaxDegree> {
   //             | c₀ * ω⁰ + c₁ * ω³ + c₂ * ω⁶ + c₃ * ω⁹ |
   // Note that the coefficients are out of order the evaluations are in order(should be swapped before).
   // clang-format on
-  constexpr static void ButterflyFnOutIn(F& lo, F& hi, const F& root) {
+  template <typename FTy>
+  constexpr static void ButterflyFnOutIn(FTy& lo, FTy& hi, const FTy& root) {
     hi *= root;
 
-    F neg = lo - hi;
+    FTy neg = lo - hi;
 
     lo += hi;
 

From df90944915b16fbbdd8fb68d78c74cdf7073b5e6 Mon Sep 17 00:00:00 2001
From: ashjeong <ashjeong@umich.edu>
Date: Thu, 27 Jun 2024 14:51:36 +0900
Subject: [PATCH 09/12] feat(math): implement `TwoAdicSubgroup`

---
 .../math/polynomials/univariate/BUILD.bazel   |  9 +++
 .../univariate/two_adic_subgroup.h            | 78 +++++++++++++++++++
 2 files changed, 87 insertions(+)
 create mode 100644 tachyon/math/polynomials/univariate/two_adic_subgroup.h

diff --git a/tachyon/math/polynomials/univariate/BUILD.bazel b/tachyon/math/polynomials/univariate/BUILD.bazel
index 097288cec..51878b94d 100644
--- a/tachyon/math/polynomials/univariate/BUILD.bazel
+++ b/tachyon/math/polynomials/univariate/BUILD.bazel
@@ -37,6 +37,15 @@ tachyon_cc_library(
     ],
 )
 
+tachyon_cc_library(
+    name = "two_adic_subgroup",
+    hdrs = ["two_adic_subgroup.h"],
+    deps = [
+        "//tachyon/base:optional",
+        "//tachyon/math/matrix:matrix_types",
+    ],
+)
+
 tachyon_cc_library(
     name = "univariate_evaluation_domain",
     hdrs = ["univariate_evaluation_domain.h"],
diff --git a/tachyon/math/polynomials/univariate/two_adic_subgroup.h b/tachyon/math/polynomials/univariate/two_adic_subgroup.h
new file mode 100644
index 000000000..045ccfcf8
--- /dev/null
+++ b/tachyon/math/polynomials/univariate/two_adic_subgroup.h
@@ -0,0 +1,78 @@
+// Copyright (c) 2022 The Plonky3 Authors
+// Use of this source code is governed by a MIT/Apache-2.0 style license that
+// can be found in the LICENSE-MIT.plonky3 and the LICENCE-APACHE.plonky3
+// file.
+
+#ifndef TACHYON_MATH_POLYNOMIALS_UNIVARIATE_TWO_ADIC_SUBGROUP_H_
+#define TACHYON_MATH_POLYNOMIALS_UNIVARIATE_TWO_ADIC_SUBGROUP_H_
+
+#include <optional>
+#include <vector>
+
+#include "tachyon/base/optional.h"
+#include "tachyon/math/matrix/matrix_types.h"
+
+namespace tachyon::math {
+
+template <typename F>
+class TwoAdicSubgroup {
+ public:
+  virtual ~TwoAdicSubgroup() = default;
+
+  // Compute the discrete Fourier transform (DFT) of each column in |mat|.
+  virtual void FFTBatch(RowMajorMatrix<F>& mat) = 0;
+
+  // Compute the inverse DFT of each column in |mat|.
+  void IFFTBatch(RowMajorMatrix<F>& mat) {
+    static_assert(F::Config::kModulusBits <= 32);
+    FFTBatch(mat);
+    Eigen::Index rows = mat.rows();
+    F inv = unwrap(F(rows).Inverse());
+
+    mat *= inv;
+
+    for (Eigen::Index row = 1; row < rows / 2; ++row) {
+      mat.row(row).swap(mat.row(rows - row));
+    }
+  }
+
+  // Compute the "coset DFT" of each column in |mat|. This can be viewed as
+  // interpolation onto a coset of a multiplicative subgroup, rather than the
+  // subgroup itself.
+  void CosetFFTBatch(RowMajorMatrix<F>& mat) {
+    static_assert(F::Config::kModulusBits <= 32);
+    // Observe that
+    // yᵢ = ∑ⱼ cⱼ (s gⁱ)ʲ
+    //    = ∑ⱼ (cⱼ sʲ) (gⁱ)ʲ
+    // which has the structure of an ordinary DFT, except each coefficient cⱼ
+    // is first replaced by cⱼ s.
+    Eigen::Index rows = mat.rows();
+    Eigen::Index cols = mat.cols();
+
+    std::vector<F> weights = F::GetSuccessivePowers(
+        rows, F::FromMontgomery(F::Config::kSubgroupGenerator));
+    OPENMP_PARALLEL_NESTED_FOR(Eigen::Index row = 0; row < rows; ++row) {
+      for (Eigen::Index col = 0; col < cols; ++col) {
+        mat(row, col) *= weights[row];
+      }
+    }
+    FFTBatch(mat);
+  }
+
+  // Compute the low-degree extension of each column in |mat| onto a coset of
+  // a larger subgroup.
+  void CosetLDEBatch(RowMajorMatrix<F>& mat, size_t added_bits) {
+    static_assert(F::Config::kModulusBits <= 32);
+    IFFTBatch(mat);
+    Eigen::Index rows = mat.rows();
+    Eigen::Index cols = mat.cols();
+
+    // Possible crash if the new resized length overflows
+    mat.conservativeResizeLike(
+        RowMajorMatrix<F>::Zero(rows << added_bits, cols));
+    CosetFFTBatch(mat);
+  }
+};
+
+}  // namespace tachyon::math
+#endif  // TACHYON_MATH_POLYNOMIALS_UNIVARIATE_TWO_ADIC_SUBGROUP_H_

From 0f02c0267e24fe9ada13a1ff50ed7f5e9b55b7a5 Mon Sep 17 00:00:00 2001
From: ashjeong <ashjeong@umich.edu>
Date: Thu, 20 Jun 2024 03:12:05 +0900
Subject: [PATCH 10/12] feat(math): impl `radix_2_dit_parallel` functions in
 `radix2_evaluation_domain.h`

[`Radix2DitParallel`](https://github.com/Plonky3/Plonky3/blob/main/dft/src/radix_2_dit_parallel.rs)
[`ReverseSliceIndexBits`](https://github.com/Plonky3/Plonky3/blob/33b94a8ebaf8fbaace3def6e792cba0dd97b3c42/util/src/lib.rs#L70-L83)
[Butterfly Function](https://github.com/Plonky3/Plonky3/blob/33b94a8ebaf8fbaace3def6e792cba0dd97b3c42/dft/src/butterflies.rs)
---
 .../math/polynomials/univariate/BUILD.bazel   |   7 +
 .../univariate/radix2_evaluation_domain.h     | 255 +++++++++++++++++-
 2 files changed, 257 insertions(+), 5 deletions(-)

diff --git a/tachyon/math/polynomials/univariate/BUILD.bazel b/tachyon/math/polynomials/univariate/BUILD.bazel
index 51878b94d..cedf54350 100644
--- a/tachyon/math/polynomials/univariate/BUILD.bazel
+++ b/tachyon/math/polynomials/univariate/BUILD.bazel
@@ -28,12 +28,19 @@ tachyon_cc_library(
     name = "radix2_evaluation_domain",
     hdrs = ["radix2_evaluation_domain.h"],
     deps = [
+        ":two_adic_subgroup",
         ":univariate_evaluation_domain",
+        "//tachyon/base:bits",
+        "//tachyon/base:openmp_util",
         "//tachyon/base:parallelize",
         "//tachyon/base/containers:container_util",
+        "//tachyon/math/finite_fields:packed_prime_field_base",
+        "//tachyon/math/matrix:matrix_types",
+        "//tachyon/math/matrix:matrix_utils",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest_prod",
+        "@eigen_archive//:eigen3",
     ],
 )
 
diff --git a/tachyon/math/polynomials/univariate/radix2_evaluation_domain.h b/tachyon/math/polynomials/univariate/radix2_evaluation_domain.h
index e0ba38e70..372d1c306 100644
--- a/tachyon/math/polynomials/univariate/radix2_evaluation_domain.h
+++ b/tachyon/math/polynomials/univariate/radix2_evaluation_domain.h
@@ -3,6 +3,11 @@
 // can be found in the LICENSE-MIT.arkworks and the LICENCE-APACHE.arkworks
 // file.
 
+// Copyright (c) 2022 The Plonky3 Authors
+// Use of this source code is governed by a MIT/Apache-2.0 style license that
+// can be found in the LICENSE-MIT.plonky3 and the LICENCE-APACHE.plonky3
+// file.
+
 // This header defines |Radix2EvaluationDomain|, an |UnivariateEvaluationDomain|
 // for performing various kinds of polynomial arithmetic on top of fields that
 // are FFT-friendly. |Radix2EvaluationDomain| supports FFTs of size at most
@@ -17,16 +22,24 @@
 #include <algorithm>
 #include <cmath>
 #include <memory>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
 #include "absl/memory/memory.h"
 #include "absl/types/span.h"
 #include "gtest/gtest_prod.h"
+#include "third_party/eigen3/Eigen/Core"
 
+#include "tachyon/base/bits.h"
 #include "tachyon/base/containers/container_util.h"
 #include "tachyon/base/logging.h"
+#include "tachyon/base/openmp_util.h"
 #include "tachyon/base/parallelize.h"
+#include "tachyon/math/finite_fields/packed_prime_field_traits_forward.h"
+#include "tachyon/math/matrix/matrix_types.h"
+#include "tachyon/math/matrix/matrix_utils.h"
+#include "tachyon/math/polynomials/univariate/two_adic_subgroup.h"
 #include "tachyon/math/polynomials/univariate/univariate_evaluation_domain.h"
 #include "tachyon/math/polynomials/univariate/univariate_polynomial.h"
 
@@ -37,13 +50,19 @@ namespace tachyon::math {
 // power-of-2.
 template <typename F,
           size_t MaxDegree = (size_t{1} << F::Config::kTwoAdicity) - 1>
-class Radix2EvaluationDomain : public UnivariateEvaluationDomain<F, MaxDegree> {
+class Radix2EvaluationDomain : public UnivariateEvaluationDomain<F, MaxDegree>,
+                               public TwoAdicSubgroup<F> {
  public:
   using Base = UnivariateEvaluationDomain<F, MaxDegree>;
   using Field = F;
   using Evals = UnivariateEvaluations<F, MaxDegree>;
   using DensePoly = UnivariateDensePolynomial<F, MaxDegree>;
   using SparsePoly = UnivariateSparsePolynomial<F, MaxDegree>;
+  using PackedPrimeField =
+      // NOLINTNEXTLINE(whitespace/operators)
+      std::conditional_t<F::Config::kModulusBits <= 32,
+                         typename PackedPrimeFieldTraits<F>::PackedPrimeField,
+                         F>;
 
   constexpr static size_t kMaxDegree = MaxDegree;
   // Factor that determines if a the degree aware FFT should be called.
@@ -73,6 +92,76 @@ class Radix2EvaluationDomain : public UnivariateEvaluationDomain<F, MaxDegree> {
     return base::bits::SafeLog2Ceiling(num_coeffs) <= F::Config::kTwoAdicity;
   }
 
+  void FFTBatch(RowMajorMatrix<F>& mat) override {
+    if constexpr (F::Config::kModulusBits > 32) {
+      NOTREACHED();
+    }
+    CHECK_EQ(this->size_, static_cast<size_t>(mat.rows()));
+    size_t log_h = this->log_size_of_group_;
+    mid_ = log_h / 2;
+
+    // The first half looks like a normal DIT.
+    ReverseMatrixIndexBits(mat);
+    RunParallelRowChunks(mat, roots_vec_[log_h - 1], packed_roots_vec_[0]);
+
+    // For the second half, we flip the DIT, working in bit-reversed order.
+    ReverseMatrixIndexBits(mat);
+    RunParallelRowChunksReversed(mat, rev_roots_vec_, packed_roots_vec_[1]);
+    ReverseMatrixIndexBits(mat);
+  }
+
+  CONSTEXPR_IF_NOT_OPENMP void CosetLDEBatch(RowMajorMatrix<F>& mat,
+                                             size_t added_bits) {
+    if constexpr (F::Config::kModulusBits > 32) {
+      NOTREACHED();
+    }
+    CHECK_EQ(this->size_, static_cast<size_t>(mat.rows()));
+    size_t log_h = this->log_size_of_group_;
+    mid_ = log_h / 2;
+
+    // The first half looks like a normal DIT.
+    ReverseMatrixIndexBits(mat);
+    RunParallelRowChunks(mat, inv_roots_vec_[0], packed_inv_roots_vec_[0]);
+
+    // For the second half, we flip the DIT, working in bit-reversed order.
+    ReverseMatrixIndexBits(mat);
+    RunParallelRowChunksReversed(mat, rev_inv_roots_vec_,
+                                 packed_inv_roots_vec_[1]);
+    // We skip the final bit-reversal, since the next FFT expects bit-reversed
+    // input.
+
+    // Rescale coefficients in two ways:
+    // - divide by number of rows (since we're doing an inverse DFT)
+    // - multiply by powers of the coset shift (see default coset LDE impl for
+    // an explanation)
+    std::vector<F> weights = F::GetSuccessivePowers(
+        this->size_, F::FromMontgomery(F::Config::kSubgroupGenerator),
+        this->size_inv_);
+    OPENMP_PARALLEL_FOR(size_t row = 0; row < weights.size(); ++row) {
+      // Reverse bits because |mat| is encoded in bit-reversed order
+      mat.row(base::bits::BitRev(row) >>
+              (sizeof(size_t) * 8 - this->log_size_of_group_)) *= weights[row];
+    }
+    ExpandInPlaceWithZeroPad<RowMajorMatrix<F>>(mat, added_bits);
+
+    size_t rows = static_cast<size_t>(mat.rows());
+    CHECK(base::bits::IsPowerOfTwo(rows));
+    std::unique_ptr<Radix2EvaluationDomain> domain =
+        Radix2EvaluationDomain<F>::Create(rows);
+    log_h = domain->log_size_of_group_;
+    mid_ = log_h / 2;
+
+    // The first half looks like a normal DIT.
+    domain->RunParallelRowChunks(mat, domain->roots_vec_[log_h - 1],
+                                 domain->packed_roots_vec_[0]);
+
+    // For the second half, we flip the DIT, working in bit-reversed order.
+    ReverseMatrixIndexBits(mat);
+    domain->RunParallelRowChunksReversed(mat, domain->rev_roots_vec_,
+                                         domain->packed_roots_vec_[1]);
+    ReverseMatrixIndexBits(mat);
+  }
+
  private:
   template <typename T>
   FRIEND_TEST(UnivariateEvaluationDomainTest, RootsOfUnity);
@@ -226,11 +315,36 @@ class Radix2EvaluationDomain : public UnivariateEvaluationDomain<F, MaxDegree> {
     roots_vec_.resize(this->log_size_of_group_);
     inv_roots_vec_.resize(this->log_size_of_group_);
 
+    size_t vec_largest_size = this->size_ / 2;
+
     // Compute biggest vector of |root_vec_| and |inv_root_vec_| first.
-    roots_vec_[this->log_size_of_group_ - 1] =
-        this->GetRootsOfUnity(this->size_ / 2, this->group_gen_);
-    inv_roots_vec_[0] =
-        this->GetRootsOfUnity(this->size_ / 2, this->group_gen_inv_);
+    std::vector<F> largest =
+        this->GetRootsOfUnity(vec_largest_size, this->group_gen_);
+    std::vector<F> largest_inv =
+        this->GetRootsOfUnity(vec_largest_size, this->group_gen_inv_);
+
+    if constexpr (F::Config::kModulusBits <= 32) {
+      packed_roots_vec_.resize(2);
+      packed_inv_roots_vec_.resize(2);
+      packed_roots_vec_[0].resize(vec_largest_size);
+      packed_inv_roots_vec_[0].resize(vec_largest_size);
+      packed_roots_vec_[1].resize(vec_largest_size);
+      packed_inv_roots_vec_[1].resize(vec_largest_size);
+      rev_roots_vec_ = ReverseSliceIndexBits(largest);
+      rev_inv_roots_vec_ = ReverseSliceIndexBits(largest_inv);
+      for (size_t i = 0; i < vec_largest_size; ++i) {
+        packed_roots_vec_[0][i] = PackedPrimeField::Broadcast(largest[i]);
+        packed_inv_roots_vec_[0][i] =
+            PackedPrimeField::Broadcast(largest_inv[i]);
+        packed_inv_roots_vec_[1][i] =
+            PackedPrimeField::Broadcast(rev_roots_vec_[i]);
+        packed_inv_roots_vec_[1][i] =
+            PackedPrimeField::Broadcast(rev_inv_roots_vec_[i]);
+      }
+    }
+
+    roots_vec_[this->log_size_of_group_ - 1] = largest;
+    inv_roots_vec_[0] = largest_inv;
 
     // Prepare space in each vector for the others.
     size_t size = this->size_ / 2;
@@ -268,6 +382,137 @@ class Radix2EvaluationDomain : public UnivariateEvaluationDomain<F, MaxDegree> {
     }
   }
 
+  // This can be used as the first half of a parallelized butterfly network.
+  CONSTEXPR_IF_NOT_OPENMP void RunParallelRowChunks(
+      RowMajorMatrix<F>& mat, const std::vector<F>& twiddles,
+      const std::vector<PackedPrimeField>& packed_twiddles_rev) {
+    if constexpr (F::Config::kModulusBits > 32) {
+      NOTREACHED();
+    }
+    CHECK(base::bits::IsPowerOfTwo(mat.rows()));
+    size_t cols = static_cast<size_t>(mat.cols());
+    size_t chunk_rows = 1 << mid_;
+
+    // max block size: 2^|mid_|
+    // TODO(ashjeong): benchmark between |OPENMP_PARALLEL_FOR| here vs
+    // |OPENMP_NESTED_PARALLEL_FOR| in |RunDitLayers|
+    for (size_t block_start = 0; block_start < this->size_;
+         block_start += chunk_rows) {
+      size_t cur_chunk_rows = std::min(chunk_rows, this->size_ - block_start);
+      Eigen::Block<RowMajorMatrix<F>> submat =
+          mat.block(block_start, 0, cur_chunk_rows, cols);
+      for (size_t layer = 0; layer < mid_; ++layer) {
+        RunDitLayers(submat, layer, absl::MakeSpan(twiddles),
+                     absl::MakeSpan(packed_twiddles_rev), false);
+      }
+    }
+  }
+
+  // This can be used as the second half of a parallelized butterfly network.
+  CONSTEXPR_IF_NOT_OPENMP void RunParallelRowChunksReversed(
+      RowMajorMatrix<F>& mat, const std::vector<F>& twiddles_rev,
+      const std::vector<PackedPrimeField>& packed_twiddles_rev) {
+    if constexpr (F::Config::kModulusBits > 32) {
+      NOTREACHED();
+    }
+    CHECK(base::bits::IsPowerOfTwo(mat.rows()));
+    size_t cols = static_cast<size_t>(mat.cols());
+    size_t chunk_rows = 1 << (this->log_size_of_group_ - mid_);
+
+    // max block size: 2^(|this->log_size_of_group_| - |mid_|)
+    // TODO(ashjeong): benchmark between |OPENMP_PARALLEL_FOR| here vs
+    // |OPENMP_NESTED_PARALLEL_FOR| in |RunDitLayers|
+    for (size_t block_start = 0; block_start < this->size_;
+         block_start += chunk_rows) {
+      size_t thread = block_start / chunk_rows;
+      size_t cur_chunk_rows = std::min(chunk_rows, this->size_ - block_start);
+      Eigen::Block<RowMajorMatrix<F>> submat =
+          mat.block(block_start, 0, cur_chunk_rows, cols);
+      for (size_t layer = mid_; layer < this->log_size_of_group_; ++layer) {
+        size_t first_block = thread << (layer - mid_);
+        RunDitLayers(submat, layer,
+                     absl::MakeSpan(twiddles_rev.data() + first_block,
+                                    twiddles_rev.size() - first_block),
+                     absl::MakeSpan(packed_twiddles_rev.data() + first_block,
+                                    packed_twiddles_rev.size() - first_block),
+                     true);
+      }
+    }
+  }
+
+  CONSTEXPR_IF_NOT_OPENMP void RunDitLayers(
+      Eigen::Block<RowMajorMatrix<F>>& submat, size_t layer,
+      const absl::Span<const F>& twiddles,
+      const absl::Span<const PackedPrimeField>& packed_twiddles, bool rev) {
+    if constexpr (F::Config::kModulusBits > 32) {
+      NOTREACHED();
+    }
+    size_t layer_rev = this->log_size_of_group_ - 1 - layer;
+    size_t half_block_size = rev ? 1 << layer_rev : 1 << layer;
+    size_t block_size = half_block_size * 2;
+    size_t sub_rows = static_cast<size_t>(submat.rows());
+    DCHECK_GE(sub_rows, block_size);
+
+    OPENMP_PARALLEL_NESTED_FOR(size_t block_start = 0; block_start < sub_rows;
+                               block_start += block_size) {
+      for (size_t i = 0; i < half_block_size; ++i) {
+        size_t lo = block_start + i;
+        size_t hi = lo + half_block_size;
+        const F& twiddle =
+            rev ? twiddles[block_start / block_size] : twiddles[i << layer_rev];
+        const PackedPrimeField& packed_twiddle =
+            rev ? packed_twiddles[block_start / block_size]
+                : packed_twiddles[i << layer_rev];
+        ApplyButterflyToRows(submat, lo, hi, twiddle, packed_twiddle);
+      }
+    }
+  }
+
+  CONSTEXPR_IF_NOT_OPENMP std::vector<F> ReverseSliceIndexBits(
+      const std::vector<F>& vals) {
+    size_t n = vals.size();
+    if (n == 0) {
+      return vals;
+    }
+    CHECK(base::bits::IsPowerOfTwo(n));
+    size_t log_n = base::bits::Log2Ceiling(n);
+
+    std::vector<F> ret = vals;
+    this->SwapElements(ret, n, log_n);
+    return ret;
+  }
+
+  CONSTEXPR_IF_NOT_OPENMP void ApplyButterflyToRows(
+      Eigen::Block<RowMajorMatrix<F>>& mat, size_t row_1, size_t row_2,
+      const F& twiddle, const PackedPrimeField& packed_twiddle) {
+    if constexpr (F::Config::kModulusBits > 32) {
+      NOTREACHED();
+    }
+    std::vector<F*> suffix_1;
+    std::vector<F*> suffix_2;
+
+    std::vector<PackedPrimeField*> shorts_1 =
+        PackRowHorizontally<PackedPrimeField>(mat, row_1, suffix_1);
+    std::vector<PackedPrimeField*> shorts_2 =
+        PackRowHorizontally<PackedPrimeField>(mat, row_2, suffix_2);
+
+    OPENMP_PARALLEL_FOR(size_t i = 0; i < shorts_1.size(); ++i) {
+      UnivariateEvaluationDomain<F, MaxDegree>::template ButterflyFnOutIn<
+          PackedPrimeField>(*shorts_1[i], *shorts_2[i], packed_twiddle);
+    }
+    for (size_t i = 0; i < suffix_1.size(); ++i) {
+      UnivariateEvaluationDomain<F, MaxDegree>::template ButterflyFnOutIn<F>(
+          *suffix_1[i], *suffix_2[i], twiddle);
+    }
+  }
+
+  size_t mid_ = 0;
+  // For small prime fields
+  std::vector<F> rev_roots_vec_;
+  std::vector<F> rev_inv_roots_vec_;
+  std::vector<std::vector<PackedPrimeField>> packed_roots_vec_;
+  std::vector<std::vector<PackedPrimeField>> packed_inv_roots_vec_;
+  // For all finite fields
   std::vector<std::vector<F>> roots_vec_;
   std::vector<std::vector<F>> inv_roots_vec_;
 };

From 7e8cb22fb20f6faebeb46b6c5ccb95947078dfc4 Mon Sep 17 00:00:00 2001
From: ashjeong <ashjeong@umich.edu>
Date: Wed, 19 Jun 2024 14:19:39 +0900
Subject: [PATCH 11/12] feat(math): impl naive radix2 `FFTBatch` and
 `CosetLDEBatch`

[ref](https://github.com/Plonky3/Plonky3/blob/f535642f5bea6a2bfded8c4e51b1f2960b1f6da7/dft/src/naive.rs)
---
 .../math/polynomials/univariate/BUILD.bazel   |  9 ++++
 .../polynomials/univariate/naive_batch_fft.h  | 45 +++++++++++++++++++
 2 files changed, 54 insertions(+)
 create mode 100644 tachyon/math/polynomials/univariate/naive_batch_fft.h

diff --git a/tachyon/math/polynomials/univariate/BUILD.bazel b/tachyon/math/polynomials/univariate/BUILD.bazel
index cedf54350..6b6a1f440 100644
--- a/tachyon/math/polynomials/univariate/BUILD.bazel
+++ b/tachyon/math/polynomials/univariate/BUILD.bazel
@@ -24,6 +24,15 @@ tachyon_cc_library(
     ],
 )
 
+tachyon_cc_library(
+    name = "naive_batch_fft",
+    hdrs = ["naive_batch_fft.h"],
+    deps = [
+        ":two_adic_subgroup",
+        "//tachyon/base:bits",
+    ],
+)
+
 tachyon_cc_library(
     name = "radix2_evaluation_domain",
     hdrs = ["radix2_evaluation_domain.h"],
diff --git a/tachyon/math/polynomials/univariate/naive_batch_fft.h b/tachyon/math/polynomials/univariate/naive_batch_fft.h
new file mode 100644
index 000000000..f1466b2dd
--- /dev/null
+++ b/tachyon/math/polynomials/univariate/naive_batch_fft.h
@@ -0,0 +1,45 @@
+// Copyright (c) 2022 The Plonky3 Authors
+// Use of this source code is governed by a MIT/Apache-2.0 style license that
+// can be found in the LICENSE-MIT.plonky3 and the LICENCE-APACHE.plonky3
+// file.
+
+#ifndef TACHYON_MATH_POLYNOMIALS_UNIVARIATE_NAIVE_BATCH_FFT_H_
+#define TACHYON_MATH_POLYNOMIALS_UNIVARIATE_NAIVE_BATCH_FFT_H_
+
+#include <vector>
+
+#include "tachyon/base/bits.h"
+#include "tachyon/math/polynomials/univariate/two_adic_subgroup.h"
+
+namespace tachyon::math {
+
+template <typename F>
+class NaiveBatchFFT : public TwoAdicSubgroup<F> {
+ public:
+  void FFTBatch(RowMajorMatrix<F>& mat) override {
+    Eigen::Index rows = mat.rows();
+    Eigen::Index cols = mat.cols();
+    CHECK(base::bits::IsPowerOfTwo(rows));
+    F g;
+    CHECK(F::GetRootOfUnity(rows, &g));
+
+    RowMajorMatrix<F> res = RowMajorMatrix<F>::Zero(rows, cols);
+
+    std::vector<F> points = F::GetSuccessivePowers(rows, g);
+    size_t num_points = points.size();
+
+    for (size_t res_r = 0; res_r < num_points; ++res_r) {
+      std::vector<F> point_powers = F::GetSuccessivePowers(rows, points[res_r]);
+      for (size_t src_r = 0; src_r < num_points; ++src_r) {
+        for (Eigen::Index col = 0; col < cols; ++col) {
+          res(res_r, col) += point_powers[src_r] * mat(src_r, col);
+        }
+      }
+    }
+    mat = res;
+  }
+};
+
+}  // namespace tachyon::math
+
+#endif  // TACHYON_MATH_POLYNOMIALS_UNIVARIATE_NAIVE_BATCH_FFT_H_

From fa4c33f16fc54c1491391fc81e97a8b3550c6c96 Mon Sep 17 00:00:00 2001
From: ashjeong <ashjeong@umich.edu>
Date: Thu, 20 Jun 2024 03:12:05 +0900
Subject: [PATCH 12/12] test(math): test `FFTBatch` and `CosetLDEBatch`

---
 .../math/polynomials/univariate/BUILD.bazel   |  4 ++
 .../radix2_evaluation_domain_unittest.cc      | 55 +++++++++++++++++++
 2 files changed, 59 insertions(+)
 create mode 100644 tachyon/math/polynomials/univariate/radix2_evaluation_domain_unittest.cc

diff --git a/tachyon/math/polynomials/univariate/BUILD.bazel b/tachyon/math/polynomials/univariate/BUILD.bazel
index 6b6a1f440..404419e03 100644
--- a/tachyon/math/polynomials/univariate/BUILD.bazel
+++ b/tachyon/math/polynomials/univariate/BUILD.bazel
@@ -143,6 +143,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",
@@ -151,6 +152,7 @@ tachyon_cc_unittest(
     deps = [
         ":lagrange_interpolation",
         ":mixed_radix_evaluation_domain",
+        ":naive_batch_fft",
         ":radix2_evaluation_domain",
         ":univariate_polynomial",
         "//tachyon/base:optional",
@@ -161,6 +163,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
new file mode 100644
index 000000000..5021e2f5d
--- /dev/null
+++ 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