[mlir][vector] cast away leading one dims for insert ops

[ThomasRaoux]

LGTM

mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp

16

do you need that?

[antiagainst]

mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp

16

removed

mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp

 //===- VectorDropLeadUnitDim.cpp - Conversion within the Vector dialect ---===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/TypeUtilities.h"
 
 #define DEBUG_TYPE "vector-drop-unit-dim"
 
 using namespace mlir;
 using namespace mlir::vector;
 
 // Trims leading one dimensions from `oldType` and returns the result type.
 // Returns `vector<1xT>` if `oldType` only has one element.
 static VectorType trimLeadingOneDims(VectorType oldType) {
   ArrayRef<int64_t> oldShape = oldType.getShape();
   ArrayRef<int64_t> newShape =
       oldShape.drop_while([](int64_t dim) { return dim == 1; });
   // Make sure we have at least 1 dimension per vector type requirements.
   if (newShape.empty())
     newShape = oldShape.take_back();
   return VectorType::get(newShape, oldType.getElementType());
 }
 
 /// Return a smallVector of size `rank` containing all zeros.
 static SmallVector<int64_t> splatZero(int64_t rank) {
   return SmallVector<int64_t>(rank, 0);
 }
 namespace {
 
 // Casts away leading one dimensions in vector.extract_strided_slice's vector
-// input by inserting vector.shape_cast.
+// input by inserting vector.broadcast.
 struct CastAwayExtractStridedSliceLeadingOneDim
     : public OpRewritePattern<vector::ExtractStridedSliceOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::ExtractStridedSliceOp extractOp,
                                 PatternRewriter &rewriter) const override {
     // vector.extract_strided_slice requires the input and output vector to have
     // the same rank. Here we drop leading one dimensions from the input vector
LogicalResult matchAndRewrite(vector::ExtractStridedSliceOp extractOp,
 
     rewriter.replaceOpWithNewOp<vector::BroadcastOp>(extractOp, oldDstType,
                                                      newExtractOp);
 
     return success();
   }
 };
 
-// Casts away leading one dimensions in vector.extract_strided_slice's vector
-// inputs by inserting vector.shape_cast.
+// Casts away leading one dimensions in vector.insert_strided_slice's vector
+// inputs by inserting vector.broadcast.
 struct CastAwayInsertStridedSliceLeadingOneDim
     : public OpRewritePattern<vector::InsertStridedSliceOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::InsertStridedSliceOp insertOp,
                                 PatternRewriter &rewriter) const override {
     VectorType oldSrcType = insertOp.getSourceVectorType();
     VectorType newSrcType = trimLeadingOneDims(oldSrcType);
LogicalResult matchAndRewrite(vector::InsertStridedSliceOp insertOp,
 
     rewriter.replaceOpWithNewOp<vector::BroadcastOp>(insertOp, oldDstType,
                                                      newInsertOp);
 
     return success();
   }
 };
 
+// Casts away leading one dimensions in vector.insert's vector inputs by
+// inserting vector.broadcast.
+struct CastAwayInsertLeadingOneDim : public OpRewritePattern<vector::InsertOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::InsertOp insertOp,
+                                PatternRewriter &rewriter) const override {
+    Type oldSrcType = insertOp.getSourceType();
+    Type newSrcType = oldSrcType;
+    int64_t oldSrcRank = 0, newSrcRank = 0;
+    if (auto type = oldSrcType.dyn_cast<VectorType>()) {
+      newSrcType = trimLeadingOneDims(type);
+      oldSrcRank = type.getRank();
+      newSrcRank = newSrcType.cast<VectorType>().getRank();
+    }
+
+    VectorType oldDstType = insertOp.getDestVectorType();
+    VectorType newDstType = trimLeadingOneDims(oldDstType);
+
+    int64_t srcDropCount = oldSrcRank - newSrcRank;
+    int64_t dstDropCount = oldDstType.getRank() - newDstType.getRank();
+    if (srcDropCount == 0 && dstDropCount == 0)
+      return failure();
+
+    // Trim leading one dimensions from both operands.
+    Location loc = insertOp.getLoc();
+
+    Value newSrcVector = insertOp.getSource();
+    if (oldSrcRank != 0) {
+      newSrcVector = rewriter.create<vector::ExtractOp>(
+          loc, insertOp.getSource(), splatZero(srcDropCount));
+    }
+    Value newDstVector = rewriter.create<vector::ExtractOp>(
+        loc, insertOp.getDest(), splatZero(dstDropCount));
+
+    unsigned oldPosRank = insertOp.getPosition().getValue().size();
+    unsigned newPosRank = newDstType.getRank() - newSrcRank;
+    SmallVector<Attribute> newPositions = llvm::to_vector(
+        insertOp.getPosition().getValue().take_back(newPosRank));
+    if (newPosRank > oldPosRank) {
+      auto zeroAttr = rewriter.getZeroAttr(rewriter.getI64Type());
+      newPositions.resize(newPosRank, zeroAttr);
+    }
+
+    auto newInsertOp = rewriter.create<vector::InsertOp>(
+        loc, newDstType, newSrcVector, newDstVector,
+        rewriter.getArrayAttr(newPositions));
+
+    rewriter.replaceOpWithNewOp<vector::BroadcastOp>(insertOp, oldDstType,
+                                                     newInsertOp);
+
+    return success();
+  }
+};
+
 // Turns vector.transfer_read on vector with leading 1 dimensions into
 // vector.shape_cast followed by vector.transfer_read on vector without leading
 // 1 dimensions.
 struct CastAwayTransferReadLeadingOneDim
     : public OpRewritePattern<vector::TransferReadOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::TransferReadOp read,

 };
 
 } // namespace
 
 void mlir::vector::populateCastAwayVectorLeadingOneDimPatterns(
     RewritePatternSet &patterns) {
   patterns
       .add<CastAwayExtractStridedSliceLeadingOneDim,
-           CastAwayInsertStridedSliceLeadingOneDim,
+           CastAwayInsertStridedSliceLeadingOneDim, CastAwayInsertLeadingOneDim,
            CastAwayTransferReadLeadingOneDim,
            CastAwayTransferWriteLeadingOneDim, CastAwayElementwiseLeadingOneDim,
            CastAwayContractionLeadingOneDim>(patterns.getContext());
   populateShapeCastFoldingPatterns(patterns);
 }

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

