Address review comments.

Author: weiweichen

mlir/lib/Dialect/Index/IR/IndexOps.cpp

@@ -138,39 +138,22 @@ OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
 }
 /// Canonicalize
 /// ` x = v + c1; y = x + c2` to `x = v + (c1 + c2)`
-/// ` x = v + c1; y = c2 + x` to `x = v + (c1 + c2)`
-/// ` x = c1 + v; y = x + c2` to `x = v + (c1 + c2)`
-/// ` x = c1 + v; y = c2 + x` to `x = v + (c1 + c2)`
 LogicalResult AddOp::canonicalize(AddOp op, PatternRewriter &rewriter) {
-
-  auto matchConstant = [](mlir::index::AddOp op, Value &v, IntegerAttr &c) {
-    v = op.getLhs();
-    if (!mlir::matchPattern(op.getRhs(), mlir::m_Constant(&c))) {
-      v = op.getRhs();
-      if (!mlir::matchPattern(op.getLhs(), mlir::m_Constant(&c)))
-        return false;
-    }
-    return true;
-  };
-
   IntegerAttr c1, c2;
-  Value v1, v2;
+  if (!mlir::matchPattern(op.getRhs(), mlir::m_Constant(&c1)))
+    return rewriter.notifyMatchFailure(op.getLoc(), "RHS is not a constant");
 
-  if (!matchConstant(op, v1, c1))
-    return rewriter.notifyMatchFailure(op.getLoc(),
-                                       "neither LHS nor RHS is constant");
-
-  auto add = v1.getDefiningOp<mlir::index::AddOp>();
+  auto add = op.getLhs().getDefiningOp<mlir::index::AddOp>();
   if (!add)
     return rewriter.notifyMatchFailure(op.getLoc(), "LHS is not a add");
 
-  if (!matchConstant(add, v2, c2))
-    return rewriter.notifyMatchFailure(op.getLoc(),
-                                       "neither LHS nor RHS is constant");
+  if (!mlir::matchPattern(add.getRhs(), mlir::m_Constant(&c2)))
+    return rewriter.notifyMatchFailure(op.getLoc(), "RHS is not a constant");
 
   auto c = rewriter.create<mlir::index::ConstantOp>(op->getLoc(),
                                                     c1.getInt() + c2.getInt());
-  auto newAdd = rewriter.create<mlir::index::AddOp>(op->getLoc(), v2, c);
+  auto newAdd =
+      rewriter.create<mlir::index::AddOp>(op->getLoc(), add.getLhs(), c);
 
   rewriter.replaceOp(op, newAdd);
   return success();

mlir/test/Dialect/Index/index-canonicalize.mlir

@@ -37,14 +37,12 @@ func.func @add_fold_constants(%arg: index) -> (index) {
   %0 = index.constant 1
   %1 = index.constant 2
   %2 = index.add %arg, %0
-  %3 = index.add %1, %2
-  %4 = index.add %3, %1
-  %5 = index.add %4, %0
+  %3 = index.add %2, %1
 
-  // CHECK-DAG: [[A:%.*]] = index.constant 6
-  // CHECK-DAG: [[V0:%.*]] = index.add %arg0, [[A]]
+  // CHECK-DAG: [[C3:%.*]] = index.constant 3
+  // CHECK-DAG: [[V0:%.*]] = index.add %arg0, [[C3]]
   // CHECK: return [[V0]]
-  return %5 : index
+  return %3 : index
 }
 
 // CHECK-LABEL: @sub