[Mlir-commits] [mlir] [MLIR][Tensor] Remove FoldDimOf[Expand|Collapse]Shape Pattern (PR #134219)

[Vivek Khandelwal]

https://github.com/vivekkhandelwal1 created https://github.com/llvm/llvm-project/pull/134219

This commit removes the FoldDimOfExpandShape and FoldDimOfCollapseShape pattern from the list of ExpandShapeOp's canonicalization patterns.

The above pattern were resulting in crash while folding the dims of an expanded tensor. The issue can be reproduced by undoing the changes done in this commit and by running the command:

```
mlir-opt --linalg-fuse-elementwise-ops repro.mlir
```

over the IR: https://gist.github.com/vivekkhandelwal1/56a1a398c21d739df77a67ce372b9366.

>From dad4284571c7bc233f4deeae5825d03ddb9a8d75 Mon Sep 17 00:00:00 2001
From: Vivek Khandelwal <vivekkhandelwal1424 at gmail.com>
Date: Thu, 3 Apr 2025 13:59:47 +0530
Subject: [PATCH] [MLIR][Tensor] Remove FoldDimOf[Expand|Collapse]Shape Pattern

This commit removes the FoldDimOfExpandShape and FoldDimOfCollapseShape
pattern from the list of ExpandShapeOp's canonicalization patterns.

The above pattern were resulting in crash while folding the dims of
an expanded tensor. The issue can be reproduced by undoing the changes
done in this commit and by running the command:

```
mlir-opt --linalg-fuse-elementwise-ops repro.mlir
```

over the IR: https://gist.github.com/vivekkhandelwal1/56a1a398c21d739df77a67ce372b9366.

Signed-off-by: Vivek Khandelwal <vivekkhandelwal1424 at gmail.com>
---
 mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index d589f627d896e..35265ac49fcbf 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2158,8 +2158,7 @@ void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
       ComposeExpandOfCollapseOp<ExpandShapeOp, CollapseShapeOp>,
       ConvertToStaticExpandShape, FoldReshapeWithConstant<ExpandShapeOp>,
       FoldReshapeWithSplat<ExpandShapeOp>,
-      FoldReshapeWithFromElements<ExpandShapeOp>, FoldDimOfExpandShape,
-      FoldDimOfCollapseShape>(context);
+      FoldReshapeWithFromElements<ExpandShapeOp>>(context);
 }
 
 void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,