[mlir][vector] Extend pattern to trim lead unit dimension to Splat Op

[ThomasRaoux]

Summary

Extend the broadcast pattern to support also SplatOp. SplatOp on vector handles a subset of what vector.broadcast does.

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 CastAwayBroadcastLeadingOneDim
-    : 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())
+template <typename BroadCastType>
+struct CastAwayBroadcastLeadingOneDim : public OpRewritePattern<BroadCastType> {
+  using OpRewritePattern<BroadCastType>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(BroadCastType broadcastOp,
+                                PatternRewriter &rewriter) const override {
+    VectorType dstType =
+        broadcastOp.getResult().getType().template dyn_cast<VectorType>();
+    if (!dstType)
+      return failure();
+    VectorType newDstType = trimLeadingOneDims(dstType);
+    if (newDstType == dstType)
       return failure();
     Location loc = broadcastOp.getLoc();
-    VectorType srcVecType = broadcastOp.getSourceType().dyn_cast<VectorType>();
+    Value source = broadcastOp->getOperand(0);
+    VectorType srcVecType = source.getType().template dyn_cast<VectorType>();
     if (srcVecType)
       srcVecType = trimLeadingOneDims(srcVecType);
-    Value source = broadcastOp.source();
-    if (srcVecType && srcVecType != broadcastOp.getSourceType()) {
+    if (srcVecType && srcVecType != source.getType()) {
       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);
+        rewriter.create<BroadCastType>(loc, newDstType, source);
+    rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(broadcastOp, dstType,
+                                                     newBroadcastOp);
     return success();
   }
 };
 
 class CastAwayElementwiseLeadingOneDim : public RewritePattern {
 public:
   CastAwayElementwiseLeadingOneDim(MLIRContext *context)
       : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
void mlir::vector::populateSplitVectorTransferPatterns(
     RewritePatternSet &patterns,
     std::function<bool(Operation *)> ignoreFilter) {
   patterns.add<SplitTransferReadOp, SplitTransferWriteOp>(patterns.getContext(),
                                                           ignoreFilter);
 }
 
 void mlir::vector::populateCastAwayVectorLeadingOneDimPatterns(
     RewritePatternSet &patterns) {
-  patterns
-      .add<CastAwayExtractStridedSliceLeadingOneDim,
-           CastAwayInsertStridedSliceLeadingOneDim,
-           CastAwayTransferReadLeadingOneDim,
-           CastAwayTransferWriteLeadingOneDim, CastAwayBroadcastLeadingOneDim,
-           CastAwayElementwiseLeadingOneDim, ShapeCastOpFolder>(
-          patterns.getContext());
+  patterns.add<
+      CastAwayExtractStridedSliceLeadingOneDim,
+      CastAwayTransferReadLeadingOneDim, CastAwayTransferWriteLeadingOneDim,
+      CastAwayBroadcastLeadingOneDim<vector::BroadcastOp>,
+      CastAwayBroadcastLeadingOneDim<SplatOp> CastAwayElementwiseLeadingOneDim,
+      ShapeCastOpFolder>(patterns.getContext());
 }
 
 void mlir::vector::populateBubbleVectorBitCastOpPatterns(
     RewritePatternSet &patterns) {
   patterns.add<BubbleDownVectorBitCastForExtract,
                BubbleDownBitCastForStridedSliceExtract,
                BubbleUpBitCastForStridedSliceInsert>(patterns.getContext());
 }

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

func @cast_away_transfer_write_leading_one_dims_one_element(%arg0: memref<1x1x1x1xf16>, %arg1: vector<1x1xf16>) {
   // 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>) {
+  (vector<1x1x8xf32>, vector<1x1x4xf32>, vector<1x3x4xf32>, vector<1x1x4xf32>) {
   // 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:  splat %{{.*}} : vector<4xf32>
+  // CHECK:  vector.shape_cast %{{.*}} : vector<4xf32> to vector<1x1x4xf32>
+  %3 = splat %arg1 : vector<1x1x4xf32>
+  return %0, %1, %2, %3: vector<1x1x8xf32>, vector<1x1x4xf32>, vector<1x3x4xf32>, vector<1x1x4xf32>
 }
 
 // CHECK-LABEL: func @cast_away_elementwise_leading_one_dims
 func @cast_away_elementwise_leading_one_dims(
   %arg0: vector<1x1x8xf32>, %arg1: f32, %arg2: vector<1x4xf32>,
   %arg3: vector<1x4xf32>, %arg4: i1) ->
   (vector<1x1x8xf32>, vector<1x4xi1>, vector<1x4xf32>, vector<1x4xf32>) {
   // CHECK:  vector.shape_cast %{{.*}} : vector<1x1x8xf32> to vector<8xf32>