[mlir][vector] add pattern to cast away lead unit dimension for broadcast op

[mravishankar]

mlir/test/Dialect/Vector/vector-transforms.mlir

679

Shouldnt this be just dropped? It seems like a no-op.

[ThomasRaoux]

mlir/test/Dialect/Vector/vector-transforms.mlir

679

Best would be to add a separate folding. Look like it isn't there right now. It would be better than adding it to this pattern. What do you think?

mlir/lib/Dialect/Vector/VectorTransforms.cpp

auto newVector = rewriter.create<vector::ShapeCastOp>(
         write.getLoc(), newType, write.vector());
     rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
         write, newVector, write.source(), write.indices(), newMap, inBounds);
 
     return success();
   }
 };
 
+struct CastAwayBrodcastLeadingOneDim
+    : public OpRewritePattern<vector::BroadcastOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::BroadcastOp broadcastOp,
+                                PatternRewriter &rewriter) const override {
+    VectorType newDstType = trimLeadingOneDims(broadcastOp.getVectorType());
+    if (newDstType == broadcastOp.getVectorType())
+      return failure();
+    Location loc = broadcastOp.getLoc();
+    VectorType srcVecType = broadcastOp.getSourceType().dyn_cast<VectorType>();
+    if (srcVecType)
+      srcVecType = trimLeadingOneDims(srcVecType);
+    Value source = broadcastOp.source();
+    if (srcVecType && srcVecType != broadcastOp.getSourceType()) {
+      source = rewriter.create<vector::ShapeCastOp>(loc, srcVecType, source);
+    }
+    Value newBroadcastOp =
+        rewriter.create<vector::BroadcastOp>(loc, newDstType, source);
+    rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(
+        broadcastOp, broadcastOp.getVectorType(), newBroadcastOp);
+    return success();
+  }
+};
+
 // Returns the values in `arrayAttr` as an integer vector.
 static SmallVector<int64_t, 4> getIntValueVector(ArrayAttr arrayAttr) {
   return llvm::to_vector<4>(
       llvm::map_range(arrayAttr.getAsRange<IntegerAttr>(),
                       [](IntegerAttr attr) { return attr.getInt(); }));
 }
 
 // Shuffles vector.bitcast op after vector.extract op.
patterns.add<SplitTransferReadOp, SplitTransferWriteOp>(patterns.getContext(),
                                                           ignoreFilter);
 }
 
 void mlir::vector::populateCastAwayVectorLeadingOneDimPatterns(
     RewritePatternSet &patterns) {
   patterns.add<CastAwayExtractStridedSliceLeadingOneDim,
                CastAwayInsertStridedSliceLeadingOneDim,
                CastAwayTransferReadLeadingOneDim,
-               CastAwayTransferWriteLeadingOneDim, ShapeCastOpFolder>(
+               CastAwayTransferWriteLeadingOneDim,
+               CastAwayBrodcastLeadingOneDim, ShapeCastOpFolder>(
       patterns.getContext());
 }
 
 void mlir::vector::populateBubbleVectorBitCastOpPatterns(
     RewritePatternSet &patterns) {
   patterns.add<BubbleDownVectorBitCastForExtract,
                BubbleDownBitCastForStridedSliceExtract,
                BubbleUpBitCastForStridedSliceInsert>(patterns.getContext());

mlir/test/Dialect/Vector/vector-transforms.mlir

 // CHECK-LABEL: func @cast_away_transfer_write_leading_one_dims_one_element
 func @cast_away_transfer_write_leading_one_dims_one_element(%arg0: memref<1x1x1x1xf16>, %arg1: vector<1x1xf16>) {
   %c0 = constant 0 : index
   // CHECK: vector.shape_cast %{{.+}} : vector<1x1xf16> to vector<1xf16>
   vector.transfer_write %arg1, %arg0[%c0, %c0, %c0, %c0] {in_bounds = [true, true]} : vector<1x1xf16>, memref<1x1x1x1xf16>
   return
 }
 
+// CHECK-LABEL: func @cast_away_broadcast_leading_one_dims
+func @cast_away_broadcast_leading_one_dims(
+  %arg0: vector<8xf32>, %arg1: f32, %arg2: vector<1x4xf32>) ->
+  (vector<1x1x8xf32>, vector<1x1x4xf32>, vector<1x3x4xf32>) {
+  // CHECK:  vector.broadcast %{{.*}} : vector<8xf32> to vector<8xf32>
+  // CHECK:  vector.shape_cast %{{.*}} : vector<8xf32> to vector<1x1x8xf32>
+  %0 = vector.broadcast %arg0 : vector<8xf32> to vector<1x1x8xf32>
+  // CHECK:  vector.broadcast %{{.*}} : f32 to vector<4xf32>
+  // CHECK:  vector.shape_cast %{{.*}} : vector<4xf32> to vector<1x1x4xf32>
+  %1 = vector.broadcast %arg1 : f32 to vector<1x1x4xf32>
+  // CHECK:  vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
+  // CHECK:  vector.broadcast %{{.*}} : vector<4xf32> to vector<3x4xf32>
+  // CHECK:  vector.shape_cast %{{.*}} : vector<3x4xf32> to vector<1x3x4xf32>
+  %2 = vector.broadcast %arg2 : vector<1x4xf32> to vector<1x3x4xf32>
+  return %0, %1, %2: vector<1x1x8xf32>, vector<1x1x4xf32>, vector<1x3x4xf32>
+}
+
 // CHECK-LABEL: func @bubble_down_bitcast_in_extract
 //  CHECK-SAME: %[[SRC:.+]]: vector<4xf32>
 func @bubble_down_bitcast_in_extract(%src: vector<4xf32>) -> (f16, f16) {
   %0 = vector.bitcast %src : vector<4xf32> to vector<8xf16>
   // CHECK: %[[EXTRACT1:.+]] = vector.extract %[[SRC]][1] : vector<4xf32>
   // CHECK:    %[[CAST1:.+]] = vector.bitcast %[[EXTRACT1]] : vector<1xf32> to vector<2xf16>
   // CHECK: %[[EXTRACT2:.+]] = vector.extract %[[CAST1]][1] : vector<2xf16>
   %1 = vector.extract %0[3] : vector<8xf16>