[mlir][arith] Canonicalization patterns for subi.

[vzakhari]

Summary

subi(addi(a, b), b) -> a
subi(addi(a, b), a) -> b
subi(subi(a, b), a) -> subi(0, b)

[jeanPerier]

Thanks, looks good to me.

[vzakhari]

In D133615#3783399, @jeanPerier wrote:

Thanks, looks good to me.

Thank you for the review, Jean! I will merge this in a couple of days to let other reviewers give their feedback.

mlir/lib/Dialect/Arithmetic/IR/ArithmeticCanonicalization.td

 //===- ArithmeticPatterns.td - Arithmetic dialect patterns -*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef ARITHMETIC_PATTERNS
 #define ARITHMETIC_PATTERNS
 
 include "mlir/IR/PatternBase.td"
 include "mlir/Dialect/Arithmetic/IR/ArithmeticOps.td"
 
+// Create zero attribute of type matching the argument's type.
+def GetZeroAttr : NativeCodeCall<"$_builder.getZeroAttr($0.getType())">;
+
 // Add two integer attributes and create a new one with the result.
 def AddIntAttrs : NativeCodeCall<"addIntegerAttrs($_builder, $0, $1, $2)">;
 
 // Subtract two integer attributes and createa a new one with the result.
 def SubIntAttrs : NativeCodeCall<"subIntegerAttrs($_builder, $0, $1, $2)">;
 
 //===----------------------------------------------------------------------===//
 // AddIOp

 
 // subi(c1, subi(c0, x)) -> addi(x, c1 - c0)
 def SubILHSSubConstantLHS :
     Pat<(Arith_SubIOp:$res
           (Arith_ConstantOp APIntAttr:$c1),
           (Arith_SubIOp (Arith_ConstantOp APIntAttr:$c0), $x)),
         (Arith_AddIOp $x, (Arith_ConstantOp (SubIntAttrs $res, $c1, $c0)))>;
 
+// subi(subi(a, b), a) -> subi(0, b)
+def SubISubILHSRHSLHS :
+    Pat<(Arith_SubIOp:$res (Arith_SubIOp $x, $y), $x),
+        (Arith_SubIOp (Arith_ConstantOp (GetZeroAttr $y)), $y)>;
+
 //===----------------------------------------------------------------------===//
 // XOrIOp
 //===----------------------------------------------------------------------===//
 
 // xori is commutative and will be canonicalized to have its constants appear
 // as the second operand.
 
 // not(cmpi(pred, a, b)) -> cmpi(~pred, a, b), where not(x) is xori(x, 1)

mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp

 OpFoldResult arith::SubIOp::fold(ArrayRef<Attribute> operands) {
   // subi(x,x) -> 0
   if (getOperand(0) == getOperand(1))
     return Builder(getContext()).getZeroAttr(getType());
   // subi(x,0) -> x
   if (matchPattern(getRhs(), m_Zero()))
     return getLhs();
 
+  if (auto add = getLhs().getDefiningOp<AddIOp>()) {
+    // subi(addi(a, b), b) -> a
+    if (getRhs() == add.getRhs())
+      return add.getLhs();
+    // subi(addi(a, b), a) -> b
+    if (getRhs() == add.getLhs())
+      return add.getRhs();
+  }
+
   return constFoldBinaryOp<IntegerAttr>(
       operands, [](APInt a, const APInt &b) { return std::move(a) - b; });
 }
 
 void arith::SubIOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                                 MLIRContext *context) {
-  patterns
-      .add<SubIRHSAddConstant, SubILHSAddConstant, SubIRHSSubConstantRHS,
-           SubIRHSSubConstantLHS, SubILHSSubConstantRHS, SubILHSSubConstantLHS>(
-          context);
+  patterns.add<SubIRHSAddConstant, SubILHSAddConstant, SubIRHSSubConstantRHS,
+               SubIRHSSubConstantLHS, SubILHSSubConstantRHS,
+               SubILHSSubConstantLHS, SubISubILHSRHSLHS>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // MulIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::MulIOp::fold(ArrayRef<Attribute> operands) {
   // muli(x, 0) -> 0

mlir/test/Dialect/Arithmetic/canonicalize.mlir

 func.func @tripleSubAdd1(%arg0: index) -> index {
   %c17 = arith.constant 17 : index
   %c42 = arith.constant 42 : index
   %add1 = arith.addi %c17, %arg0 : index
   %add2 = arith.subi %add1, %c42 : index
   return %add2 : index
 }
 
+// CHECK-LABEL: @subSub0
+//       CHECK:   %[[c0:.+]] = arith.constant 0 : index
+//       CHECK:   %[[add:.+]] = arith.subi %[[c0]], %arg1 : index
+//       CHECK:   return %[[add]]
+func.func @subSub0(%arg0: index, %arg1: index) -> index {
+  %sub1 = arith.subi %arg0, %arg1 : index
+  %sub2 = arith.subi %sub1, %arg0 : index
+  return %sub2 : index
+}
+
 // CHECK-LABEL: @tripleSubSub0
 //       CHECK:   %[[cres:.+]] = arith.constant 25 : index
 //       CHECK:   %[[add:.+]] = arith.addi %arg0, %[[cres]] : index
 //       CHECK:   return %[[add]]
 func.func @tripleSubSub0(%arg0: index) -> index {
   %c17 = arith.constant 17 : index
   %c42 = arith.constant 42 : index
   %add1 = arith.subi %c17, %arg0 : index

 func.func @tripleSubSub3(%arg0: index) -> index {
   %c17 = arith.constant 17 : index
   %c42 = arith.constant 42 : index
   %add1 = arith.subi %arg0, %c17 : index
   %add2 = arith.subi %add1, %c42 : index
   return %add2 : index
 }
 
+// CHECK-LABEL: @subAdd1
+//  CHECK-NEXT:   return %arg0
+func.func @subAdd1(%arg0: index, %arg1 : index) -> index {
+  %add = arith.addi %arg0, %arg1 : index
+  %sub = arith.subi %add, %arg1 : index
+  return %sub : index
+}
+
+// CHECK-LABEL: @subAdd2
+//  CHECK-NEXT:   return %arg1
+func.func @subAdd2(%arg0: index, %arg1 : index) -> index {
+  %add = arith.addi %arg0, %arg1 : index
+  %sub = arith.subi %add, %arg0 : index
+  return %sub : index
+}
+
 // CHECK-LABEL: @doubleAddSub1
 //  CHECK-NEXT:   return %arg0
 func.func @doubleAddSub1(%arg0: index, %arg1 : index) -> index {
   %sub = arith.subi %arg0, %arg1 : index
   %add = arith.addi %sub, %arg1 : index
   return %add : index
 }