[Mlir-commits] [mlir] 21aa2a1 - [MLIR] Create add of sub folder

[William S. Moses]

Author: William S. Moses
Date: 2022-01-03T09:28:20-05:00
New Revision: 21aa2a1b09118c4678c198672ef8eb23e0cfd8e7

URL: https://github.com/llvm/llvm-project/commit/21aa2a1b09118c4678c198672ef8eb23e0cfd8e7
DIFF: https://github.com/llvm/llvm-project/commit/21aa2a1b09118c4678c198672ef8eb23e0cfd8e7.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp
    mlir/test/Dialect/Arithmetic/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp b/mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp
index 59af0a5c999b3..1536eeaf48af8 100644
--- a/mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp
+++ b/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; });
 }

diff  --git a/mlir/test/Dialect/Arithmetic/canonicalize.mlir b/mlir/test/Dialect/Arithmetic/canonicalize.mlir
index dff4141a38387..e4cbe21710ba8 100644
--- a/mlir/test/Dialect/Arithmetic/canonicalize.mlir
+++ 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]]