[Mlir-commits] [mlir] [mlir][tensor] Fold unpadding collapse_shape into extract_slice (PR #93554)

[Quinn Dawkins]

================
@@ -48,6 +48,39 @@ struct FoldExpandOfRankReducingExtract
   }
 };
 
+/// Fold collapse_shape which only removes static dimensions of size `1`
+/// into extract_slice.
+struct FoldUnPaddingCollapseIntoExtract
+    : public OpRewritePattern<tensor::CollapseShapeOp> {
+  using OpRewritePattern<tensor::CollapseShapeOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::CollapseShapeOp collapseShapeOp,
+                                PatternRewriter &rewriter) const override {
+    auto extractSliceOp =
+        collapseShapeOp.getSrc().getDefiningOp<tensor::ExtractSliceOp>();
+    // Collapse cannot be folded away with multiple users of the extract slice
+    // and it is not necessarily beneficial to only convert the collapse into
+    // another extract slice.
+    if (!extractSliceOp || !extractSliceOp.getResult().hasOneUse())
----------------
qedawkins wrote:

nit:
```suggestion
    if (!extractSliceOp || !extractSliceOp.hasOneUse())
```

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