feat(elliptic_curve): add JacobianA0Double

Jacobian double when A=0

Author: ashjeong

zkir/Dialect/EllipticCurve/Conversions/EllipticCurveToField/PointOperations/Jacobian/BUILD.bazel

@@ -32,6 +32,7 @@ cc_library(
         "//zkir/Dialect/Field/IR:Field",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:TensorDialect",
     ],
 )

zkir/Dialect/EllipticCurve/Conversions/EllipticCurveToField/PointOperations/Jacobian/Double.cpp

@@ -1,6 +1,7 @@
 #include "zkir/Dialect/EllipticCurve/Conversions/EllipticCurveToField/PointOperations/Jacobian/Double.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "zkir/Dialect/EllipticCurve/IR/EllipticCurveTypes.h"
@@ -13,17 +14,14 @@ namespace mlir::zkir::elliptic_curve {
 // https://www.hyperelliptic.org/EFD/g1p/auto-shortw-jacobian.html#doubling-mdbl-2007-bl
 // Cost: 1M + 5S
 // Assumption: Z == 1
-Value affineToJacobianDouble(const Value &point, Type affineType,
+Value affineToJacobianDouble(const Value &point, const Value &a,
                              ImplicitLocOpBuilder &b) {
   Value zero = b.create<arith::ConstantIndexOp>(0);
   Value one = b.create<arith::ConstantIndexOp>(1);
 
   auto x = b.create<tensor::ExtractOp>(point, zero);
   auto y = b.create<tensor::ExtractOp>(point, one);
 
-  field::PrimeFieldType basefield = cast<field::PrimeFieldType>(x.getType());
-  field::PrimeFieldAttr aAttr;
-
   // XX = X²
   auto xx = b.create<field::SquareOp>(x);
   // YY = Y²
@@ -39,8 +37,6 @@ Value affineToJacobianDouble(const Value &point, Type affineType,
   // M = 3*XX+a
   auto mTmp1 = b.create<field::DoubleOp>(xx);
   auto mTmp2 = b.create<field::AddOp>(mTmp1, xx);
-  aAttr = cast<AffineType>(affineType).getCurve().getA();
-  auto a = b.create<field::ConstantOp>(basefield, aAttr.getValue());
   auto m = b.create<field::AddOp>(mTmp2, a);
   // X3 = M²-2*S
   auto x3Tmp1 = b.create<field::SquareOp>(m);
@@ -59,10 +55,56 @@ Value affineToJacobianDouble(const Value &point, Type affineType,
   return b.create<tensor::FromElementsOp>(SmallVector<Value>({x3, y3, z3}));
 }
 
+// dbl-2009-l
+// https://www.hyperelliptic.org/EFD/g1p/auto-shortw-jacobian-0.html#doubling-dbl-2009-l
+// Cost: 2M + 5S
+Value jacobianA0Double(const Value &point, ImplicitLocOpBuilder &b) {
+  Value zero = b.create<arith::ConstantIndexOp>(0);
+  Value one = b.create<arith::ConstantIndexOp>(1);
+  Value two = b.create<arith::ConstantIndexOp>(2);
+
+  auto x = b.create<tensor::ExtractOp>(point, zero);
+  auto y = b.create<tensor::ExtractOp>(point, one);
+  auto z = b.create<tensor::ExtractOp>(point, two);
+
+  // A = X²
+  auto a = b.create<field::SquareOp>(x);
+  // B = Y²
+  auto yy = b.create<field::SquareOp>(y);
+  // C = B²
+  auto c = b.create<field::SquareOp>(yy);
+  // D = 2*((X+B)²-A-C)
+  auto dTmp1 = b.create<field::AddOp>(x, yy);
+  auto dTmp2 = b.create<field::SquareOp>(dTmp1);
+  auto dTmp3 = b.create<field::SubOp>(dTmp2, a);
+  auto dTmp4 = b.create<field::SubOp>(dTmp3, c);
+  auto d = b.create<field::DoubleOp>(dTmp4);
+  // E = 3*A
+  auto eTmp1 = b.create<field::DoubleOp>(a);
+  auto e = b.create<field::AddOp>(eTmp1, a);
+  // F = E²
+  auto f = b.create<field::SquareOp>(e);
+  // X3 = F-2*D
+  auto x3Tmp1 = b.create<field::DoubleOp>(d);
+  auto x3 = b.create<field::SubOp>(f, x3Tmp1);
+  // Y3 = E*(D-X3)-8*C
+  auto y3Tmp1 = b.create<field::SubOp>(d, x3);
+  auto y3Tmp2 = b.create<field::MulOp>(e, y3Tmp1);
+  auto y3Tmp3 = b.create<field::DoubleOp>(c);
+  auto y3Tmp4 = b.create<field::DoubleOp>(y3Tmp3);
+  auto y3Tmp5 = b.create<field::DoubleOp>(y3Tmp4);
+  auto y3 = b.create<field::SubOp>(y3Tmp2, y3Tmp5);
+  // Z3 = 2*Y*Z
+  auto z3Tmp1 = b.create<field::DoubleOp>(y);
+  auto z3 = b.create<field::MulOp>(z3Tmp1, z);
+
+  return b.create<tensor::FromElementsOp>(SmallVector<Value>({x3, y3, z3}));
+}
+
 // dbl-2007-bl
 // https://www.hyperelliptic.org/EFD/g1p/auto-shortw-jacobian.html#doubling-dbl-2007-bl
 // Cost: 1M + 8S + 1*a
-Value jacobianDefaultDouble(const Value &point, Type jacobianType,
+Value jacobianDefaultDouble(const Value &point, const Value &a,
                             ImplicitLocOpBuilder &b) {
   Value zero = b.create<arith::ConstantIndexOp>(0);
   Value one = b.create<arith::ConstantIndexOp>(1);
@@ -72,9 +114,6 @@ Value jacobianDefaultDouble(const Value &point, Type jacobianType,
   auto y = b.create<tensor::ExtractOp>(point, one);
   auto z = b.create<tensor::ExtractOp>(point, two);
 
-  field::PrimeFieldType basefield = cast<field::PrimeFieldType>(x.getType());
-  field::PrimeFieldAttr aAttr;
-
   // XX = X²
   auto xx = b.create<field::SquareOp>(x);
   // YY = Y²
@@ -93,8 +132,6 @@ Value jacobianDefaultDouble(const Value &point, Type jacobianType,
   auto mTmp1 = b.create<field::DoubleOp>(xx);
   auto mTmp2 = b.create<field::AddOp>(mTmp1, xx);
   auto mTmp3 = b.create<field::SquareOp>(zz);
-  aAttr = cast<JacobianType>(jacobianType).getCurve().getA();
-  auto a = b.create<field::ConstantOp>(basefield, aAttr.getValue());
   auto mTmp4 = b.create<field::MulOp>(a, mTmp3);
   auto m = b.create<field::AddOp>(mTmp2, mTmp4);
   // X3 = M²-2*S
@@ -119,10 +156,33 @@ Value jacobianDefaultDouble(const Value &point, Type jacobianType,
 
 Value jacobianDouble(const Value &point, Type inputType,
                      ImplicitLocOpBuilder &b) {
-  if (isa<AffineType>(inputType)) {
-    return affineToJacobianDouble(point, inputType, b);
-  } else if (isa<JacobianType>(inputType)) {
-    return jacobianDefaultDouble(point, inputType, b);
+  auto baseField = cast<field::PrimeFieldType>(
+      cast<RankedTensorType>(point.getType()).getElementType());
+  field::PrimeFieldAttr aAttr;
+
+  if (auto affineType = dyn_cast<AffineType>(inputType)) {
+    aAttr = affineType.getCurve().getA();
+    auto a = b.create<field::ConstantOp>(baseField, aAttr.getValue());
+    return affineToJacobianDouble(point, a, b);
+  } else if (auto jacobianType = dyn_cast<JacobianType>(inputType)) {
+    aAttr = jacobianType.getCurve().getA();
+    auto a = b.create<field::ConstantOp>(baseField, aAttr.getValue());
+    auto zero = b.create<field::ConstantOp>(baseField, 0);
+
+    auto cmpEq = b.create<field::CmpOp>(arith::CmpIPredicate::eq, a, zero);
+    auto ifOp = b.create<scf::IfOp>(
+        cmpEq,
+        /*thenBuilder=*/
+        [&](OpBuilder &builder, Location loc) {
+          ImplicitLocOpBuilder b(loc, builder);
+          b.create<scf::YieldOp>(jacobianA0Double(point, b));
+        },
+        /*elseBuilder=*/
+        [&](OpBuilder &builder, Location loc) {
+          ImplicitLocOpBuilder b(loc, builder);
+          b.create<scf::YieldOp>(jacobianDefaultDouble(point, a, b));
+        });
+    return ifOp.getResult(0);
   } else {
     assert(false && "Unsupported point type for Jacobian doubling");
   }