Merge pull request #441 from kroma-network/feat/impl-radix2ditparallel

feat: impl `Radix2DitParallel`

Author: Ryan Kim

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",

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,
 )
 

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 {

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,
 )
 

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 {

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 {

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

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 {

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",

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 {
@@ -45,22 +48,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);
   });
 }
 
@@ -74,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_