move all to-shape-casts to single file, where sensible

Author: newling

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

@@ -2267,27 +2267,6 @@ class ExtractOpFromCreateMask final : public OpRewritePattern<ExtractOp> {
   }
 };
 
-// Folds extract(shape_cast(..)) into shape_cast when the total element count
-// does not change.
-LogicalResult foldExtractFromShapeCastToShapeCast(ExtractOp extractOp,
-                                                  PatternRewriter &rewriter) {
-  auto castOp = extractOp.getVector().getDefiningOp<ShapeCastOp>();
-  if (!castOp)
-    return failure();
-
-  VectorType sourceType = castOp.getSourceVectorType();
-  auto targetType = dyn_cast<VectorType>(extractOp.getResult().getType());
-  if (!targetType)
-    return failure();
-
-  if (sourceType.getNumElements() != targetType.getNumElements())
-    return failure();
-
-  rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(extractOp, targetType,
-                                                   castOp.getSource());
-  return success();
-}
-
 /// Try to canonicalize the extraction of a subvector from a vector defined by
 /// vector.from_elements. E.g.:
 ///
@@ -2335,14 +2314,14 @@ LogicalResult foldExtractFromFromElements(ExtractOp extractOp,
   return success();
 }
 
-/// The canonical form of vector operations that just reshape vectors is
-/// vector.shape_cast. This pattern canonicalizes vector.extract ops of this
-/// kind.
+/// Replace `vector.extract` to `vector.shape_cast`.
 ///
 /// BEFORE:
 /// %0 = vector.extract %arg0[0] : vector<4xf32> from vector<1x4xf32>
 /// AFTER:
 /// %0 = vector.shape_cast %arg0 : vector<1x4xf32> to vector<4xf32>
+///
+/// The canonical form of vector operations that reshape vectors is shape_cast.
 struct ExtractToShapeCast final : public OpRewritePattern<vector::ExtractOp> {
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(vector::ExtractOp extractOp,
@@ -2376,7 +2355,6 @@ void ExtractOp::getCanonicalizationPatterns(RewritePatternSet &results,
   results
       .add<ExtractOpFromBroadcast, ExtractOpFromCreateMask, ExtractToShapeCast>(
           context);
-  results.add(foldExtractFromShapeCastToShapeCast);
   results.add(foldExtractFromFromElements);
 }
 
@@ -2966,14 +2944,14 @@ struct BroadcastFolder : public OpRewritePattern<BroadcastOp> {
   }
 };
 
-/// The canonical form of vector operations that just reshape vectors is
-/// vector.shape_cast. This pattern canonicalizes vector.broadcast ops of this
-/// kind.
+/// Replace `vector.broadcast` with `vector.shape_cast`.
 ///
 /// BEFORE:
 /// %0 = vector.broadcast %arg0 : vector<4xi8> to vector<1x1x4xi8>
 /// AFTER:
 /// %0 = vector.shape_cast %arg0 : vector<4xi8> to vector<1x1x4xi8>
+///
+/// The canonical form of vector operations that reshape vectors is shape_cast.
 struct BroadcastToShapeCast final
     : public OpRewritePattern<vector::BroadcastOp> {
   using OpRewritePattern::OpRewritePattern;
@@ -6615,16 +6593,16 @@ class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
   }
 };
 
-/// The canonical form of operations that just reshape a vector is
-/// vector.shape_cast. This pattern canonicalizes vector.transpose operations of
-/// this kind.
+/// Replace `vector.transpose` with `vector.shape_cast`.
 ///
 /// BEFORE:
 /// %0 = vector.transpose %arg0, [0, 2, 1] :
 ///                   vector<2x1x2xf32> to vector<2x2x1xf32>
 /// AFTER:
 /// %0 = vector.shape_cast %arg0 :
 ///                   vector<2x1x2xf32> to vector<2x2x1xf32>
