[Mlir-commits] [mlir] [mlir] Add bubbling patterns for non intersecting reshapes (PR #94637)

[llvmlistbot at llvm.org]

================
@@ -1086,6 +1086,76 @@ struct FoldReshapeWithGenericOpByExpansion
 private:
   ControlFusionFn controlFoldingReshapes;
 };
+
+/// Pattern to bubble up a tensor.expand_shape op through a producer
+/// tensor.collapse_shape op that has non intersecting reassociations.
+struct BubbleUpExpandThroughParallelCollapse
+    : public OpRewritePattern<tensor::ExpandShapeOp> {
+  using OpRewritePattern<tensor::ExpandShapeOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::ExpandShapeOp expandOp,
+                                PatternRewriter &rewriter) const override {
+    auto collapseOp =
+        expandOp.getSrc().getDefiningOp<tensor::CollapseShapeOp>();
+    if (!collapseOp || !collapseOp->hasOneUse())
+      return failure();
+    auto expandReInds = expandOp.getReassociationIndices();
+    auto collapseReInds = collapseOp.getReassociationIndices();
+
+    // Reshapes are parallel to each other if none of the reassociation indices
+    // have greater than 1 index for both reshapes.
+    for (auto [expandReassociation, collapseReassociation] :
----------------
Max191 wrote:

This comment means that there are no reassociations where the size is greater than 1 for both the expand and collapse at the same time. There could be cases where only one of the collapse or expand shape have size > 1, which would be parallel reshapes, but not identity reshapes. I can update the comment to be more clear.

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