perf: create constexpr MulByA() and kAIsZero for faster computation

`CurveConfig::A()` and `CurveConfig::B()` are used to check whether
point is on a given curve and group operation like addition and
doubling. `IsOnCurve()` is not a critical path in terms of performance.
But addition and doubling is. But when seeing `CurveConfig::A()` for
bn254 and bls12_381, they are zero! We can skip some of instruction if
we can determine at compile time. Plus, using `ScalarMul()` in
semigroups, it needs to iterate all the bits, if you know `A()` in
advance, you can make a fater multiplication code from generator.

After all these, I found that there's no point to differentiate
`SWCurveGpu` from `SWCurve`. So this commit also incorporates those
twos.

tachyon/c/math/elliptic_curves/msm/msm_gpu.h

@@ -12,15 +12,14 @@
 #include "tachyon/math/elliptic_curves/affine_point.h"
 #include "tachyon/math/elliptic_curves/msm/variable_base_msm_gpu.h"
 #include "tachyon/math/elliptic_curves/point_conversions.h"
-#include "tachyon/math/elliptic_curves/short_weierstrass/sw_curve_traits.h"
 
 namespace tachyon::c::math {
 
 template <typename GpuCurve>
 struct MSMGpuApi {
   using GpuAffinePointTy = tachyon::math::AffinePoint<GpuCurve>;
   using GpuScalarField = typename GpuAffinePointTy::ScalarField;
-  using CpuCurve = typename tachyon::math::SWCurveTraits<GpuCurve>::CpuCurve;
+  using CpuCurve = typename GpuCurve::CpuCurve;
   using CpuAffinePointTy = tachyon::math::AffinePoint<CpuCurve>;
 
   tachyon::device::gpu::ScopedMemPool mem_pool;
@@ -91,8 +90,7 @@ template <typename RetPointTy, typename GpuCurve, typename CPointTy,
           typename CScalarField,
           typename CRetPointTy =
               typename cc::math::PointTraits<RetPointTy>::CCurvePointTy,
-          typename CpuCurve =
-              typename tachyon::math::SWCurveTraits<GpuCurve>::CpuCurve>
+          typename CpuCurve = typename GpuCurve::CpuCurve>
 CRetPointTy* DoMSMGpu(MSMGpuApi<GpuCurve>& msm_api, const CPointTy* bases,
                       const CScalarField* scalars, size_t size) {
   msm_api.provider.Inject(bases, scalars, size);

tachyon/math/elliptic_curves/msm/algorithms/BUILD.bazel

@@ -8,6 +8,5 @@ tachyon_cc_library(
     deps = [
         "//tachyon/device/gpu:gpu_memory",
         "//tachyon/math/elliptic_curves:points",
-        "//tachyon/math/elliptic_curves/short_weierstrass:sw_curve_traits",
     ],
 )

tachyon/math/elliptic_curves/msm/algorithms/bellman/bellman_msm.h

@@ -4,7 +4,6 @@
 #include "tachyon/math/elliptic_curves/msm/algorithms/bellman/bellman_msm_impl.h"
 #include "tachyon/math/elliptic_curves/msm/algorithms/msm_algorithm.h"
 #include "tachyon/math/elliptic_curves/msm/algorithms/pippenger/pippenger_base.h"
-#include "tachyon/math/elliptic_curves/short_weierstrass/sw_curve_traits.h"
 
 namespace tachyon::math {
 
@@ -15,7 +14,7 @@ class BellmanMSM : public PippengerBase<AffinePoint<GpuCurve>>,
   using BaseField = typename AffinePoint<GpuCurve>::BaseField;
   using ScalarField = typename AffinePoint<GpuCurve>::ScalarField;
   using Bucket = typename PippengerBase<GpuCurve>::Bucket;
-  using CpuCurve = typename SWCurveTraits<GpuCurve>::CpuCurve;
+  using CpuCurve = typename GpuCurve::CpuCurve;
 
   BellmanMSM() : BellmanMSM(nullptr, nullptr) {}
   BellmanMSM(gpuMemPool_t mem_pool, gpuStream_t stream)

tachyon/math/elliptic_curves/msm/algorithms/cuzk/cuzk.h

@@ -11,7 +11,6 @@
 #include "tachyon/math/elliptic_curves/msm/algorithms/pippenger/pippenger_base.h"
 #include "tachyon/math/elliptic_curves/msm/algorithms/pippenger/pippenger_ctx.h"
 #include "tachyon/math/elliptic_curves/msm/kernels/cuzk/cuzk_kernels.cu.h"
-#include "tachyon/math/elliptic_curves/short_weierstrass/sw_curve_traits.h"
 
 namespace tachyon::math {
 
@@ -22,7 +21,7 @@ class CUZK : public PippengerBase<AffinePoint<GpuCurve>>,
   using BaseField = typename AffinePoint<GpuCurve>::BaseField;
   using ScalarField = typename AffinePoint<GpuCurve>::ScalarField;
   using Bucket = typename PippengerBase<GpuCurve>::Bucket;
-  using CpuCurve = typename SWCurveTraits<GpuCurve>::CpuCurve;
+  using CpuCurve = typename GpuCurve::CpuCurve;
 
   CUZK() : CUZK(nullptr, nullptr) {}
   CUZK(gpuMemPool_t mem_pool, gpuStream_t stream)

tachyon/math/elliptic_curves/msm/algorithms/msm_algorithm.h

@@ -4,7 +4,6 @@
 #include "tachyon/device/gpu/gpu_memory.h"
 #include "tachyon/math/elliptic_curves/affine_point.h"
 #include "tachyon/math/elliptic_curves/jacobian_point.h"
-#include "tachyon/math/elliptic_curves/short_weierstrass/sw_curve_traits.h"
 
 namespace tachyon::math {
 
@@ -18,7 +17,7 @@ template <typename GpuCurve>
 class MSMGpuAlgorithm {
  public:
   using ScalarField = typename JacobianPoint<GpuCurve>::ScalarField;
-  using CpuCurve = typename SWCurveTraits<GpuCurve>::CpuCurve;
+  using CpuCurve = typename GpuCurve::CpuCurve;
 
   virtual bool Run(const device::gpu::GpuMemory<AffinePoint<GpuCurve>>& bases,
                    const device::gpu::GpuMemory<ScalarField>& scalars,

tachyon/math/elliptic_curves/msm/variable_base_msm_gpu.h

@@ -10,7 +10,7 @@ template <typename GpuCurve>
 class VariableBaseMSMGpu {
  public:
   using ScalarField = typename JacobianPoint<GpuCurve>::ScalarField;
-  using CpuCurve = typename SWCurveTraits<GpuCurve>::CpuCurve;
+  using CpuCurve = typename GpuCurve::CpuCurve;
 
   VariableBaseMSMGpu(MSMAlgorithmKind kind, gpuMemPool_t mem_pool,
                      gpuStream_t stream) {

tachyon/math/elliptic_curves/short_weierstrass/BUILD.bazel

@@ -33,30 +33,14 @@ tachyon_cc_library(
     name = "sw_curve",
     srcs = ["sw_curve.h"],
     deps = [
-        ":sw_curve_base",
-        "//tachyon/base:static_storage",
-        "//tachyon/math/elliptic_curves:points",
-    ],
-)
-
-tachyon_cc_library(
-    name = "sw_curve_base",
-    hdrs = ["sw_curve_base.h"],
-    deps = [":sw_curve_traits"],
-)
-
-tachyon_cc_library(
-    name = "sw_curve_gpu",
-    srcs = ["sw_curve_gpu.h"],
-    deps = [
-        ":sw_curve_base",
+        ":sw_curve_traits",
         "//tachyon/math/elliptic_curves:points",
     ],
 )
 
 tachyon_cc_library(
     name = "sw_curve_traits",
-    hdrs = ["sw_curve_traits.h"],
+    srcs = ["sw_curve_traits.h"],
 )
 
 tachyon_cc_test(

tachyon/math/elliptic_curves/short_weierstrass/affine_point_impl.h

@@ -21,7 +21,9 @@ constexpr ProjectivePoint<Curve> CLASS::DoubleProjective() const {
   // w = a + 3 * XX
   BaseField w = xx;
   w += w.Double();
-  w += Curve::A();
+  if constexpr (!Curve::Config::kAIsZero) {
+    w += Curve::A();
+  }
 
   // Y1Y1 = Y1²
   BaseField y1y1 = y_;
@@ -96,7 +98,9 @@ constexpr PointXYZZ<Curve> CLASS::DoubleXYZZ() const {
   BaseField m = x_;
   m.SquareInPlace();
   m += m.Double();
-  m += Curve::A();
+  if constexpr (!Curve::Config::kAIsZero) {
+    m += Curve::A();
+  }
 
   // X3 = M² - 2 * S
   BaseField x = m;

tachyon/math/elliptic_curves/short_weierstrass/generator/build_defs.bzl

@@ -109,7 +109,6 @@ def generate_ec_points(
             ":{}".format(name),
             ":fq_gpu",
             ":fr_gpu",
-            "//tachyon/math/elliptic_curves/short_weierstrass:sw_curve_gpu",
         ],
         **kwargs
     )

tachyon/math/elliptic_curves/short_weierstrass/generator/generator.cc

@@ -57,6 +57,8 @@ int GenerationConfig::GenerateConfigHdr() const {
       "  using CpuCurveConfig = %{class}CurveConfig<CpuBaseField, CpuScalarField>;",
       "  using GpuCurveConfig = %{class}CurveConfig<GpuBaseField, GpuScalarField>;",
       "",
+      "  constexpr static bool kAIsZero = %{a_is_zero};",
+      "",
       "  constexpr static BigInt<%{fq_n}> kA = BigInt<%{fq_n}>({",
       "    %{a_mont_form}",
       "  });",
@@ -89,6 +91,10 @@ int GenerationConfig::GenerateConfigHdr() const {
       "  constexpr static BigInt<%{fr_n}> kGLVCoeff11 = BigInt<%{fr_n}>({",
       "    %{glv_coeff_11_mont_form}",
       "  });",
+      "",
+      "  constexpr static BaseField MulByA(const BaseField& v) {",
+      "    return %{mul_by_a};",
+      "  }",
       "};",
       "",
       "using %{class}Curve = SWCurve<%{class}CurveConfig<Fq, Fr>>;",
@@ -127,18 +133,26 @@ int GenerationConfig::GenerateConfigHdr() const {
   mpz_class x = math::gmp::FromDecString(this->x);
   mpz_class y = math::gmp::FromDecString(this->y);
 
+  std::string mul_by_a;
+  if (a == mpz_class(0)) {
+    mul_by_a = "BaseField::Zero()";
+  } else {
+    mul_by_a = math::GenerateFastMultiplication(a.get_si());
+  }
+
   size_t fq_n = math::gmp::GetLimbSize(fq_modulus);
 
-  std::map<std::string, std::string> replacements = {{
-      {"%{header_dir_name}", GetHdrPath().DirName().value()},
-      {"%{namespace}", ns_name},
-      {"%{class}", class_name},
-      {"%{fq_n}", base::NumberToString(fq_n)},
-      {"%{a_mont_form}", math::MpzClassToMontString(a, fq_modulus)},
-      {"%{b_mont_form}", math::MpzClassToMontString(b, fq_modulus)},
-      {"%{x_mont_form}", math::MpzClassToMontString(x, fq_modulus)},
-      {"%{y_mont_form}", math::MpzClassToMontString(y, fq_modulus)},
-  }};
+  std::map<std::string, std::string> replacements = {
+      {{"%{header_dir_name}", GetHdrPath().DirName().value()},
+       {"%{namespace}", ns_name},
+       {"%{class}", class_name},
+       {"%{fq_n}", base::NumberToString(fq_n)},
+       {"%{a_is_zero}", base::BoolToString(a == mpz_class(0))},
+       {"%{a_mont_form}", math::MpzClassToMontString(a, fq_modulus)},
+       {"%{b_mont_form}", math::MpzClassToMontString(b, fq_modulus)},
+       {"%{x_mont_form}", math::MpzClassToMontString(x, fq_modulus)},
+       {"%{y_mont_form}", math::MpzClassToMontString(y, fq_modulus)},
+       {"%{mul_by_a}", mul_by_a}}};
 
   if (!glv_coefficients.empty()) {
     mpz_class fr_modulus = math::gmp::FromDecString(this->fr_modulus);
@@ -178,11 +192,10 @@ int GenerationConfig::GenerateConfigGpuHdr() const {
       "#include \"%{header_dir_name}/fq_gpu.h\"",
       "#include \"%{header_dir_name}/fr_gpu.h\"",
       "#include \"%{header_path}\"",
-      "#include \"tachyon/math/elliptic_curves/short_weierstrass/sw_curve_gpu.h\"",
       "",
       "namespace %{namespace} {",
       "",
-      "using %{class}CurveGpu = SWCurveGpu<%{class}CurveConfig<FqGpu, FrGpu>>;",
+      "using %{class}CurveGpu = SWCurve<%{class}CurveConfig<FqGpu, FrGpu>>;",
       "using %{class}AffinePointGpu = AffinePoint<%{class}CurveGpu>;",
       "using %{class}ProjectivePointGpu = ProjectivePoint<%{class}CurveGpu>;",
       "using %{class}JacobianPointGpu = JacobianPoint<%{class}CurveGpu>;",

tachyon/math/elliptic_curves/short_weierstrass/jacobian_point_impl.h

@@ -172,7 +172,7 @@ constexpr CLASS& CLASS::DoubleInPlace() {
     return *this;
   }
 
-  if (Curve::A().IsZero()) {
+  if constexpr (Curve::Config::kAIsZero) {
     // http://www.hyperelliptic.org/EFD/g1p/auto-shortw-jacobian-0.html#doubling-dbl-2009-l
     // A = X1²
     BaseField a = x_;
@@ -245,7 +245,9 @@ constexpr CLASS& CLASS::DoubleInPlace() {
     BaseField m = xx;
     m.DoubleInPlace();
     m += xx;
-    m += Curve::A() * zz.Square();
+    if constexpr (!Curve::Config::kAIsZero) {
+      m += Curve::Config::MulByA(zz.Square());
+    }
 
     // T = M² - 2 * S
     // X3 = T

tachyon/math/elliptic_curves/short_weierstrass/point_xyzz_impl.h

@@ -170,7 +170,9 @@ constexpr CLASS& CLASS::DoubleInPlace() {
   BaseField m = std::move(x_);
   m.SquareInPlace();
   m += m.Double();
-  m += Curve::A() * zz_.Square();
+  if constexpr (!Curve::Config::kAIsZero) {
+    m += Curve::Config::MulByA(zz_.Square());
+  }
 
   // X3 = M² - 2 * S
   x_ = m;

tachyon/math/elliptic_curves/short_weierstrass/projective_point_impl.h

@@ -162,7 +162,9 @@ constexpr CLASS& CLASS::DoubleInPlace() {
   // w = a * ZZ + 3 * XX
   BaseField w = xx;
   w += w.Double();
-  w += Curve::A() * zz;
+  if constexpr (!Curve::Config::kAIsZero) {
+    w += Curve::Config::MulByA(zz);
+  }
 
   // s = 2 * Y1 * Z1
   BaseField s = y_;

tachyon/math/elliptic_curves/short_weierstrass/sw_curve.h

@@ -1,54 +1,103 @@
 #ifndef TACHYON_MATH_ELLIPTIC_CURVES_SHORT_WEIERSTRASS_SW_CURVE_H_
 #define TACHYON_MATH_ELLIPTIC_CURVES_SHORT_WEIERSTRASS_SW_CURVE_H_
 
-#include "tachyon/base/static_storage.h"
-#include "tachyon/math/elliptic_curves/short_weierstrass/sw_curve_base.h"
+#include "tachyon/math/elliptic_curves/affine_point.h"
+#include "tachyon/math/elliptic_curves/jacobian_point.h"
+#include "tachyon/math/elliptic_curves/point_xyzz.h"
+#include "tachyon/math/elliptic_curves/projective_point.h"
 #include "tachyon/math/elliptic_curves/short_weierstrass/sw_curve_traits.h"
 
 namespace tachyon::math {
 
-template <typename Config>
-class SWCurveGpu;
-
-template <typename _Config>
-class SWCurve : public SWCurveBase<SWCurve<_Config>> {
+// Config for Short Weierstrass model.
+// See https://www.hyperelliptic.org/EFD/g1p/auto-shortw.html for more details.
+// This config represents `y² = x³ + a * x + b`, where `a` and `b` are
+// constants.
+template <typename SWCurveConfig>
+class SWCurve {
  public:
-  using Config = _Config;
+  using Config = SWCurveConfig;
+
   using BaseField = typename Config::BaseField;
   using ScalarField = typename Config::ScalarField;
-  using AffinePointTy = AffinePoint<SWCurve<Config>>;
-  using ProjectivePointTy = ProjectivePoint<SWCurve<Config>>;
-  using JacobianPointTy = JacobianPoint<SWCurve<Config>>;
+  using AffinePointTy = typename SWCurveTraits<Config>::AffinePointTy;
+  using ProjectivePointTy = typename SWCurveTraits<Config>::ProjectivePointTy;
+  using JacobianPointTy = typename SWCurveTraits<Config>::JacobianPointTy;
+  using PointXYZZTy = typename SWCurveTraits<Config>::PointXYZZTy;
+
+  using CpuCurve = SWCurve<typename Config::CpuCurveConfig>;
+  using GpuCurve = SWCurve<typename Config::GpuCurveConfig>;
 
   constexpr static bool kIsSWCurve = true;
 
-  DEFINE_STATIC_STORAGE_TEMPLATE_METHOD(BaseField, A)
-  DEFINE_STATIC_STORAGE_TEMPLATE_METHOD(BaseField, B)
-  DEFINE_STATIC_STORAGE_TEMPLATE_METHOD(JacobianPointTy, Generator)
+  constexpr static BaseField A() {
+    return BaseField::FromMontgomery(Config::kA);
+  }
+
+  constexpr static BaseField B() {
+    return BaseField::FromMontgomery(Config::kB);
+  }
+
+  constexpr static JacobianPointTy Generator() {
+    return JacobianPointTy(BaseField::FromMontgomery(Config::kGenerator.x),
+                           BaseField::FromMontgomery(Config::kGenerator.y),
+                           BaseField::One());
+  }
 
   static void Init() {
     BaseField::Init();
     ScalarField::Init();
+  }
 
-    A() = BaseField::FromMontgomery(Config::kA);
-    B() = BaseField::FromMontgomery(Config::kB);
-    Generator() = JacobianPointTy(
-        BaseField::FromMontgomery(Config::kGenerator.x),
-        BaseField::FromMontgomery(Config::kGenerator.y), BaseField::One());
+  constexpr static bool IsOnCurve(const AffinePointTy& point) {
+    if (point.infinity()) return false;
+    BaseField right = point.x().Square() * point.x() + B();
+    if constexpr (!Config::kAIsZero) {
+      right += A() * point.x();
+    }
+    return point.y().Square() == right;
+  }
+
+  constexpr static bool IsOnCurve(const ProjectivePointTy& point) {
+    if (point.z().IsZero()) return false;
+    BaseField z2 = point.z().Square();
+    BaseField z3 = z2 * point.z();
+    BaseField right = point.x().Square() * point.x() + B() * z3;
+    if constexpr (!Config::kAIsZero) {
+      right += A() * point.x() * z2;
+    }
+    return point.y().Square() * point.z() == right;
+  }
+
+  constexpr static bool IsOnCurve(const JacobianPointTy& point) {
+    if (point.z().IsZero()) return false;
+    BaseField z3 = point.z().Square() * point.z();
+    BaseField right = point.x().Square() * point.x() + B() * z3.Square();
+    if constexpr (!Config::kAIsZero) {
+      right += A() * point.x() * z3 * point.z();
+    }
+    return point.y().Square() == right;
+  }
+
+  constexpr static bool IsOnCurve(const PointXYZZTy& point) {
+    if (point.zzz().IsZero()) return false;
+    BaseField right =
+        point.x().Square() * point.x() + B() * point.zzz().Square();
+    if constexpr (!Config::kAIsZero) {
+      right += A() * point.x() * point.zz().Square();
+    }
+    return point.y().Square() == right;
   }
 };
 
 template <typename Config>
-struct SWCurveTraits<SWCurve<Config>> {
+struct SWCurveTraits {
   using BaseField = typename Config::BaseField;
   using ScalarField = typename Config::ScalarField;
   using AffinePointTy = AffinePoint<SWCurve<Config>>;
   using ProjectivePointTy = ProjectivePoint<SWCurve<Config>>;
   using JacobianPointTy = JacobianPoint<SWCurve<Config>>;
   using PointXYZZTy = PointXYZZ<SWCurve<Config>>;
-
-  using CpuCurve = SWCurve<typename Config::CpuCurveConfig>;
-  using GpuCurve = SWCurveGpu<typename Config::GpuCurveConfig>;
 };
 
 }  // namespace tachyon::math