[Mlir-commits] [mlir] [mlir][affine] Cancel exactly-matching delinearize/linearize pairs (PR #115758)

[llvmlistbot at llvm.org]

================
@@ -4751,12 +4779,45 @@ struct DropLinearizeOneBasisElement final
     return success();
   }
 };
+
+/// Cancel out linearize_index(delinearize_index(x, B), B).
+///
+/// That is, rewrite
+/// ```
+/// %0:N = affine.delinearize_index %x by (%b1, %b2, ... %bN)
+/// %y = affine.linearize_index [%0#0, %0#1, ... %0#(N-1)] by (%b1, %b2, ...
+/// %bN)
+/// ```
+/// to replacing `%y` with `%x`.
+struct CancelLinearizeOfDelinearizeExact final
+    : OpRewritePattern<affine::AffineLinearizeIndexOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(affine::AffineLinearizeIndexOp linearizeOp,
+                                PatternRewriter &rewriter) const override {
+    auto delinearizeOp = linearizeOp.getMultiIndex()
+                             .front()
+                             .getDefiningOp<affine::AffineDelinearizeIndexOp>();
+    if (!delinearizeOp)
+      return failure();
+
+    if (linearizeOp.getMixedBasis() != delinearizeOp.getMixedBasis())
+      return failure();
+
+    if (delinearizeOp.getResults() != linearizeOp.getMultiIndex())
+      return failure();
+
+    rewriter.replaceOp(linearizeOp, delinearizeOp.getLinearIndex());
+    return success();
+  }
+};
 } // namespace
 
 void affine::AffineLinearizeIndexOp::getCanonicalizationPatterns(
     RewritePatternSet &patterns, MLIRContext *context) {
   patterns.add<DropLinearizeUnitComponentsIfDisjointOrZero,
-               DropLinearizeOneBasisElement>(context);
+               DropLinearizeOneBasisElement, CancelLinearizeOfDelinearizeExact>(
----------------
MaheshRavishankar wrote:

Nit: Please maintain these alphabetically

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