feat: add IsZeroOp

Author: ashjeong

zkir/Dialect/EllipticCurve/Conversions/EllipticCurveToField/EllipticCurveToField.cpp

@@ -147,6 +147,37 @@ struct ConvertPoint : public OpConversionPattern<PointOp> {
   }
 };
 
+struct ConvertIsZero : public OpConversionPattern<IsZeroOp> {
+  explicit ConvertIsZero(MLIRContext *context)
+      : OpConversionPattern<IsZeroOp>(context) {}
+
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      IsZeroOp op, OneToNOpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    ImplicitLocOpBuilder b(op.getLoc(), rewriter);
+
+    ValueRange coords = adaptor.getInput();
+    field::PrimeFieldType baseField =
+        cast<field::PrimeFieldType>(coords[0].getType());
+    Value zeroPF = b.create<field::ConstantOp>(baseField, 0);
+
+    Value cmp;
+    if (isa<AffineType>(op.getInput().getType())) {
+      Value xIsZero =
+          b.create<field::CmpOp>(arith::CmpIPredicate::eq, coords[0], zeroPF);
+      Value yIsZero =
+          b.create<field::CmpOp>(arith::CmpIPredicate::eq, coords[1], zeroPF);
+      cmp = b.create<arith::AndIOp>(xIsZero, yIsZero);
+    } else {
+      cmp = b.create<field::CmpOp>(arith::CmpIPredicate::eq, coords[2], zeroPF);
+    }
+    rewriter.replaceOp(op, cmp);
+    return success();
+  }
+};
+
 struct ConvertExtract : public OpConversionPattern<ExtractOp> {
   explicit ConvertExtract(MLIRContext *context)
       : OpConversionPattern<ExtractOp>(context) {}
@@ -554,10 +585,11 @@ void EllipticCurveToField::runOnOperation() {
 
   RewritePatternSet patterns(context);
   rewrites::populateWithGenerated(patterns);
-  patterns.add<ConvertPoint, ConvertExtract, ConvertConvertPointType,
-               ConvertAdd, ConvertDouble, ConvertNegate, ConvertSub,
-               ConvertScalarMul, ConvertMSM, ConvertAny<tensor::FromElementsOp>,
-               ConvertAny<tensor::ExtractOp>>(typeConverter, context);
+  patterns
+      .add<ConvertPoint, ConvertIsZero, ConvertExtract, ConvertConvertPointType,
+           ConvertAdd, ConvertDouble, ConvertNegate, ConvertSub,
+           ConvertScalarMul, ConvertMSM, ConvertAny<tensor::FromElementsOp>,
+           ConvertAny<tensor::ExtractOp>>(typeConverter, context);
   target.addDynamicallyLegalOp<tensor::FromElementsOp, tensor::ExtractOp>(
       [&](auto op) { return typeConverter.isLegal(op); });
 

zkir/Dialect/EllipticCurve/IR/BUILD.bazel

@@ -23,6 +23,7 @@ cc_library(
         ":types_inc_gen",
         "//zkir/Dialect/Field/IR:Field",
         "//zkir/Dialect/Poly/IR:Poly",
+        "//zkir/Utils:OpUtils",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:InferTypeOpInterface",

zkir/Dialect/EllipticCurve/IR/EllipticCurveDialect.cpp

@@ -179,6 +179,15 @@ LogicalResult PointOp::verify() {
 
 /////////////// VERIFY OPS /////////////////
 
+LogicalResult IsZeroOp::verify() {
+  Type inputType = getInput().getType();
+  if (isa<AffineType>(getElementTypeOrSelf(inputType)) ||
+      isa<JacobianType>(getElementTypeOrSelf(inputType)) ||
+      isa<XYZZType>(getElementTypeOrSelf(inputType)))
+    return success();
+  return emitError() << "invalid input type";
+}
+
 template <typename OpType>
 LogicalResult verifyBinaryOp(OpType op) {
   Type lhsType = op.getLhs().getType();

zkir/Dialect/EllipticCurve/IR/EllipticCurveOps.h

@@ -7,6 +7,7 @@
 #include "zkir/Dialect/EllipticCurve/IR/EllipticCurveTypes.h"
 #include "zkir/Dialect/Field/IR/FieldAttributes.h"
 #include "zkir/Dialect/Field/IR/FieldTypes.h"
+#include "zkir/Utils/OpUtils.h"
 
 #define GET_OP_CLASSES
 #include "zkir/Dialect/EllipticCurve/IR/EllipticCurveOps.h.inc"

zkir/Dialect/EllipticCurve/IR/EllipticCurveOps.td

@@ -51,6 +51,26 @@ def EllipticCurve_PointOp : EllipticCurve_Op<"point"> {
   let hasVerifier = 1;
 }
 
+////////////// POINT CHECKS //////////////
+
+def EllipticCurve_IsZeroOp : EllipticCurve_Op<"is_zero", [TypesMatchWith<
+  "result type has i1 element type and same shape as operands",
+  "input", "output", "getI1SameShape($_self)">]> {
+  let summary = "Checks whether an elliptic curve point is zero";
+  let description = [{
+    Outputs a bool true (1) if the elliptic curve point input is zero, and a bool false (0) otherwise.
+
+    Example:
+    ```
+    %0 = elliptic_curve.is_zero %point : !jacobian
+    ```
+  }];
+  let arguments = (ins PointLike:$input);
+  let results = (outs BoolLike:$output);
+  let hasVerifier = 1;
+  let assemblyFormat = "operands attr-dict `:` type($input)";
+}
+
 /////////// POINT EXTRACTION //////////////
 
 def EllipticCurve_ExtractOp : EllipticCurve_Op<"extract"> {