func @cast_away_elementwise_leading_one_dims(
   // CHECK:  vector.extract %{{.*}}[0] : vector<1x4xf32>
   // CHECK:  vector.extract %{{.*}}[0] : vector<1x4xf32>
   // CHECK:  select %arg4, %12, %{{.*}} : vector<4xf32>
   // CHECK:  vector.broadcast %{{.*}} : vector<4xf32> to vector<1x4xf32>
   %3 = arith.select %arg4, %arg3, %arg2 : vector<1x4xf32>
   return %0, %1, %2, %3: vector<1x1x8xf32>, vector<1x4xi1>, vector<1x4xf32>, vector<1x4xf32>
 }
 
+// CHECK-LABEL: func @cast_away_insert_leading_one_dims_scalar
+//  CHECK-SAME: (%[[S:.+]]: f32, %[[V:.+]]: vector<1x1x4xf32>)
+//       CHECK:   %[[EXTRACT:.+]] = vector.extract %[[V]][0, 0] : vector<1x1x4xf32>
+//       CHECK:   %[[INSERT:.+]] = vector.insert %[[S]], %[[EXTRACT]] [0] : f32 into vector<4xf32>
+//       CHECK:   %[[BCAST:.+]] = vector.broadcast %[[INSERT]] : vector<4xf32> to vector<1x1x4xf32>
+//       CHECK:   return %[[BCAST]]
+func @cast_away_insert_leading_one_dims_scalar(%s: f32, %v: vector<1x1x4xf32>) -> vector<1x1x4xf32> {
+  %0 = vector.insert %s, %v [0, 0, 0] : f32 into vector<1x1x4xf32>
+  return %0: vector<1x1x4xf32>
+}
+
+// CHECK-LABEL: func @cast_away_insert_leading_one_dims_rank1
+//  CHECK-SAME: (%[[S:.+]]: vector<4xf32>, %[[V:.+]]: vector<1x1x4xf32>)
+//       CHECK:   %[[BCAST:.+]] = vector.broadcast %[[S]] : vector<4xf32> to vector<1x1x4xf32>
+//       CHECK:   return %[[BCAST]]
+func @cast_away_insert_leading_one_dims_rank1(%s: vector<4xf32>, %v: vector<1x1x4xf32>) -> vector<1x1x4xf32> {
+  %0 = vector.insert %s, %v [0, 0] : vector<4xf32> into vector<1x1x4xf32>
+  return %0: vector<1x1x4xf32>
+}
+
+// CHECK-LABEL: func @cast_away_insert_leading_one_dims_rank2
+//  CHECK-SAME: (%[[S:.+]]: vector<1x4xf32>, %[[V:.+]]: vector<1x1x4xf32>)
+//       CHECK:   %[[EXTRACT:.+]] = vector.extract %[[S]][0] : vector<1x4xf32>
+//       CHECK:   %[[BCAST:.+]] = vector.broadcast %[[EXTRACT]] : vector<4xf32> to vector<1x1x4xf32>
+//       CHECK:   return %[[BCAST]]
+func @cast_away_insert_leading_one_dims_rank2(%s: vector<1x4xf32>, %v: vector<1x1x4xf32>) -> vector<1x1x4xf32> {
+  %0 = vector.insert %s, %v [0] : vector<1x4xf32> into vector<1x1x4xf32>
+  return %0: vector<1x1x4xf32>
+}
+
+// CHECK-LABEL: func @cast_away_insert_leading_one_dims_non_one_dest
+//  CHECK-SAME: (%[[S:.+]]: vector<1x4xf32>, %[[V:.+]]: vector<8x1x4xf32>)
+//       CHECK:   %[[EXTRACT:.+]] = vector.extract %[[S]][0] : vector<1x4xf32>
+//       CHECK:   %[[INSERT:.+]] = vector.insert %[[EXTRACT]], %[[V]] [5, 0] : vector<4xf32> into vector<8x1x4xf32>
+//       CHECK:   return %[[INSERT]]
+func @cast_away_insert_leading_one_dims_non_one_dest(%s: vector<1x4xf32>, %v: vector<8x1x4xf32>) -> vector<8x1x4xf32> {
+  %0 = vector.insert %s, %v [5] : vector<1x4xf32> into vector<8x1x4xf32>
+  return %0: vector<8x1x4xf32>
+}