Skip to content

Commit 40f7aa3

Browse files
chokobolechokobole
authored
Merge pull request #24 from zk-rabbit/fix/point-add-and-doubles
fix: Fix point add and double operations
2 parents 238a5db + 7f084a4 commit 40f7aa3

File tree

8 files changed

+414
-194
lines changed

8 files changed

+414
-194
lines changed

‎third_party/llvm/mlir_runner.patch‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,16 @@
11
diff --git utils/bazel/llvm-project-overlay/mlir/BUILD.bazel utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
2-
index 8559cc12cbaf..7b5078b5a02f 100644
2+
index 8559cc12cbaf..9cb9bb4b0609 100644
33
--- utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
44
+++ utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
5-
@@ -10147,6 +10147,7 @@ cc_binary(
5+
@@ -10109,6 +10109,7 @@ cc_library(
6+
7+
cc_library(
8+
name = "_mlir_c_runner_utils",
9+
+ alwayslink = True,
10+
srcs = [
11+
"lib/ExecutionEngine/CRunnerUtils.cpp",
12+
"lib/ExecutionEngine/SparseTensorRuntime.cpp",
13+
@@ -10147,6 +10148,7 @@ cc_binary(
614

715
cc_library(
816
name = "_mlir_runner_utils",

‎zkir/Dialect/EllipticCurve/Conversions/EllipticCurveToField/PointOperations/Jacobian/Add.cpp‎

Lines changed: 162 additions & 104 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
#include "zkir/Dialect/EllipticCurve/Conversions/EllipticCurveToField/PointOperations/Jacobian/Add.h"
22

33
#include "mlir/Dialect/Arith/IR/Arith.h"
4+
#include "mlir/Dialect/SCF/IR/SCF.h"
45
#include "mlir/Dialect/Tensor/IR/Tensor.h"
56
#include "mlir/IR/ImplicitLocOpBuilder.h"
7+
#include "zkir/Dialect/EllipticCurve/Conversions/EllipticCurveToField/PointOperations/Jacobian/Double.h"
68
#include "zkir/Dialect/EllipticCurve/IR/EllipticCurveTypes.h"
79
#include "zkir/Dialect/Field/IR/FieldOps.h"
810

@@ -12,8 +14,8 @@ namespace mlir::zkir::elliptic_curve {
1214
// http://www.hyperelliptic.org/EFD/g1p/auto-shortw-jacobian-0.html#addition-mmadd-2007-bl
1315
// Cost: 4M + 2S
1416
// Assumption: Z1 == Z2 == 1
15-
Value affineAndAffine(const Value &p1, const Value &p2,
16-
ImplicitLocOpBuilder &b) {
17+
static Value affineAndAffine(const Value &p1, const Value &p2, Type affineType,
18+
ImplicitLocOpBuilder &b) {
1719
Value zero = b.create<arith::ConstantIndexOp>(0);
1820
Value one = b.create<arith::ConstantIndexOp>(1);
1921

@@ -23,43 +25,61 @@ Value affineAndAffine(const Value &p1, const Value &p2,
2325
auto x2 = b.create<tensor::ExtractOp>(p2, zero);
2426
auto y2 = b.create<tensor::ExtractOp>(p2, one);
2527

26-
// H = X2-X1
27-
auto h = b.create<field::SubOp>(x2, x1);
28-
// HH = H²
29-
auto hh = b.create<field::SquareOp>(h);
30-
// I = 4*HH
31-
auto iTmp = b.create<field::DoubleOp>(hh);
32-
auto i = b.create<field::DoubleOp>(iTmp);
33-
// J = H*I
34-
auto j = b.create<field::MulOp>(h, i);
35-
// r = 2*(Y2-Y1)
36-
auto rTmp = b.create<field::SubOp>(y2, y1);
37-
auto r = b.create<field::DoubleOp>(rTmp);
38-
// V = X1*I
39-
auto v = b.create<field::MulOp>(x1, i);
40-
// X3 = r²-J-2*V
41-
auto x3Tmp1 = b.create<field::SquareOp>(r);
42-
auto x3Tmp2 = b.create<field::DoubleOp>(v);
43-
auto x3Tmp3 = b.create<field::SubOp>(x3Tmp1, j);
44-
auto x3 = b.create<field::SubOp>(x3Tmp3, x3Tmp2);
45-
// Y3 = r*(V-X3)-2*Y1*J
46-
auto y3Tmp1 = b.create<field::SubOp>(v, x3);
47-
auto y3Tmp2 = b.create<field::MulOp>(r, y3Tmp1);
48-
auto y3Tmp3 = b.create<field::DoubleOp>(y1);
49-
auto y3Tmp4 = b.create<field::MulOp>(y3Tmp3, j);
50-
auto y3 = b.create<field::SubOp>(y3Tmp2, y3Tmp4);
51-
// Z3 = 2*H
52-
auto z3 = b.create<field::DoubleOp>(h);
53-
54-
return b.create<tensor::FromElementsOp>(SmallVector<Value>({x3, y3, z3}));
28+
// if x1 == x2 && y1 == y2
29+
auto cmpEq1 = b.create<field::CmpOp>(arith::CmpIPredicate::eq, x1, x2);
30+
auto cmpEq2 = b.create<field::CmpOp>(arith::CmpIPredicate::eq, y1, y2);
31+
auto combined_condition = b.create<arith::AndIOp>(cmpEq1, cmpEq2);
32+
auto ifOp = b.create<scf::IfOp>(
33+
combined_condition,
34+
/*thenBuilder=*/
35+
[&](OpBuilder &builder, Location loc) {
36+
ImplicitLocOpBuilder b(loc, builder);
37+
b.create<scf::YieldOp>(jacobianDouble(p1, affineType, b));
38+
},
39+
/*elseBuilder=*/
40+
[&](OpBuilder &builder, Location loc) {
41+
ImplicitLocOpBuilder b(loc, builder);
42+
// H = X2-X1
43+
auto h = b.create<field::SubOp>(x2, x1);
44+
// HH = H²
45+
auto hh = b.create<field::SquareOp>(h);
46+
// I = 4*HH
47+
auto iTmp = b.create<field::DoubleOp>(hh);
48+
auto i = b.create<field::DoubleOp>(iTmp);
49+
// J = H*I
50+
auto j = b.create<field::MulOp>(h, i);
51+
// r = 2*(Y2-Y1)
52+
auto rTmp = b.create<field::SubOp>(y2, y1);
53+
auto r = b.create<field::DoubleOp>(rTmp);
54+
// V = X1*I
55+
auto v = b.create<field::MulOp>(x1, i);
56+
// X3 = r²-J-2*V
57+
auto x3Tmp1 = b.create<field::SquareOp>(r);
58+
auto x3Tmp2 = b.create<field::DoubleOp>(v);
59+
auto x3Tmp3 = b.create<field::SubOp>(x3Tmp1, j);
60+
auto x3 = b.create<field::SubOp>(x3Tmp3, x3Tmp2);
61+
// Y3 = r*(V-X3)-2*Y1*J
62+
auto y3Tmp1 = b.create<field::SubOp>(v, x3);
63+
auto y3Tmp2 = b.create<field::MulOp>(r, y3Tmp1);
64+
auto y3Tmp3 = b.create<field::DoubleOp>(y1);
65+
auto y3Tmp4 = b.create<field::MulOp>(y3Tmp3, j);
66+
auto y3 = b.create<field::SubOp>(y3Tmp2, y3Tmp4);
67+
// Z3 = 2*H
68+
auto z3 = b.create<field::DoubleOp>(h);
69+
70+
auto makePoint =
71+
b.create<tensor::FromElementsOp>(SmallVector<Value>({x3, y3, z3}));
72+
b.create<scf::YieldOp>(ValueRange{makePoint});
73+
});
74+
return ifOp.getResult(0);
5575
}
5676

5777
// madd-2007-bl
5878
// http://www.hyperelliptic.org/EFD/g1p/auto-shortw-jacobian-0.html#addition-madd-2007-bl
5979
// Cost: 7M + 4S
60-
// Assumption: Z1 == 1
61-
Value jacobianAndAffine(const Value &p1, const Value &p2,
62-
ImplicitLocOpBuilder &b) {
80+
// Assumption: Z2 == 1
81+
static Value jacobianAndAffine(const Value &p1, const Value &p2,
82+
Type affineType, ImplicitLocOpBuilder &b) {
6383
Value zero = b.create<arith::ConstantIndexOp>(0);
6484
Value one = b.create<arith::ConstantIndexOp>(1);
6585
Value two = b.create<arith::ConstantIndexOp>(2);
@@ -78,45 +98,64 @@ Value jacobianAndAffine(const Value &p1, const Value &p2,
7898
// S2 = Y2*Z1*Z1Z1
7999
auto s2Tmp = b.create<field::MulOp>(y2, z1);
80100
auto s2 = b.create<field::MulOp>(s2Tmp, z1z1);
81-
// H = U2-X1
82-
auto h = b.create<field::SubOp>(u2, x1);
83-
// HH = H²
84-
auto hh = b.create<field::SquareOp>(h);
85-
// I = 4*HH
86-
auto iTmp = b.create<field::DoubleOp>(hh);
87-
auto i = b.create<field::DoubleOp>(iTmp);
88-
// J = H*I
89-
auto j = b.create<field::MulOp>(h, i);
90-
// r = 2*(S2-Y1)
91-
auto rTmp = b.create<field::SubOp>(s2, y1);
92-
auto r = b.create<field::DoubleOp>(rTmp);
93-
// V = X1*I
94-
auto v = b.create<field::MulOp>(x1, i);
95-
// X3 = r²-J-2*V
96-
auto x3Tmp1 = b.create<field::SquareOp>(r);
97-
auto x3Tmp2 = b.create<field::DoubleOp>(v);
98-
auto x3Tmp3 = b.create<field::SubOp>(x3Tmp1, j);
99-
auto x3 = b.create<field::SubOp>(x3Tmp3, x3Tmp2);
100-
// Y3 = r*(V-X3)-2*Y1*J
101-
auto y3Tmp1 = b.create<field::SubOp>(v, x3);
102-
auto y3Tmp2 = b.create<field::MulOp>(r, y3Tmp1);
103-
auto y3Tmp3 = b.create<field::DoubleOp>(y1);
104-
auto y3Tmp4 = b.create<field::MulOp>(y3Tmp3, j);
105-
auto y3 = b.create<field::SubOp>(y3Tmp2, y3Tmp4);
106-
// Z3 = (Z1+H)²-Z1Z1-HH
107-
auto z3Tmp1 = b.create<field::AddOp>(z1, h);
108-
auto z3Tmp2 = b.create<field::SquareOp>(z3Tmp1);
109-
auto z3Tmp3 = b.create<field::SubOp>(z3Tmp2, z1z1);
110-
auto z3 = b.create<field::SubOp>(z3Tmp3, hh);
111-
112-
return b.create<tensor::FromElementsOp>(SmallVector<Value>({x3, y3, z3}));
101+
102+
// if x1 == u2 && y1 == s2
103+
auto cmpEq1 = b.create<field::CmpOp>(arith::CmpIPredicate::eq, x1, u2);
104+
auto cmpEq2 = b.create<field::CmpOp>(arith::CmpIPredicate::eq, y1, s2);
105+
auto combined_condition = b.create<arith::AndIOp>(cmpEq1, cmpEq2);
106+
auto ifOp = b.create<scf::IfOp>(
107+
combined_condition,
108+
/*thenBuilder=*/
109+
[&](OpBuilder &builder, Location loc) {
110+
ImplicitLocOpBuilder b(loc, builder);
111+
b.create<scf::YieldOp>(jacobianDouble(p2, affineType, b));
112+
},
113+
/*elseBuilder=*/
114+
[&](OpBuilder &builder, Location loc) {
115+
ImplicitLocOpBuilder b(loc, builder);
116+
// H = U2-X1
117+
auto h = b.create<field::SubOp>(u2, x1);
118+
// HH = H²
119+
auto hh = b.create<field::SquareOp>(h);
120+
// I = 4*HH
121+
auto iTmp = b.create<field::DoubleOp>(hh);
122+
auto i = b.create<field::DoubleOp>(iTmp);
123+
// J = H*I
124+
auto j = b.create<field::MulOp>(h, i);
125+
// r = 2*(S2-Y1)
126+
auto rTmp = b.create<field::SubOp>(s2, y1);
127+
auto r = b.create<field::DoubleOp>(rTmp);
128+
// V = X1*I
129+
auto v = b.create<field::MulOp>(x1, i);
130+
// X3 = r²-J-2*V
131+
auto x3Tmp1 = b.create<field::SquareOp>(r);
132+
auto x3Tmp2 = b.create<field::DoubleOp>(v);
133+
auto x3Tmp3 = b.create<field::SubOp>(x3Tmp1, j);
134+
auto x3 = b.create<field::SubOp>(x3Tmp3, x3Tmp2);
135+
// Y3 = r*(V-X3)-2*Y1*J
136+
auto y3Tmp1 = b.create<field::SubOp>(v, x3);
137+
auto y3Tmp2 = b.create<field::MulOp>(r, y3Tmp1);
138+
auto y3Tmp3 = b.create<field::DoubleOp>(y1);
139+
auto y3Tmp4 = b.create<field::MulOp>(y3Tmp3, j);
140+
auto y3 = b.create<field::SubOp>(y3Tmp2, y3Tmp4);
141+
// Z3 = (Z1+H)²-Z1Z1-HH
142+
auto z3Tmp1 = b.create<field::AddOp>(z1, h);
143+
auto z3Tmp2 = b.create<field::SquareOp>(z3Tmp1);
144+
auto z3Tmp3 = b.create<field::SubOp>(z3Tmp2, z1z1);
145+
auto z3 = b.create<field::SubOp>(z3Tmp3, hh);
146+
147+
auto makePoint =
148+
b.create<tensor::FromElementsOp>(SmallVector<Value>({x3, y3, z3}));
149+
b.create<scf::YieldOp>(ValueRange{makePoint});
150+
});
151+
return ifOp.getResult(0);
113152
}
114153

115154
// add-2007-bl
116155
// http://www.hyperelliptic.org/EFD/g1p/auto-shortw-jacobian-0.html#addition-add-2007-bl
117156
// Cost: 11M + 5S
118-
Value jacobianAndJacobian(const Value &p1, const Value &p2,
119-
ImplicitLocOpBuilder &b) {
157+
static Value jacobianAndJacobian(const Value &p1, const Value &p2,
158+
Type jacobianType, ImplicitLocOpBuilder &b) {
120159
Value zero = b.create<arith::ConstantIndexOp>(0);
121160
Value one = b.create<arith::ConstantIndexOp>(1);
122161
Value two = b.create<arith::ConstantIndexOp>(2);
@@ -143,51 +182,70 @@ Value jacobianAndJacobian(const Value &p1, const Value &p2,
143182
// S2 = Y2*Z1*Z1Z1
144183
auto s2Tmp = b.create<field::MulOp>(y2, z1);
145184
auto s2 = b.create<field::MulOp>(s2Tmp, z1z1);
146-
// H = U2-U1
147-
auto h = b.create<field::SubOp>(u2, u1);
148-
// I = (2*H)²
149-
auto iTmp = b.create<field::DoubleOp>(h);
150-
auto i = b.create<field::SquareOp>(iTmp);
151-
// J = -H*I
152-
auto jTmp = b.create<field::NegateOp>(h);
153-
auto j = b.create<field::MulOp>(jTmp, i);
154-
// r = 2*(S2-S1)
155-
auto rTmp = b.create<field::SubOp>(s2, s1);
156-
auto r = b.create<field::DoubleOp>(rTmp);
157-
// V = U1*I
158-
auto v = b.create<field::MulOp>(u1, i);
159-
// X3 = r²+J-2*V
160-
auto x3Tmp1 = b.create<field::SquareOp>(r);
161-
auto x3Tmp2 = b.create<field::AddOp>(x3Tmp1, j);
162-
auto x3Tmp3 = b.create<field::DoubleOp>(v);
163-
auto x3 = b.create<field::SubOp>(x3Tmp2, x3Tmp3);
164-
// Y3 = r*(V-X3)+2*S1*J
165-
auto y3Tmp1 = b.create<field::SubOp>(v, x3);
166-
auto y3Tmp2 = b.create<field::MulOp>(r, y3Tmp1);
167-
auto y3Tmp3 = b.create<field::DoubleOp>(s1);
168-
auto y3Tmp4 = b.create<field::MulOp>(y3Tmp3, j);
169-
auto y3 = b.create<field::AddOp>(y3Tmp2, y3Tmp4);
170-
// Z3 = ((Z1+Z2)²-Z1Z1-Z2Z2)*H
171-
auto z3Tmp1 = b.create<field::AddOp>(z1, z2);
172-
auto z3Tmp2 = b.create<field::SquareOp>(z3Tmp1);
173-
auto z3Tmp3 = b.create<field::SubOp>(z3Tmp2, z1z1);
174-
auto z3Tmp4 = b.create<field::SubOp>(z3Tmp3, z2z2);
175-
auto z3 = b.create<field::MulOp>(z3Tmp4, h);
176-
177-
return b.create<tensor::FromElementsOp>(SmallVector<Value>({x3, y3, z3}));
185+
186+
// if u1 == u2 && s1 == s2
187+
auto cmpEq1 = b.create<field::CmpOp>(arith::CmpIPredicate::eq, u1, u2);
188+
auto cmpEq2 = b.create<field::CmpOp>(arith::CmpIPredicate::eq, s1, s2);
189+
auto combined_condition = b.create<arith::AndIOp>(cmpEq1, cmpEq2);
190+
auto ifOp = b.create<scf::IfOp>(
191+
combined_condition,
192+
/*thenBuilder=*/
193+
[&](OpBuilder &builder, Location loc) {
194+
ImplicitLocOpBuilder b(loc, builder);
195+
b.create<scf::YieldOp>(jacobianDouble(p1, jacobianType, b));
196+
},
197+
/*elseBuilder=*/
198+
[&](OpBuilder &builder, Location loc) {
199+
ImplicitLocOpBuilder b(loc, builder);
200+
// H = U2-U1
201+
auto h = b.create<field::SubOp>(u2, u1);
202+
// I = (2*H)²
203+
auto iTmp = b.create<field::DoubleOp>(h);
204+
auto i = b.create<field::SquareOp>(iTmp);
205+
// J = -H*I
206+
auto jTmp = b.create<field::NegateOp>(h);
207+
auto j = b.create<field::MulOp>(jTmp, i);
208+
// r = 2*(S2-S1)
209+
auto rTmp = b.create<field::SubOp>(s2, s1);
210+
auto r = b.create<field::DoubleOp>(rTmp);
211+
// V = U1*I
212+
auto v = b.create<field::MulOp>(u1, i);
213+
// X3 = r²+J-2*V
214+
auto x3Tmp1 = b.create<field::SquareOp>(r);
215+
auto x3Tmp2 = b.create<field::AddOp>(x3Tmp1, j);
216+
auto x3Tmp3 = b.create<field::DoubleOp>(v);
217+
auto x3 = b.create<field::SubOp>(x3Tmp2, x3Tmp3);
218+
// Y3 = r*(V-X3)+2*S1*J
219+
auto y3Tmp1 = b.create<field::SubOp>(v, x3);
220+
auto y3Tmp2 = b.create<field::MulOp>(r, y3Tmp1);
221+
auto y3Tmp3 = b.create<field::DoubleOp>(s1);
222+
auto y3Tmp4 = b.create<field::MulOp>(y3Tmp3, j);
223+
auto y3 = b.create<field::AddOp>(y3Tmp2, y3Tmp4);
224+
// Z3 = ((Z1+Z2)²-Z1Z1-Z2Z2)*H
225+
auto z3Tmp1 = b.create<field::AddOp>(z1, z2);
226+
auto z3Tmp2 = b.create<field::SquareOp>(z3Tmp1);
227+
auto z3Tmp3 = b.create<field::SubOp>(z3Tmp2, z1z1);
228+
auto z3Tmp4 = b.create<field::SubOp>(z3Tmp3, z2z2);
229+
auto z3 = b.create<field::MulOp>(z3Tmp4, h);
230+
231+
auto makePoint =
232+
b.create<tensor::FromElementsOp>(SmallVector<Value>({x3, y3, z3}));
233+
b.create<scf::YieldOp>(ValueRange{makePoint});
234+
});
235+
return ifOp.getResult(0);
178236
}
179237

180238
Value jacobianAdd(const Value &p1, const Value &p2, Type p1Type, Type p2Type,
181239
ImplicitLocOpBuilder &b) {
182240
if (isa<AffineType>(p1Type)) {
183241
if (isa<AffineType>(p2Type)) {
184-
return affineAndAffine(p1, p2, b);
242+
return affineAndAffine(p1, p2, p1Type, b);
185243
}
186-
return jacobianAndAffine(p2, p1, b);
244+
return jacobianAndAffine(p2, p1, p1Type, b);
187245
} else if (isa<AffineType>(p2Type)) {
188-
return jacobianAndAffine(p1, p2, b);
246+
return jacobianAndAffine(p1, p2, p2Type, b);
189247
} else if (isa<JacobianType>(p1Type) && isa<JacobianType>(p2Type)) {
190-
return jacobianAndJacobian(p1, p2, b);
248+
return jacobianAndJacobian(p1, p2, p1Type, b);
191249
} else {
192250
assert(false && "Unsupported point types for Jacobian addition");
193251
}

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,10 +11,12 @@ cc_library(
1111
"Add.h",
1212
],
1313
deps = [
14+
":Double",
1415
"//zkir/Dialect/EllipticCurve/IR:EllipticCurve",
1516
"//zkir/Dialect/Field/IR:Field",
1617
"@llvm-project//mlir:ArithDialect",
1718
"@llvm-project//mlir:IR",
19+
"@llvm-project//mlir:SCFDialect",
1820
"@llvm-project//mlir:TensorDialect",
1921
],
2022
)
@@ -30,6 +32,7 @@ cc_library(
3032
"//zkir/Dialect/Field/IR:Field",
3133
"@llvm-project//mlir:ArithDialect",
3234
"@llvm-project//mlir:IR",
35+
"@llvm-project//mlir:SCFDialect",
3336
"@llvm-project//mlir:TensorDialect",
3437
],
3538
)

0 commit comments

Comments
 (0)