feat(math): add quartic extension for BabyBear and KoalaBear

See https://github.com/Plonky3/Plonky3/blob/11ac8745b21295d5699f47089dbd927836f53fff/baby-bear/src/extension.rs#L6-L11 and https://github.com/Plonky3/Plonky3/blob/11ac8745b21295d5699f47089dbd927836f53fff/koala-bear/src/extension.rs#L40-L45.
kroma-network · Jun 25, 2024 · 081fa8d · 081fa8d
1 parent 1d8b026
commit 081fa8d
Show file tree

Hide file tree

Showing 11 changed files with 1,014 additions and 6 deletions.
diff --git a/tachyon/math/circle/stark/BUILD.bazel b/tachyon/math/circle/stark/BUILD.bazel
@@ -48,6 +48,7 @@ generate_fp2s(
 generate_fp4s(
     name = "fq4",
     base_field = "Fq2",
+    base_field_degree = 2,
     base_field_hdr = "tachyon/math/circle/stark/fq2.h",
     class_name = "Fq4",
     namespace = "tachyon::math::stark",

diff --git a/tachyon/math/finite_fields/BUILD.bazel b/tachyon/math/finite_fields/BUILD.bazel
@@ -72,6 +72,7 @@ tachyon_cc_library(
     hdrs = ["fp4.h"],
     deps = [
         ":quadratic_extension_field",
+        ":quartic_extension_field",
         "//tachyon/math/base/gmp:gmp_util",
     ],
 )
@@ -235,6 +236,18 @@ tachyon_cc_library(
     ],
 )
 
+tachyon_cc_library(
+    name = "quartic_extension_field",
+    hdrs = ["quartic_extension_field.h"],
+    deps = [
+        ":cyclotomic_multiplicative_subgroup",
+        "//tachyon/base/buffer:copyable",
+        "//tachyon/base/json",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
 tachyon_cc_library(
     name = "small_prime_field",
     hdrs = ["small_prime_field.h"],
@@ -275,6 +288,7 @@ tachyon_cc_unittest(
         "prime_field_generator_unittest.cc",
         "prime_field_unittest.cc",
         "quadratic_extension_field_unittest.cc",
+        "quartic_extension_field_unittest.cc",
     ] + select({
         "@platforms//cpu:x86_64": ["packed_prime_field_unittest.cc"],
         "@platforms//cpu:aarch64": ["packed_prime_field_unittest.cc"],
@@ -301,6 +315,7 @@ tachyon_cc_unittest(
         "//tachyon/math/elliptic_curves/secp/secp256k1:fq",
         "//tachyon/math/elliptic_curves/secp/secp256k1:fr",
         "//tachyon/math/finite_fields/baby_bear",
+        "//tachyon/math/finite_fields/baby_bear:baby_bear4",
         "//tachyon/math/finite_fields/binary_fields",
         "//tachyon/math/finite_fields/goldilocks:goldilocks_prime_field",
         "//tachyon/math/finite_fields/koala_bear",

diff --git a/tachyon/math/finite_fields/baby_bear/BUILD.bazel b/tachyon/math/finite_fields/baby_bear/BUILD.bazel
@@ -1,5 +1,6 @@
 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")
 
 package(default_visibility = ["//visibility:public"])
@@ -14,6 +15,17 @@ generate_prime_fields(
     use_montgomery = True,
 )
 
+generate_fp4s(
+    name = "baby_bear4",
+    base_field = "BabyBear",
+    base_field_degree = 1,
+    base_field_hdr = "tachyon/math/finite_fields/baby_bear/baby_bear.h",
+    class_name = "BabyBear4",
+    namespace = "tachyon::math",
+    non_residue = ["11"],
+    deps = [":baby_bear"],
+)
+
 tachyon_cc_library(
     name = "packed_baby_bear",
     hdrs = ["packed_baby_bear.h"],

diff --git a/tachyon/math/finite_fields/finite_field_forwards.h b/tachyon/math/finite_fields/finite_field_forwards.h
@@ -21,7 +21,7 @@ class Fp2;
 template <typename Config>
 class Fp3;
 
-template <typename Config>
+template <typename Config, typename SFINAE = void>
 class Fp4;
 
 template <typename Config, typename SFINAE = void>

diff --git a/tachyon/math/finite_fields/fp4.h b/tachyon/math/finite_fields/fp4.h
@@ -8,11 +8,13 @@
 
 #include "tachyon/math/base/gmp/gmp_util.h"
 #include "tachyon/math/finite_fields/quadratic_extension_field.h"
+#include "tachyon/math/finite_fields/quartic_extension_field.h"
 
 namespace tachyon::math {
 
 template <typename Config>
-class Fp4 final : public QuadraticExtensionField<Fp4<Config>> {
+class Fp4<Config, std::enable_if_t<Config::kDegreeOverBaseField == 2>> final
+    : public QuadraticExtensionField<Fp4<Config>> {
  public:
   using BaseField = typename Config::BaseField;
   using BasePrimeField = typename Config::BasePrimeField;
@@ -89,6 +91,114 @@ class Fp4 final : public QuadraticExtensionField<Fp4<Config>> {
   }
 };
 
+template <typename Config>
+class Fp4<Config, std::enable_if_t<Config::kDegreeOverBaseField == 4>> final
+    : public QuarticExtensionField<Fp4<Config>> {
+ public:
+  using BaseField = typename Config::BaseField;
+  using BasePrimeField = typename Config::BasePrimeField;
+  using FrobeniusCoefficient = typename Config::FrobeniusCoefficient;
+
+  using CpuField = Fp4<Config>;
+  // TODO(chokobole): Implement Fp4Gpu
+  using GpuField = Fp4<Config>;
+
+  using QuarticExtensionField<Fp4<Config>>::QuarticExtensionField;
+
+  static_assert(Config::kDegreeOverBaseField == 4);
+  static_assert(BaseField::ExtensionDegree() == 1);
+
+  constexpr static uint64_t kDegreeOverBasePrimeField = 4;
+
+  static void Init() {
+    Config::Init();
+    // x⁴ = q = |Config::kNonResidue|
+
+    // αᴾ = (α₀ + α₁x + α₂x² + α₃x³)ᴾ
+    //    = α₀ᴾ + α₁ᴾxᴾ + α₂ᴾx²ᴾ + α₃ᴾx³ᴾ
+    //    = α₀ + α₁xᴾ + α₂x²ᴾ + α₃x³ᴾ <- Fermat's little theorem
+    //    = α₀ + α₁xᴾ⁻¹x + α₂x²ᴾ⁻²x² + α₃x³ᴾ⁻³x³
+    //    = α₀ + α₁(x⁴)^((P - 1) / 4) * x + α₂(x⁴)^(2 * (P - 1) / 4) * x² +
+    //      α₃(x⁴)^(3 * (P - 1) / 4) * x³
+    //    = α₀ + α₁ωx + α₂ω²x² + α₃ω³x³, where ω is a quartic root of unity.
+
+    constexpr uint64_t N = BasePrimeField::kLimbNums;
+    // m₁ = P
+    mpz_class m1;
+    if constexpr (BasePrimeField::Config::kModulusBits <= 32) {
+      m1 = mpz_class(BasePrimeField::Config::kModulus);
+    } else {
+      gmp::WriteLimbs(BasePrimeField::Config::kModulus.limbs, N, &m1);
+    }
+
+#define SET_M(d, d_prev) mpz_class m##d = m##d_prev * m1
+
+    // m₂ = m₁ * P = P²
+    SET_M(2, 1);
+    // m₃ = m₂ * P = P³
+    SET_M(3, 2);
+
+#undef SET_M
+
+#define SET_EXP_GMP(d) mpz_class exp##d##_gmp = (m##d - 1) / mpz_class(4)
+
+    // exp₁ = (m₁ - 1) / 4 = (P¹ - 1) / 4
+    SET_EXP_GMP(1);
+    // exp₂ = (m₂ - 1) / 4 = (P² - 1) / 4
+    SET_EXP_GMP(2);
+    // exp₃ = (m₃ - 1) / 4 = (P³ - 1) / 4
+    SET_EXP_GMP(3);
+
+#undef SET_EXP_GMP
+
+    // |kFrobeniusCoeffs[0]| = q^((P⁰ - 1) / 4) = 1
+    Config::kFrobeniusCoeffs[0] = FrobeniusCoefficient::One();
+#define SET_FROBENIUS_COEFF(d)                \
+  BigInt<d * N> exp##d;                       \
+  gmp::CopyLimbs(exp##d##_gmp, exp##d.limbs); \
+  Config::kFrobeniusCoeffs[d] = Config::kNonResidue.Pow(exp##d)
+
+    // |kFrobeniusCoeffs[1]| = q^(exp₁) = q^((P¹ - 1) / 4) = ω
+    SET_FROBENIUS_COEFF(1);
+    // |kFrobeniusCoeffs[2]| = q^(exp₂) = q^((P² - 1) / 4)
+    SET_FROBENIUS_COEFF(2);
+    // |kFrobeniusCoeffs[3]| = q^(exp₃) = q^((P³ - 1) / 4)
+    SET_FROBENIUS_COEFF(3);
+
+#undef SET_FROBENIUS_COEFF
+
+    // |kFrobeniusCoeffs2[0]| = q^(2 * (P⁰ - 1) / 4) = 1
+    Config::kFrobeniusCoeffs2[0] = FrobeniusCoefficient::One();
+#define SET_FROBENIUS_COEFF2(d)                              \
+  gmp::CopyLimbs(mpz_class(2) * exp##d##_gmp, exp##d.limbs); \
+  Config::kFrobeniusCoeffs2[d] = Config::kNonResidue.Pow(exp##d)
+
+    // |kFrobeniusCoeffs2[1]| = q^(2 * exp₁) = q^(2 * (P¹ - 1) / 4) = ω²
+    SET_FROBENIUS_COEFF2(1);
+    // |kFrobeniusCoeffs2[2]| = q^(2 * exp₂) = q^(2 * (P² - 1) / 4)
+    SET_FROBENIUS_COEFF2(2);
+    // |kFrobeniusCoeffs2[3]| = q^(2 * exp₃) = q^(2 * (P³ - 1) / 4)
+    SET_FROBENIUS_COEFF2(3);
+
+#undef SET_FROBENIUS_COEFF2
+
+    // |kFrobeniusCoeffs3[0]| = q^(3 * (P⁰ - 1) / 4) = 1
+    Config::kFrobeniusCoeffs3[0] = FrobeniusCoefficient::One();
+#define SET_FROBENIUS_COEFF3(d)                              \
+  gmp::CopyLimbs(mpz_class(3) * exp##d##_gmp, exp##d.limbs); \
+  Config::kFrobeniusCoeffs3[d] = Config::kNonResidue.Pow(exp##d)
+
+    // |kFrobeniusCoeffs3[1]| = q^(3 * exp₁) = q^(3 * (P¹ - 1) / 4) = ω³
+    SET_FROBENIUS_COEFF3(1);
+    // |kFrobeniusCoeffs3[2]| = q^(3 * exp₂) = q^(3 * (P² - 1) / 4)
+    SET_FROBENIUS_COEFF3(2);
+    // |kFrobeniusCoeffs3[3]| = q^(3 * exp₃) = q^(3 * (P³ - 1) / 4)
+    SET_FROBENIUS_COEFF3(3);
+
+#undef SET_FROBENIUS_COEFF3
+  }
+};
+
 }  // namespace tachyon::math
 
 #endif  // TACHYON_MATH_FINITE_FIELDS_FP4_H_
diff --git a/tachyon/math/finite_fields/generator/ext_prime_field_generator/build_defs.bzl b/tachyon/math/finite_fields/generator/ext_prime_field_generator/build_defs.bzl
@@ -122,7 +122,6 @@ def generate_fp4s(
     _generate_ext_prime_fields(
         name = name,
         degree = 4,
-        base_field_degree = 2,
         ext_prime_field_deps = ["//tachyon/math/finite_fields:fp4"],
         **kwargs
     )

diff --git a/tachyon/math/finite_fields/generator/ext_prime_field_generator/ext_prime_field_generator.cc b/tachyon/math/finite_fields/generator/ext_prime_field_generator/ext_prime_field_generator.cc
@@ -112,7 +112,8 @@ int GenerationConfig::GenerateConfigHdr() const {
   }
 
   replacements["%{frobenius_coefficient}"] =
-      (degree == 4 || (degree == 6 && base_field_degree == 3) || degree == 12)
+      ((degree == 4 && base_field_degree == 2) ||
+       (degree == 6 && base_field_degree == 3) || degree == 12)
           ? "typename BaseField::BaseField"
           : "BaseField";
 
@@ -122,7 +123,9 @@ int GenerationConfig::GenerateConfigHdr() const {
   std::vector<std::string> tpl_lines = absl::StrSplit(tpl_content, '\n');
 
   RemoveOptionalLines(tpl_lines, "FrobeniusCoefficient2",
-                      degree_over_base_field == 3);
+                      degree_over_base_field >= 3);
+  RemoveOptionalLines(tpl_lines, "FrobeniusCoefficient3",
+                      degree_over_base_field >= 4);
 
   tpl_content = absl::StrJoin(tpl_lines, "\n");
   std::string content = absl::StrReplaceAll(tpl_content, replacements);

diff --git a/tachyon/math/finite_fields/generator/ext_prime_field_generator/fq.h.tpl b/tachyon/math/finite_fields/generator/ext_prime_field_generator/fq.h.tpl
@@ -18,6 +18,9 @@ class %{class}Config {
 %{if FrobeniusCoefficient2}
   static FrobeniusCoefficient kFrobeniusCoeffs2[%{frobenius_coeffs_size}];
 %{endif FrobeniusCoefficient2}
+%{if FrobeniusCoefficient3}
+  static FrobeniusCoefficient kFrobeniusCoeffs3[%{frobenius_coeffs_size}];
+%{endif FrobeniusCoefficient3}
 
   constexpr static bool kNonResidueIsMinusOne = %{non_residue_is_minus_one};
   constexpr static uint64_t kDegreeOverBaseField = %{degree_over_base_field};
@@ -42,6 +45,10 @@ typename %{class}Config<BaseField>::FrobeniusCoefficient %{class}Config<BaseFiel
 template <typename BaseField>
 typename %{class}Config<BaseField>::FrobeniusCoefficient %{class}Config<BaseField>::kFrobeniusCoeffs2[%{frobenius_coeffs_size}];
 %{endif FrobeniusCoefficient2}
+%{if FrobeniusCoefficient3}
+template <typename BaseField>
+typename %{class}Config<BaseField>::FrobeniusCoefficient %{class}Config<BaseField>::kFrobeniusCoeffs3[%{frobenius_coeffs_size}];
+%{endif FrobeniusCoefficient3}
 using %{class} = Fp%{degree}<%{class}Config<%{base_field}>>;
 
 }  // namespace %{namespace}

diff --git a/tachyon/math/finite_fields/koala_bear/BUILD.bazel b/tachyon/math/finite_fields/koala_bear/BUILD.bazel
@@ -1,6 +1,6 @@
 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")
+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")
 
 package(default_visibility = ["//visibility:public"])
@@ -25,6 +25,17 @@ generate_fp2s(
     deps = [":koala_bear"],
 )
 
+generate_fp4s(
+    name = "koala_bear4",
+    base_field = "KoalaBear",
+    base_field_degree = 1,
+    base_field_hdr = "tachyon/math/finite_fields/koala_bear/koala_bear.h",
+    class_name = "KoalaBear4",
+    namespace = "tachyon::math",
+    non_residue = ["3"],
+    deps = [":koala_bear"],
+)
+
 tachyon_cc_library(
     name = "packed_koala_bear",
     hdrs = ["packed_koala_bear.h"],