[mlir][vector] Fold cancelling vector.shape_cast(vector.broadcast)

[antiagainst]

Summary

vector.broadcast can inject all size one dimensions. If it's
followed by a vector.shape_cast to the original type, we can
cancel the op pair, like cancelling consecutive shape_cast ops.

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

LogicalResult ShapeCastOp::verify() {
   // Check if source/result are of vector type.
   if (sourceVectorType && resultVectorType)
     return verifyVectorShapeCast(*this, sourceVectorType, resultVectorType);
 
   return success();
 }
 
 OpFoldResult ShapeCastOp::fold(ArrayRef<Attribute> operands) {
-  // Nop shape cast.
+  // No-op shape cast.
   if (getSource().getType() == getResult().getType())
     return getSource();
 
   // Canceling shape casts.
   if (auto otherOp = getSource().getDefiningOp<ShapeCastOp>()) {
     if (getResult().getType() == otherOp.getSource().getType())
       return otherOp.getSource();
 
     // Only allows valid transitive folding.
     VectorType srcType = otherOp.getSource().getType().cast<VectorType>();
     VectorType resultType = getResult().getType().cast<VectorType>();
     if (srcType.getRank() < resultType.getRank()) {
       if (!isValidShapeCast(srcType.getShape(), resultType.getShape()))
         return {};
     } else if (srcType.getRank() > resultType.getRank()) {
       if (!isValidShapeCast(resultType.getShape(), srcType.getShape()))
         return {};
     } else {
       return {};
     }
 
     setOperand(otherOp.getSource());
     return getResult();
   }
+
+  // Cancelling broadcast and shape cast ops.
+  if (auto bcastOp = getSource().getDefiningOp<BroadcastOp>()) {
+    if (bcastOp.getSourceType() == getType())
+      return bcastOp.getSource();
+  }
+
   return {};
 }
 
 namespace {
 // Pattern to rewrite a ShapeCast(splat ConstantOp) -> ConstantOp.
 class ShapeCastConstantFolder final : public OpRewritePattern<ShapeCastOp> {
 public:
   using OpRewritePattern<ShapeCastOp>::OpRewritePattern;

mlir/test/Dialect/Vector/canonicalize.mlir

 func.func @dont_fold_expand_collapse(%arg0: vector<1x1x64xf32>) -> vector<8x8xf32> {
     %0 = vector.shape_cast %arg0 : vector<1x1x64xf32> to vector<1x1x8x8xf32>
     %1 = vector.shape_cast %0 : vector<1x1x8x8xf32> to vector<8x8xf32>
     return %1 : vector<8x8xf32>
 }
 
 // -----
 
+// CHECK-LABEL: func @fold_broadcast_shapecast
+//  CHECK-SAME: (%[[V:.+]]: vector<4xf32>)
+//       CHECK:   return %[[V]]
+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:   vector.broadcast
+//       CHECK:   vector.shape_cast
+func @dont_fold_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 @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: 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>