[MLIR] Create add of sub folder

Create folders for add(sub(a, b), b) -> a and add(b, sub(a, b)) -> a

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D116471

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

@@ -194,6 +194,16 @@ OpFoldResult arith::AddIOp::fold(ArrayRef<Attribute> operands) {
   if (matchPattern(getRhs(), m_Zero()))
     return getLhs();
 
+  // add(sub(a, b), b) -> a
+  if (auto sub = getLhs().getDefiningOp<SubIOp>())
+    if (getRhs() == sub.getRhs())
+      return sub.getLhs();
+
+  // add(b, sub(a, b)) -> a
+  if (auto sub = getRhs().getDefiningOp<SubIOp>())
+    if (getLhs() == sub.getRhs())
+      return sub.getLhs();
+
   return constFoldBinaryOp<IntegerAttr>(
       operands, [](APInt a, const APInt &b) { return std::move(a) + b; });
 }

mlir/test/Dialect/Arithmetic/canonicalize.mlir

@@ -299,6 +299,22 @@ func @tripleSubSub3(%arg0: index) -> index {
   return %add2 : index
 }
 
+// CHECK-LABEL: @doubleAddSub1
+//  CHECK-NEXT:   return %arg0
+func @doubleAddSub1(%arg0: index, %arg1 : index) -> index {
+  %sub = arith.subi %arg0, %arg1 : index
+  %add = arith.addi %sub, %arg1 : index
+  return %add : index
+}
+
+// CHECK-LABEL: @doubleAddSub2
+//  CHECK-NEXT:   return %arg0
+func @doubleAddSub2(%arg0: index, %arg1 : index) -> index {
+  %sub = arith.subi %arg0, %arg1 : index
+  %add = arith.addi %arg1, %sub : index
+  return %add : index
+}
+
 // CHECK-LABEL: @notCmpEQ
 //       CHECK:   %[[cres:.+]] = arith.cmpi ne, %arg0, %arg1 : i8
 //       CHECK:   return %[[cres]]