feat(mod_arith): impl NegateOp

Author: batzor

tests/Dialect/ModArith/mod_arith_to_arith.mlir

@@ -13,6 +13,28 @@ func.func @test_lower_constant() -> !mod_arith.int<3 : i5> {
   return %res: !mod_arith.int<3 : i5>
 }
 
+// CHECK-LABEL: @test_lower_negate
+// CHECK-SAME: (%[[LHS:.*]]: [[T:.*]]) -> [[T]] {
+func.func @test_lower_negate(%lhs : !Zp) -> !Zp {
+  // CHECK-NOT: mod_arith.negate
+  // CHECK: %[[CMOD:.*]] = arith.constant 65537 : [[T]]
+  // CHECK: %[[SUB:.*]] = arith.subi %[[CMOD]], %[[LHS]] : [[T]]
+  // CHECK: return %[[SUB]] : [[T]]
+  %res = mod_arith.negate %lhs: !Zp
+  return %res : !Zp
+}
+
+// CHECK-LABEL: @test_lower_negate_vec
+// CHECK-SAME: (%[[LHS:.*]]: [[T:.*]]) -> [[T]] {
+func.func @test_lower_negate_vec(%lhs : !Zpv) -> !Zpv {
+  // CHECK-NOT: mod_arith.negate
+  // CHECK: %[[CMOD:.*]] = arith.constant dense<65537> : [[T]]
+  // CHECK: %[[SUB:.*]] = arith.subi %[[CMOD]], %[[LHS]] : [[T]]
+  // CHECK: return %[[SUB]] : [[T]]
+  %res = mod_arith.negate %lhs: !Zpv
+  return %res : !Zpv
+}
+
 // CHECK-LABEL: @test_lower_encapsulate
 // CHECK-SAME: (%[[LHS:.*]]: [[T:.*]]) -> [[T]] {
 func.func @test_lower_encapsulate(%lhs : i32) -> !Zp {

zkir/Dialect/ModArith/Conversions/ModArithToArith/ModArithToArith.cpp

@@ -128,6 +128,24 @@ struct ConvertConstant : public OpConversionPattern<ConstantOp> {
   }
 };
 
+struct ConvertNegate : public OpConversionPattern<NegateOp> {
+  explicit ConvertNegate(mlir::MLIRContext *context)
+      : OpConversionPattern<NegateOp>(context) {}
+
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      NegateOp op, OpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    ImplicitLocOpBuilder b(op.getLoc(), rewriter);
+
+    auto cmod = b.create<arith::ConstantOp>(modulusAttr(op));
+    auto sub = b.create<arith::SubIOp>(cmod, adaptor.getOperands()[0]);
+    rewriter.replaceOp(op, sub);
+    return success();
+  }
+};
+
 struct ConvertReduce : public OpConversionPattern<ReduceOp> {
   explicit ConvertReduce(mlir::MLIRContext *context)
       : OpConversionPattern<ReduceOp>(context) {}
@@ -530,15 +548,15 @@ void ModArithToArith::runOnOperation() {
 
   RewritePatternSet patterns(context);
   rewrites::populateWithGenerated(patterns);
-  patterns
-      .add<ConvertEncapsulate, ConvertExtract, ConvertReduce, ConvertMontReduce,
-           ConvertToMont, ConvertFromMont, ConvertAdd, ConvertSub, ConvertMul,
-           ConvertMontMul, ConvertMac, ConvertConstant, ConvertInverse,
-           ConvertAny<affine::AffineForOp>, ConvertAny<affine::AffineYieldOp>,
-           ConvertAny<linalg::GenericOp>, ConvertAny<linalg::YieldOp>,
-           ConvertAny<tensor::CastOp>, ConvertAny<tensor::ExtractOp>,
-           ConvertAny<tensor::FromElementsOp>, ConvertAny<tensor::InsertOp>>(
-          typeConverter, context);
+  patterns.add<
+      ConvertNegate, ConvertEncapsulate, ConvertExtract, ConvertReduce,
+      ConvertMontReduce, ConvertToMont, ConvertFromMont, ConvertAdd, ConvertSub,
+      ConvertMul, ConvertMontMul, ConvertMac, ConvertConstant, ConvertInverse,
+      ConvertAny<affine::AffineForOp>, ConvertAny<affine::AffineYieldOp>,
+      ConvertAny<linalg::GenericOp>, ConvertAny<linalg::YieldOp>,
+      ConvertAny<tensor::CastOp>, ConvertAny<tensor::ExtractOp>,
+      ConvertAny<tensor::FromElementsOp>, ConvertAny<tensor::InsertOp>>(
+      typeConverter, context);
 
   addStructuralConversionPatterns(typeConverter, patterns, target);
 

zkir/Dialect/ModArith/IR/ModArithDialect.cpp

@@ -95,6 +95,10 @@ LogicalResult verifySameWidth(OpType op, ModArithType modArithType,
   return success();
 }
 
+LogicalResult NegateOp::verify() {
+  return verifyModArithType(*this, getResultModArithType(*this));
+}
+
 LogicalResult ExtractOp::verify() {
   auto modArithType = getOperandModArithType(*this);
   auto integerType = getResultIntegerType(*this);

zkir/Dialect/ModArith/IR/ModArithOps.td

@@ -97,6 +97,24 @@ def ModArith_ConstantOp : Op<ModArith_Dialect, "constant",
   let hasCustomAssemblyFormat = 1;
 }
 
+def ModArith_NegateOp : ModArith_Op<"negate", [Pure, ElementwiseMappable, SameOperandsAndResultType, Involution]> {
+  let summary = "negate a mod arith element";
+
+  let description = [{
+    `mod_arith.negate x` computes $-x \mod q$
+    Examples:
+    ```
+    %1 = mod_arith.negate %0 : mod_arith.int<65537: i32>
+    ```
+  }];
+
+  let arguments = (ins
+    ModArithLike:$input
+  );
+  let results = (outs ModArithLike:$output);
+  let hasVerifier = 1;
+  let assemblyFormat = "operands attr-dict `:` type($output)";
+}
 
 def ModArith_ReduceOp : ModArith_Op<"reduce", [Pure, ElementwiseMappable, SameOperandsAndResultType]> {
   let summary = "reduce the mod arith type to its canonical representative";