feat(math): implement QuadraticExtensionField

kroma-network · Sep 25, 2023 · d9a894d · d9a894d
1 parent 9585ee8
commit d9a894d
Show file tree

Hide file tree

Showing 17 changed files with 722 additions and 33 deletions.
diff --git a/tachyon/math/elliptic_curves/bls/bls12_381/BUILD.bazel b/tachyon/math/elliptic_curves/bls/bls12_381/BUILD.bazel
@@ -1,5 +1,6 @@
 load("//tachyon/math/elliptic_curves/short_weierstrass/generator:build_defs.bzl", "generate_ec_points")
-load("//tachyon/math/finite_fields/generator:build_defs.bzl", "generate_prime_fields")
+load("//tachyon/math/finite_fields/generator/prime_field_generator:build_defs.bzl", "generate_prime_fields")
+load("//tachyon/math/finite_fields/generator/quadratic_extension_field_generator:build_defs.bzl", "generate_quadratic_extension_fields")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -71,3 +72,13 @@ generate_ec_points(
     # Hex: 0x8b3f481e3aaa0f1a09e30ed741d8ae4fcf5e095d5d00af600db18cb2c04b3edd03cc744a2888ae40caa232946c5e7e1
     y = "1339506544944476473020471379941921221584933875938349620426543736416511423956333506472724655353366534992391756441569",
 )
+
+generate_quadratic_extension_fields(
+    name = "fp2",
+    base_field = "Fq",
+    base_field_hdr = "tachyon/math/elliptic_curves/bls/bls12_381/fq.h",
+    class_name = "Fp2",
+    namespace = "tachyon::math::bls12_381",
+    non_residue = -1,
+    deps = [":fq"],
+)
diff --git a/tachyon/math/elliptic_curves/bn/bn254/BUILD.bazel b/tachyon/math/elliptic_curves/bn/bn254/BUILD.bazel
@@ -1,7 +1,8 @@
 load("//bazel:tachyon.bzl", "if_polygon_zkevm_backend")
 load("//bazel:tachyon_cc.bzl", "tachyon_cc_library")
 load("//tachyon/math/elliptic_curves/short_weierstrass/generator:build_defs.bzl", "generate_ec_points")
-load("//tachyon/math/finite_fields/generator:build_defs.bzl", "generate_prime_fields")
+load("//tachyon/math/finite_fields/generator/prime_field_generator:build_defs.bzl", "generate_prime_fields")
+load("//tachyon/math/finite_fields/generator/quadratic_extension_field_generator:build_defs.bzl", "generate_quadratic_extension_fields")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -131,3 +132,13 @@ generate_ec_points(
     x = "1",
     y = "2",
 )
+
+generate_quadratic_extension_fields(
+    name = "fp2",
+    base_field = "Fq",
+    base_field_hdr = "tachyon/math/elliptic_curves/bn/bn254/fq.h",
+    class_name = "Fp2",
+    namespace = "tachyon::math::bn254",
+    non_residue = -1,
+    deps = [":fq"],
+)
diff --git a/tachyon/math/elliptic_curves/secp/secp256k1/BUILD.bazel b/tachyon/math/elliptic_curves/secp/secp256k1/BUILD.bazel
@@ -1,7 +1,7 @@
 load("//bazel:tachyon.bzl", "if_polygon_zkevm_backend")
 load("//bazel:tachyon_cc.bzl", "tachyon_cc_library")
 load("//tachyon/math/elliptic_curves/short_weierstrass/generator:build_defs.bzl", "generate_ec_points")
-load("//tachyon/math/finite_fields/generator:build_defs.bzl", "generate_prime_fields")
+load("//tachyon/math/finite_fields/generator/prime_field_generator:build_defs.bzl", "generate_prime_fields")
 
 package(default_visibility = ["//visibility:public"])
 

