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

[llvmlistbot at llvm.org]

================
@@ -4625,11 +4625,39 @@ struct DropDelinearizeOneBasisElement
   }
 };
 
+/// If a `affine.delinearize_index`'s input is a `affine.linearize_index
+/// disjoint` and the two operations have the same basis, replace the
+/// delinearizeation results with the inputs of the `affine.linearize_index`
+/// since they are exact inverses of each other.
+///
+/// The `disjoint` flag is needed on the `affine.linearize_index` because
+/// otherwise, there is no guarantee that the inputs to the linearization are
+/// in-bounds the way the outputs of the delinearization would be.
+struct CancelDelinearizeOfLinearizeDisjointExact
+    : public OpRewritePattern<affine::AffineDelinearizeIndexOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(affine::AffineDelinearizeIndexOp delinearizeOp,
+                                PatternRewriter &rewriter) const override {
+    auto linearizeOp = delinearizeOp.getLinearIndex()
+                           .getDefiningOp<affine::AffineLinearizeIndexOp>();
+    if (!linearizeOp)
+      return failure();
+
+    if (!linearizeOp.getDisjoint() ||
+        linearizeOp.getMixedBasis() != delinearizeOp.getMixedBasis())
+      return failure();
----------------
MaheshRavishankar wrote:

For such failures I prefer to return a `rewriter.notifyMatchFailure(delinearizeOp, "<message>"`. Really helps with `--debug` prints.

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