[Mlir-commits] [mlir] [MLIR][Canonicalization] Added shape_cast folding patterns (PR #183061)

[Kunwar Grover]

================
@@ -2591,9 +2591,31 @@ struct ToElementsOfBroadcast final : OpRewritePattern<ToElementsOp> {
   }
 };
 
+/// Pattern to rewrite Y = ToElements(ShapeCast(X)) as Y = ToElements(X)
+///
+/// BEFORE:
+///    %1 = vector.shape_cast %0 : vector<6xf32> to vector<2x3xf32>
+///    %2:6 = vector.to_elements %1 : vector<2x3xf32>
+/// AFTER:
+///    %2:6 = vector.to_elements %0 : vector<6xf32>
+struct FoldToElementsOfShapeCast final : public OpRewritePattern<ToElementsOp> {
+  using Base::Base;
+
+  LogicalResult matchAndRewrite(ToElementsOp toElementsOp,
+                                PatternRewriter &rewriter) const override {
+    auto shapeCast = toElementsOp.getSource().getDefiningOp<ShapeCastOp>();
+    if (!shapeCast)
+      return failure();
+
+    rewriter.replaceOpWithNewOp<ToElementsOp>(toElementsOp,
+                                              shapeCast.getSource());
+    return success();
+  }
+};
+
----------------
Groverkss wrote:

This should be a fold. Have a look at ToElementsOp::fold . You can read more about the difference here: https://mlir.llvm.org/docs/Canonicalization/

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