diff --git a/tachyon/math/elliptic_curves/short_weierstrass/jacobian_point_impl.h b/tachyon/math/elliptic_curves/short_weierstrass/jacobian_point_impl.h
@@ -189,7 +189,7 @@ constexpr CLASS& CLASS::DoubleInPlace() {
     // D = 2 * ((X1 + B)² - A - C)
     //   = 2 * ((X1 + Y1²)² - A - C)
     //   = 2 * 2 * X1 * Y1²
-    uint64_t ext_deg = BaseField::Config::ExtensionDegree();
+    uint64_t ext_deg = BaseField::ExtensionDegree();
     BaseField d;
     if (ext_deg == 1 || ext_deg == 2) {
       d = x_;

diff --git a/tachyon/math/finite_fields/BUILD.bazel b/tachyon/math/finite_fields/BUILD.bazel
@@ -119,12 +119,14 @@ tachyon_cc_test(
         "modulus_unittest.cc",
         "prime_field_base_unittest.cc",
         "prime_field_unittest.cc",
+        "quadratic_extension_field_unittest.cc",
     ],
     deps = [
         "//tachyon/base:bits",
         "//tachyon/math/elliptic_curves/bn/bn254:fq",
         "//tachyon/math/elliptic_curves/bn/bn254:fr",
         "//tachyon/math/finite_fields/test:gf7",
+        "//tachyon/math/finite_fields/test:gf7_quad_ext",
     ],
 )
 
@@ -178,3 +180,12 @@ tachyon_cc_benchmark(
         "//tachyon/math/finite_fields/goldilocks_prime:goldilocks",
     ],
 )
+
+tachyon_cc_library(
+    name = "quadratic_extension_field",
+    hdrs = ["quadratic_extension_field.h"],
+    deps = [
+        "//tachyon/math/base:field",
+        "@com_google_absl//absl/strings",
+    ],
+)
diff --git a/tachyon/math/finite_fields/generator/BUILD.bazel b/tachyon/math/finite_fields/generator/BUILD.bazel
@@ -1,4 +1,4 @@
-load("//bazel:tachyon_cc.bzl", "tachyon_cc_binary", "tachyon_cc_library")
+load("//bazel:tachyon_cc.bzl", "tachyon_cc_library")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -12,18 +12,3 @@ tachyon_cc_library(
         "//tachyon/math/finite_fields:modulus",
     ],
 )
-
-tachyon_cc_binary(
-    name = "generator",
-    srcs = ["generator.cc"],
-    deps = [
-        ":generator_util",
-        "//tachyon/base/console",
-        "//tachyon/base/files:file_path_flag",
-        "//tachyon/base/flag:flag_parser",
-        "//tachyon/build:cc_writer",
-        "//tachyon/math/base:bit_iterator",
-        "//tachyon/math/base/gmp:bit_traits",
-        "//tachyon/math/finite_fields:prime_field_util",
-    ],
-)
diff --git a/tachyon/math/finite_fields/generator/prime_field_generator/BUILD.bazel b/tachyon/math/finite_fields/generator/prime_field_generator/BUILD.bazel
@@ -0,0 +1,18 @@
+load("//bazel:tachyon_cc.bzl", "tachyon_cc_binary")
+
+package(default_visibility = ["//visibility:public"])
+
+tachyon_cc_binary(
+    name = "prime_field_generator",
+    srcs = ["prime_field_generator.cc"],
+    deps = [
+        "//tachyon/base/console",
+        "//tachyon/base/files:file_path_flag",
+        "//tachyon/base/flag:flag_parser",
+        "//tachyon/build:cc_writer",
+        "//tachyon/math/base:bit_iterator",
+        "//tachyon/math/base/gmp:bit_traits",
+        "//tachyon/math/finite_fields:prime_field_util",
+        "//tachyon/math/finite_fields/generator:generator_util",
+    ],
+)
diff --git a/...th/finite_fields/generator/build_defs.bzl → ...ator/prime_field_generator/build_defs.bzl b/...th/finite_fields/generator/build_defs.bzl → ...ator/prime_field_generator/build_defs.bzl
@@ -50,7 +50,7 @@ generate_prime_field = rule(
             cfg = "target",
             executable = True,
             allow_single_file = True,
-            default = Label("@kroma_network_tachyon//tachyon/math/finite_fields/generator"),
+            default = Label("@kroma_network_tachyon//tachyon/math/finite_fields/generator/prime_field_generator"),
         ),
     },
 )