+///
+/// The canonical form of vector operations that reshape vectors is shape_cast.
 struct TransposeToShapeCast final
     : public OpRewritePattern<vector::TransposeOp> {
   using OpRewritePattern::OpRewritePattern;

mlir/test/Dialect/Vector/canonicalize.mlir

@@ -821,7 +821,8 @@ func.func @fold_extract_broadcast_0dvec_input_scalar_output(%a : vector<f32>,
 
 // -----
 
-
+// This test is negative in the sense that the broadcast is not folded into the extract.
+// The extract is still converted into shape_cast, however.
 // CHECK-LABEL: negative_fold_extract_broadcast
 //       CHECK:   vector.broadcast %{{.*}} : vector<1x1xf32> to vector<1x1x4xf32>
 //       CHECK:   vector.shape_cast{{.*}} vector<1x1x4xf32> to vector<4xf32>
@@ -939,6 +940,10 @@ func.func @fold_extract_broadcast_to_equal_rank(%a : vector<1xf32>, %idx0 : inde
 
 // -----
 
+
+// One possible path this takes is
+// 1) Match on [ExtractOpFromBroadcast], which matches as the extract is broadcastlike.
+// 2) Match on [BroadcastToShapeCast], as the resulting broadcast just prepends a 1.
 // CHECK-LABEL: fold_extract_broadcastlike_shape_cast
 //  CHECK-SAME:   %[[A:.*]]: vector<1xf32>
 //       CHECK:   %[[R:.*]] = vector.shape_cast %[[A]] : vector<1xf32> to vector<1x1xf32>
@@ -1028,18 +1033,6 @@ func.func @negative_fold_extract_shapecast(%arg0 : vector<16xf32>) -> vector<4x2
 
 // -----
 
-// CHECK-LABEL: fold_extract_shapecast_to_shapecast
-//  CHECK-SAME: (%[[ARG:.+]]: vector<3x4xf32>)
-//       CHECK:   %[[R:.+]] = vector.shape_cast %[[ARG]] : vector<3x4xf32> to vector<12xf32>
-//       CHECK:   return %[[R]]
-func.func @fold_extract_shapecast_to_shapecast(%arg0 : vector<3x4xf32>) -> vector<12xf32> {
-  %0 = vector.shape_cast %arg0 : vector<3x4xf32> to vector<1x12xf32>
-  %r = vector.extract %0[0] : vector<12xf32> from vector<1x12xf32>
-  return %r : vector<12xf32>
-}
-
-// -----
-
 // CHECK-LABEL: func @extract_no_fold_scalar_to_0d(
 //  CHECK-SAME:     %[[v:.*]]: vector<f32>)
 //       CHECK:   %[[extract:.*]] = vector.extract %[[v]][] : f32 from vector<f32>
@@ -1154,30 +1147,6 @@ func.func @canonicalize_broadcast_shapecast_to_broadcast_scalar(%arg0: f32) -> v
 
 // -----
 
-// In this test, broadcast (2)->(1,2,1) is not legal, but shape_cast (2)->(1,2,1) is.
-// CHECK-LABEL: func @canonicalize_broadcast_shapecast_to_shapcast
-//   CHECK-NOT:   vector.broadcast
-//       CHECK:   vector.shape_cast {{.+}} : vector<2xf32> to vector<1x2x1xf32>
-func.func @canonicalize_broadcast_shapecast_to_shapcast(%arg0 : vector<2xf32>) -> vector<1x2x1xf32> {
-  %0 = vector.broadcast %arg0 : vector<2xf32> to vector<1x2xf32>
-  %1 = vector.shape_cast %0 : vector<1x2xf32> to vector<1x2x1xf32>
-  return %1 : vector<1x2x1xf32>
-}
-
-// -----
-
-// In this test, broadcast (1)->(1,1) and shape_cast (1)->(1,1) are both legal. shape_cast is chosen.
-// CHECK-LABEL: func @canonicalize_broadcast_shapecast_both_possible
-//   CHECK-NOT:   vector.broadcast
-//       CHECK:   vector.shape_cast {{.+}} : vector<1xf32> to vector<1x1xf32>
-func.func @canonicalize_broadcast_shapecast_both_possible(%arg0: vector<1xf32>) -> vector<1x1xf32> {
-    %0 = vector.broadcast %arg0 : vector<1xf32> to vector<1x1x1xf32>
-    %1 = vector.shape_cast %0 : vector<1x1x1xf32> to vector<1x1xf32>
-    return %1 : vector<1x1xf32>
-}
-
-// -----
-
 // CHECK-LABEL: func @canonicalize_shapecast_broadcast_to_broadcast_prepend_dim
 //   CHECK-NOT:   vector.shape_cast
 //       CHECK:   vector.broadcast {{.+}} : vector<2xf32> to vector<32x2xf32>
@@ -1571,7 +1540,7 @@ func.func @extract_strided_broadcast4(%arg0: f32) -> vector<1x4xf32> {
 
 // -----
 
-// Check the case where the same dimension is both broadcasted and sliced 
+// Check the case where the same dimension is both broadcasted and sliced
 // CHECK-LABEL: func @extract_strided_broadcast5
 //  CHECK-SAME: (%[[ARG:.+]]: vector<2x1xf32>)
 //       CHECK: %[[V:.+]] = vector.broadcast %[[ARG]] : vector<2x1xf32> to vector<2x4xf32>
@@ -2186,20 +2155,6 @@ func.func @extract_strided_splatlike(%arg0: f16) -> vector<2x4xf16> {
 
 // -----
 
-// CHECK-LABEL: func @insert_extract_to_shape_cast
-//  CHECK-SAME: (%[[ARG0:.*]]: vector<1x1x4xf32>, %[[ARG1:.*]]: vector<4xf32>)
-//       CHECK:   %[[V0:.*]] = vector.shape_cast %[[ARG0]] : vector<1x1x4xf32> to vector<4xf32>
-//       CHECK:   %[[V1:.*]] = vector.shape_cast %[[ARG1]] : vector<4xf32> to vector<1x1x4xf32>
-//       CHECK:   return %[[V0]], %[[V1]] : vector<4xf32>, vector<1x1x4xf32>
-func.func @insert_extract_to_shape_cast(%arg0 : vector<1x1x4xf32>,
-  %arg1 : vector<4xf32>) -> (vector<4xf32>, vector<1x1x4xf32>) {
-  %0 = vector.extract %arg0[0, 0] : vector<4xf32> from vector<1x1x4xf32>
-  %1 = vector.insert %arg1, %arg0 [0, 0] : vector<4xf32> into vector<1x1x4xf32>
-  return %0, %1 : vector<4xf32>, vector<1x1x4xf32>
-}
-
-// -----
-
 // CHECK-LABEL: func.func @extract_splat_constant
 //   CHECK-DAG:   %[[CST1:.*]] = arith.constant 1 : i32
 //   CHECK-DAG:   %[[CST0:.*]] = arith.constant dense<2.000000e+00> : vector<7xf32>
@@ -2554,6 +2509,7 @@ func.func @shuffle_1d_rhs_poison() -> vector<4xi32> {
 
 // -----
 
+// The shuffle becomes a broadcast, which is then canonicalized to a shapecast.
 // CHECK-LABEL: func @shuffle_canonicalize_0d
 func.func @shuffle_canonicalize_0d(%v0 : vector<i32>, %v1 : vector<i32>) -> vector<1xi32> {
   // CHECK: vector.shape_cast %{{.*}} : vector<i32> to vector<1xi32>
@@ -2928,15 +2884,6 @@ func.func @transfer_read_from_rank_reducing_extract_slice(%src: tensor<1x8x8x8xf
 
 // -----
 
-// CHECK-LABEL: func.func @extract_from_broadcast
-func.func @extract_from_broadcast(%src: vector<1x1x1xf32>) -> vector<1xf32> {
-  %0 = vector.broadcast %src : vector<1x1x1xf32> to vector<1x1x32x1xf32>
-  //  CHECK-NEXT:   %[[RES:.*]] = vector.shape_cast{{.*}} vector<1x1x1xf32> to vector<1xf32>
-  //  CHECK-NEXT:   return %[[RES]] : vector<1xf32>
-  %1 = vector.extract %0[0, 0, 31] : vector<1xf32> from vector<1x1x32x1xf32>
-  return %1: vector<1xf32>
-}
-
 // CHECK-LABEL: func.func @extract_from_stretch_broadcast
 func.func @extract_from_stretch_broadcast(%src: vector<3x1x2xf32>) -> f32 {
   //  CHECK-NEXT:  %0 = vector.extract {{.*}}[0, 0, 0] : f32 from vector<3x1x2xf32>
@@ -2947,6 +2894,7 @@ func.func @extract_from_stretch_broadcast(%src: vector<3x1x2xf32>) -> f32 {
 }
 
 // -----
+
 // CHECK-LABEL: func.func @extract_strided_slice_of_constant_mask
 func.func @extract_strided_slice_of_constant_mask() -> vector<5x7xi1>{
   //  CHECK-NEXT:   %[[RES:.*]] = vector.constant_mask [5, 4] : vector<5x7xi1>