[mlir][Vector] Add a ShapeCastOp(BroadcastOp) canonicalization pattern

[nicolasvasilache]

Summary

This pattern can kick in when the source of the broadcast has a shape
that is a prefix/suffix of the result of the shape_cast.

[nicolasvasilache]

Drop prefix case that is potentially dangerous.

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

LogicalResult matchAndRewrite(ShapeCastOp shapeCastOp,
     auto newAttr =
         DenseElementsAttr::get(shapeCastOp.getType().cast<VectorType>(),
                                dense.getSplatValue<Attribute>());
     rewriter.replaceOpWithNewOp<arith::ConstantOp>(shapeCastOp, newAttr);
     return success();
   }
 };
 
+/// Pattern to rewrite a ShapeCast(Broadcast) -> Broadcast.
+/// This only applies when the shape of the broadcast source is a suffix of the
+/// shape of the result (i.e. when broadcast without reshape is expressive
+/// enough to capture the result in a single op).
+class ShapeCastBroadcastFolder final : public OpRewritePattern<ShapeCastOp> {
+public:
+  using OpRewritePattern<ShapeCastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ShapeCastOp shapeCastOp,
+                                PatternRewriter &rewriter) const override {
+    auto broadcastOp =
+        shapeCastOp.getSource().getDefiningOp<vector::BroadcastOp>();
+    if (!broadcastOp)
+      return failure();
+
+    auto broadcastSourceVectorType =
+        broadcastOp.getSourceType().dyn_cast<VectorType>();
+    auto broadcastSourceShape = broadcastSourceVectorType
+                                    ? broadcastSourceVectorType.getShape()
+                                    : ArrayRef<int64_t>{};
+    auto shapeCastTargetShape = shapeCastOp.getResultVectorType().getShape();
+
+    // Bail if `broadcastSourceShape` is not a suffix of the result.
+    bool isSuffix = (broadcastSourceShape == shapeCastTargetShape.take_back(
+                                                 broadcastSourceShape.size()));
+    if (!isSuffix)
+      return failure();
+
+    rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
+        shapeCastOp, shapeCastOp.getResultVectorType(),
+        broadcastOp.getSource());
+    return success();
+  }
+};
+
 } // namespace
 
 void ShapeCastOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
-  // Pattern to rewrite a ShapeCastOp(ConstantOp) -> ConstantOp.
-  results.add<ShapeCastConstantFolder>(context);
+  results.add<ShapeCastConstantFolder, ShapeCastBroadcastFolder>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // VectorBitCastOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult BitCastOp::verify() {
   auto sourceVectorType = getSourceVectorType();

mlir/test/Dialect/Vector/canonicalize.mlir

 func.func @fold_broadcast_shapecast(%arg0: vector<4xf32>) -> vector<4xf32> {
     %0 = vector.broadcast %arg0 : vector<4xf32> to vector<1x1x4xf32>
     %1 = vector.shape_cast %0 : vector<1x1x4xf32> to vector<4xf32>
     return %1 : vector<4xf32>
 }
 
 // -----
 
-// CHECK-LABEL: func @dont_fold_broadcast_shapecast_scalar
+// CHECK-LABEL: func @canonicalize_broadcast_shapecast_scalar
 //       CHECK:   vector.broadcast
-//       CHECK:   vector.shape_cast
-func.func @dont_fold_broadcast_shapecast_scalar(%arg0: f32) -> vector<1xf32> {
+//   CHECK-NOT:   vector.shape_cast
+func.func @canonicalize_broadcast_shapecast_scalar(%arg0: f32) -> vector<1xf32> {
     %0 = vector.broadcast %arg0 : f32 to vector<1x1x1xf32>
     %1 = vector.shape_cast %0 : vector<1x1x1xf32> to vector<1xf32>
     return %1 : vector<1xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @dont_fold_broadcast_shapecast_diff_shape
 //       CHECK:   vector.broadcast
 //       CHECK:   vector.shape_cast
 func.func @dont_fold_broadcast_shapecast_diff_shape(%arg0: vector<4xf32>) -> vector<8xf32> {
     %0 = vector.broadcast %arg0 : vector<4xf32> to vector<1x2x4xf32>
     %1 = vector.shape_cast %0 : vector<1x2x4xf32> to vector<8xf32>
     return %1 : vector<8xf32>
 }
 
 // -----
 
+// CHECK-LABEL: func @canonicalize_broadcast_shapecast
+//       CHECK:   vector.broadcast
+//   CHECK-NOT:   vector.shape_cast
+func.func @canonicalize_broadcast_shapecast(%arg0: vector<3xf32>) -> vector<8x3xf32> {
+    %0 = vector.broadcast %arg0 : vector<3xf32> to vector<2x4x3xf32>
+    %1 = vector.shape_cast %0 : vector<2x4x3xf32> to vector<8x3xf32>
+    return %1 : vector<8x3xf32>
+}
+
+// -----
+
 // CHECK-LABEL: fold_vector_transfers
 func.func @fold_vector_transfers(%A: memref<?x8xf32>) -> (vector<4x8xf32>, vector<4x9xf32>) {
   %c0 = arith.constant 0 : index
   %f0 = arith.constant 0.0 : f32
 
   // CHECK: vector.transfer_read %{{.*}} {in_bounds = [false, true]}
   %1 = vector.transfer_read %A[%c0, %c0], %f0 : memref<?x8xf32>, vector<4x8xf32>