diff --git a/...math/finite_fields/generator/generator.cc → ..._field_generator/prime_field_generator.cc b/...math/finite_fields/generator/generator.cc → ..._field_generator/prime_field_generator.cc
@@ -120,8 +120,6 @@ int GenerationConfig::GenerateConfigHdr() const {
       "",
       "  constexpr static bool kHasLargeSubgroupRootOfUnity = false;",
       "",
-      "  constexpr static uint64_t ExtensionDegree() { return 1; }",
-      "",
       "  static void Init();",
       "};",
       "",
@@ -322,7 +320,7 @@ int GenerationConfig::GenerateConfigGpuHdr() const {
 
 int RealMain(int argc, char** argv) {
   GenerationConfig config;
-  config.generator = "//tachyon/math/finite_fields/generator";
+  config.generator = "//tachyon/math/finite_fields/prime_field_field_generator";
 
   base::FlagParser parser;
   parser.AddFlag<base::FilePathFlag>(&config.out)
@@ -331,9 +329,7 @@ int RealMain(int argc, char** argv) {
   parser.AddFlag<base::StringFlag>(&config.ns_name)
       .set_long_name("--namespace")
       .set_required();
-  parser.AddFlag<base::StringFlag>(&config.class_name)
-      .set_long_name("--class")
-      .set_required();
+  parser.AddFlag<base::StringFlag>(&config.class_name).set_long_name("--class");
   parser.AddFlag<base::StringFlag>(&config.modulus)
       .set_long_name("--modulus")
       .set_required();

diff --git a/tachyon/math/finite_fields/generator/quadratic_extension_field_generator/BUILD.bazel b/tachyon/math/finite_fields/generator/quadratic_extension_field_generator/BUILD.bazel
@@ -0,0 +1,16 @@
+load("//bazel:tachyon_cc.bzl", "tachyon_cc_binary")
+
+package(default_visibility = ["//visibility:public"])
+
+tachyon_cc_binary(
+    name = "quadratic_extension_field_generator",
+    srcs = ["quadratic_extension_field_generator.cc"],
+    deps = [
+        "//tachyon/base/console",
+        "//tachyon/base/files:file_path_flag",
+        "//tachyon/base/flag:flag_parser",
+        "//tachyon/build:cc_writer",
+        "//tachyon/math/base:big_int",
+        "//tachyon/math/base:bit_iterator",
+    ],
+)
diff --git a/tachyon/math/finite_fields/generator/quadratic_extension_field_generator/build_defs.bzl b/tachyon/math/finite_fields/generator/quadratic_extension_field_generator/build_defs.bzl
@@ -0,0 +1,70 @@
+load("//bazel:tachyon_cc.bzl", "tachyon_cc_library")
+
+def _generate_quadratic_extension_field_impl(ctx):
+    arguments = [
+        "--out=%s" % (ctx.outputs.out.path),
+        "--namespace=%s" % (ctx.attr.namespace),
+        "--class=%s" % (ctx.attr.class_name),
+        "--non_residue=%s" % (ctx.attr.non_residue),
+        "--base_field_hdr=%s" % (ctx.attr.base_field_hdr),
+        "--base_field=%s" % (ctx.attr.base_field),
+    ]
+
+    ctx.actions.run(
+        tools = [ctx.executable._tool],
+        executable = ctx.executable._tool,
+        outputs = [ctx.outputs.out],
+        arguments = arguments,
+    )
+
+    return [DefaultInfo(files = depset([ctx.outputs.out]))]
+
+generate_quadratic_extension_field = rule(
+    implementation = _generate_quadratic_extension_field_impl,
+    attrs = {
+        "out": attr.output(mandatory = True),
+        "namespace": attr.string(mandatory = True),
+        "class_name": attr.string(mandatory = True),
+        "non_residue": attr.int(mandatory = True),
+        "base_field_hdr": attr.string(mandatory = True),
+        "base_field": attr.string(mandatory = True),
+        "_tool": attr.label(
+            # TODO(chokobole): Change it to "exec" we can build it on macos.
+            cfg = "target",
+            executable = True,
+            allow_single_file = True,
+            default = Label("@kroma_network_tachyon//tachyon/math/finite_fields/generator/quadratic_extension_field_generator"),
+        ),
+    },
+)
+
+def generate_quadratic_extension_fields(
+        name,
+        namespace,
+        class_name,
+        non_residue,
+        base_field_hdr = "",
+        base_field = "",
+        deps = [],
+        **kwargs):
+    for n in [
+        ("{}_gen_hdr".format(name), "{}.h".format(name)),
+    ]:
+        generate_quadratic_extension_field(
+            namespace = namespace,
+            class_name = class_name,
+            non_residue = non_residue,
+            base_field_hdr = base_field_hdr,
+            base_field = base_field,
+            name = n[0],
+            out = n[1],
+        )
+
+    tachyon_cc_library(
+        name = name,
+        hdrs = [":{}_gen_hdr".format(name)],
+        deps = deps + [
+            "//tachyon/math/finite_fields:quadratic_extension_field",
+        ],
+        **kwargs
+    )