[Mlir-commits] [mlir] [MLIR] Add canonicalizations to all eligible `index` binary ops (PR #114000)

[Mehdi Amini]

================
@@ -118,6 +118,31 @@ static OpFoldResult foldBinaryOpChecked(
   return IntegerAttr::get(IndexType::get(lhs.getContext()), *result64);
 }
 
+/// Helper for associative and commutative binary ops that can be transformed:
+/// `x = op(v, c1); y = op(x, c2)` -> `tmp = op(c1, c2); y = op(v, tmp)`
+/// where c1 and c2 are constants. It is expected that `tmp` will be folded.
+template <typename BinaryOp>
+static LogicalResult
+canonicalizeAssociativeCommutativeBinaryOp(BinaryOp op,
+                                           PatternRewriter &rewriter) {
+  IntegerAttr c1, c2;
+  if (!mlir::matchPattern(op.getRhs(), mlir::m_Constant(&c1)))
+    return rewriter.notifyMatchFailure(op.getLoc(), "RHS is not a constant");
+
+  auto lhsOp = op.getLhs().template getDefiningOp<BinaryOp>();
+  if (!lhsOp)
+    return rewriter.notifyMatchFailure(op.getLoc(), "LHS is not a add");
+
+  if (!mlir::matchPattern(lhsOp.getRhs(), mlir::m_Constant(&c2)))
+    return rewriter.notifyMatchFailure(op.getLoc(), "RHS is not a constant");
+
+  auto c = rewriter.createOrFold<BinaryOp>(op->getLoc(), op.getRhs(),
+                                           lhsOp.getRhs());
+
----------------
joker-eph wrote:

We should check here that `c` is not a BinaryOp (that is: it was folded) before replacing.

https://github.com/llvm/llvm-project/pull/114000