[MLIR][Canonicalization] Add shape_cast folding patterns (llvm#183061)

### Summary

This PR adds two shape_cast-related canonicalization patterns for
`vector.to_elements` and `vector.from_elements`.

### Details

- Added` ToElements(ShapeCast(X)) -> ToElements(X)` as an in-place fold
in `ToElementsOp::fold`.
- Added `ShapeCast(FromElements(X)) -> FromElements(X)` as an
`OpRewritePattern` — it must be a pattern (not a `fold`) because we have
to create new op `FromElementsOp` with updated result type. This cannot
be done with a `fold`, because `fold` cannot create new ops and the
existing `FromElementsOp` result type differs from the `ShapeCastOp`
result type. Mutating the `FromElementsOp` (not root op) would violate
the `fold` contract and break other users.
- Added lit tests for the both ops (new `vector-to-elements.mlir`,
`vector-from-elements.mlir`)

---------

Signed-off-by: Alexandra Sidorova <asidorov@amd.com>

Author: a-sidorova

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

@@ -2559,6 +2559,13 @@ LogicalResult ToElementsOp::fold(FoldAdaptor adaptor,
                                  SmallVectorImpl<OpFoldResult> &results) {
   if (succeeded(foldToElementsFromElements(*this, results)))
     return success();
+
+  // Y = ToElements(ShapeCast(X)) -> Y = ToElements(X)
+  if (auto shapeCast = getSource().getDefiningOp<ShapeCastOp>()) {
+    setOperand(shapeCast.getSource());
+    return success();
+  }
+
   return foldToElementsOfBroadcast(*this, results);
 }
 
@@ -6723,13 +6730,43 @@ class ShapeCastBroadcastFolder final : public OpRewritePattern<ShapeCastOp> {
   }
 };
 
+/// Pattern to rewrite Y = ShapeCast(FromElements(X)) as Y = FromElements(X)
+///
+/// BEFORE:
+///    %1 = vector.from_elements %c1, %c2, %c3 : vector<3xf32>
+///    %2 = vector.shape_cast %1 : vector<3xf32> to vector<1x3xf32>
+/// AFTER:
+///    %2 = vector.from_elements %c1, %c2, %c3 : vector<1x3xf32>
+///
+/// Note: this transformation is implemented as an OpRewritePattern, not as a
+/// fold, because we have to create new op FromElementsOp with updated result
+/// type. This cannot be done with a fold, because fold cannot create new ops
+/// and the existing FromElementsOp result type differs from the ShapeCastOp
+/// result type. Mutating the FromElementsOp (not root op) would violate the
+/// fold contract and break other users.
+class FoldShapeCastOfFromElements final : public OpRewritePattern<ShapeCastOp> {
+public:
+  using Base::Base;
+
+  LogicalResult matchAndRewrite(ShapeCastOp shapeCastOp,
+                                PatternRewriter &rewriter) const override {
+    auto fromElements = shapeCastOp.getSource().getDefiningOp<FromElementsOp>();
+    if (!fromElements)
+      return failure();
+
+    rewriter.replaceOpWithNewOp<FromElementsOp>(
+        shapeCastOp, shapeCastOp.getResultVectorType(),
+        fromElements.getElements());
+    return success();
+  }
+};
+
 } // namespace
 
 void ShapeCastOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
-  results
-      .add<ShapeCastCreateMaskFolderTrailingOneDim, ShapeCastBroadcastFolder>(
-          context);
+  results.add<ShapeCastCreateMaskFolderTrailingOneDim, ShapeCastBroadcastFolder,
+              FoldShapeCastOfFromElements>(context);
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/Vector/canonicalize.mlir

@@ -1170,6 +1170,19 @@ func.func @dont_fold_expand_collapse(%arg0: vector<1x1x64xf32>) -> vector<8x8xf3
 
 // -----
 
+// CHECK-LABEL:   func.func @fold_shape_cast_from_elements(
+// CHECK-SAME:    %[[C1:.*]]: f32, %[[C2:.*]]: f32, %[[C3:.*]]: f32, %[[C4:.*]]: f32
+func.func @fold_shape_cast_from_elements(%c1: f32, %c2: f32, %c3: f32, %c4: f32) -> vector<2x2xf32>{
+// CHECK:           %[[VAL_0:.*]] = vector.from_elements %[[C1]], %[[C2]], %[[C3]], %[[C4]] : vector<2x2xf32>
+// CHECK:           return %[[VAL_0]] : vector<2x2xf32>
+// CHECK-NOT: vector.shape_cast
+  %1 = vector.from_elements %c1, %c2, %c3, %c4 : vector<4xf32>
+  %2 = vector.shape_cast %1 : vector<4xf32> to vector<2x2xf32>
+  return %2 : vector<2x2xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @fold_broadcast_shapecast
 //  CHECK-SAME: (%[[V:.+]]: vector<4xf32>)
 //       CHECK:   return %[[V]]

mlir/test/Dialect/Vector/canonicalize/vector-to-elements.mlir

@@ -0,0 +1,18 @@
+// RUN: mlir-opt %s -canonicalize="test-convergence" -split-input-file -allow-unregistered-dialect | FileCheck %s
+
+// This file contains some tests of folding/canonicalizing vector.to_elements
+
+///===----------------------------------------------===//
+///  Tests of `ToElementsOp::fold`
+///===----------------------------------------------===//
+
+// CHECK-LABEL: func @to_elements_of_shape_cast_folds
+// CHECK-SAME: (%[[VEC:.*]]: vector<4xf32>) -> (f32, f32, f32, f32)
+func.func @to_elements_of_shape_cast_folds(%v: vector<4xf32>) -> (f32, f32, f32, f32) {
+  %sc = vector.shape_cast %v : vector<4xf32> to vector<2x2xf32>
+  %e:4 = vector.to_elements %sc : vector<2x2xf32>
+  // CHECK-NOT: vector.shape_cast
+  // CHECK: %[[E:.*]]:4 = vector.to_elements %[[VEC]] : vector<4xf32>
+  // CHECK: return %[[E]]#0, %[[E]]#1, %[[E]]#2, %[[E]]#3 : f32, f32, f32, f32
+  return %e#0, %e#1, %e#2, %e#3 : f32, f32, f32, f32
+}