additional fixes

Author: newling

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

@@ -1691,10 +1691,36 @@ static bool hasZeroDimVectors(Operation *op) {
          llvm::any_of(op->getResultTypes(), hasZeroDimVectorType);
 }
 
+/// vector.splat, and vector.shape_cast that just prepends 1's are
+/// special cases of vector.broadcast. This function returns true
+/// if \p op is one of these operations.
+static bool isBroadcastLike(Operation *op) {
+
+  if (isa<vector::BroadcastOp, SplatOp>(op))
+    return true;
+
+  // a shape_cast which just prepends 1's is broadcast-like.
+  auto shapeCast = dyn_cast<vector::ShapeCastOp>(op);
+  if (!shapeCast)
+    return false;
+
+  ArrayRef<int64_t> dstShape = shapeCast.getType().getShape();
+  ArrayRef<int64_t> srcShape = shapeCast.getSourceVectorType().getShape();
+
+  // A rank-reducing shape_cast cannot be broadcast-like.
+  if (srcShape.size() > dstShape.size())
+    return false;
+
+  bool isSuffix = (srcShape == dstShape.take_back(srcShape.size()));
+  return isSuffix;
+}
+
 /// Fold extractOp with scalar result coming from BroadcastOp or SplatOp.
-static Value foldExtractFromBroadcast(ExtractOp extractOp) {
+static Value foldExtractFromBroadcastLike(ExtractOp extractOp) {
+
   Operation *defOp = extractOp.getVector().getDefiningOp();
-  if (!defOp || !isa<vector::BroadcastOp, SplatOp>(defOp))
+
+  if (!defOp || !isBroadcastLike(defOp))
     return Value();
 
   Value source = defOp->getOperand(0);
@@ -1721,14 +1747,22 @@ static Value foldExtractFromBroadcast(ExtractOp extractOp) {
           broadcastVecType.getShape().take_back(extractResultRank))
     return Value();
 
-  auto broadcastOp = cast<vector::BroadcastOp>(defOp);
-  int64_t broadcastDstRank = broadcastOp.getResultVectorType().getRank();
+  assert(defOp->getNumResults() == 1 && "all broadcast-like ops have 1 result");
+  auto dstType = dyn_cast<VectorType>(defOp->getResult(0).getType());
+  assert(dstType && "all broadcast-like ops have vector results");
+
+  int64_t broadcastDstRank = dstType.getRank();
 
   // Detect all the positions that come from "dim-1" broadcasting.
-  // These dimensions correspond to "dim-1" broadcasted dims; set the mathching
+  // These dimensions correspond to "dim-1" broadcasted dims; set the matching
   // extract position to `0` when extracting from the source operand.
-  llvm::SetVector<int64_t> broadcastedUnitDims =
-      broadcastOp.computeBroadcastedUnitDims();
+  auto broadcastedUnitDims = [&]() -> llvm::SetVector<int64_t> {
+    if (auto broadcastOp = dyn_cast<BroadcastOp>(defOp)) {
+      return broadcastOp.computeBroadcastedUnitDims();
+    }
+    return {};
+  }();
+
   SmallVector<OpFoldResult> extractPos(extractOp.getMixedPosition());
   OpBuilder b(extractOp.getContext());
   int64_t broadcastRankDiff = broadcastDstRank - broadcastSrcRank;
@@ -2163,7 +2197,7 @@ OpFoldResult ExtractOp::fold(FoldAdaptor adaptor) {
     return getResult();
   if (auto res = ExtractFromInsertTransposeChainState(*this).fold())
     return res;
-  if (auto res = foldExtractFromBroadcast(*this))
+  if (auto res = foldExtractFromBroadcastLike(*this))
     return res;
   if (auto res = foldExtractFromShuffle(*this))
     return res;
@@ -2181,15 +2215,16 @@ OpFoldResult ExtractOp::fold(FoldAdaptor adaptor) {
 
 namespace {
 
-// Pattern to rewrite a ExtractOp(Broadcast) -> Broadcast.
+// Pattern to rewrite a ExtractOp(broadcast-like) -> Broadcast.
 class ExtractOpFromBroadcast final : public OpRewritePattern<ExtractOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ExtractOp extractOp,
                                 PatternRewriter &rewriter) const override {
     Operation *defOp = extractOp.getVector().getDefiningOp();
-    if (!defOp || !isa<vector::BroadcastOp, SplatOp>(defOp))
+
+    if (!defOp || !isBroadcastLike(defOp))
       return failure();
 
     Value source = defOp->getOperand(0);

mlir/test/Dialect/Vector/canonicalize.mlir

@@ -762,35 +762,55 @@ func.func @fold_extract_broadcast_0dvec_input_scalar_output(%a : vector<f32>,
 
 // -----
 
-
-// CHECK-LABEL: negative_fold_extract_broadcast
+// CHECK-LABEL: negative_fold_partial_extract_broadcast
 //       CHECK:   vector.broadcast %{{.*}} : vector<1x1xf32> to vector<1x2x4xf32>
 //       CHECK:   vector.extract %{{.*}}[0, 0] : vector<4xf32> from vector<1x2x4xf32>
-func.func @negative_fold_extract_broadcast(%a : vector<1x1xf32>) -> vector<4xf32> {
+func.func @negative_fold_partial_extract_broadcast(%a : vector<1x1xf32>) -> vector<4xf32> {
   %b = vector.broadcast %a : vector<1x1xf32> to vector<1x2x4xf32>
   %r = vector.extract %b[0, 0] : vector<4xf32> from vector<1x2x4xf32>
   return %r : vector<4xf32>
 }
 
 // -----
 
-// CHECK-LABEL: fold_extract_splat
+// CHECK-LABEL: negative_fold_full_extract_broadcast
+//       CHECK:   vector.broadcast %{{.*}} : vector<1x1xf32> to vector<1x1x4xf32>
+//       CHECK:   vector.shape_cast %{{.*}} : vector<1x1x4xf32> to vector<4xf32>
+func.func @negative_fold_full_extract_broadcast(%a : vector<1x1xf32>) -> vector<4xf32> {
+  %b = vector.broadcast %a : vector<1x1xf32> to vector<1x1x4xf32>
+  %r = vector.extract %b[0, 0] : vector<4xf32> from vector<1x1x4xf32>
+  return %r : vector<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: fold_extract_scalar_splat
 //  CHECK-SAME:   %[[A:.*]]: f32
 //       CHECK:   return %[[A]] : f32
-func.func @fold_extract_splat(%a : f32, %idx0 : index, %idx1 : index, %idx2 : index) -> f32 {
+func.func @fold_extract_scalar_splat(%a : f32, %idx0 : index, %idx1 : index, %idx2 : index) -> f32 {
   %b = vector.splat %a : vector<1x2x4xf32>
   %r = vector.extract %b[%idx0, %idx1, %idx2] : f32 from vector<1x2x4xf32>
   return %r : f32
 }
 
 // -----
 
-// CHECK-LABEL: fold_extract_broadcast_dim1_broadcasting
+// CHECK-LABEL: fold_extract_vector_splat
+//       CHECK:   vector.broadcast {{.*}} f32 to vector<4xf32>
+func.func @fold_extract_vector_splat(%a : f32, %idx0 : index, %idx1 : index) -> vector<4xf32> {
+  %b = vector.splat %a : vector<1x2x4xf32>
+  %r = vector.extract %b[%idx0, %idx1] : vector<4xf32> from vector<1x2x4xf32>
+  return %r : vector<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: fold_extract_broadcast_21_to_124
 //  CHECK-SAME:   %[[A:.*]]: vector<2x1xf32>
 //  CHECK-SAME:   %[[IDX:.*]]: index, %[[IDX1:.*]]: index, %[[IDX2:.*]]: index
 //       CHECK:   %[[R:.*]] = vector.extract %[[A]][%[[IDX1]], 0] : f32 from vector<2x1xf32>
 //       CHECK:   return %[[R]] : f32
-func.func @fold_extract_broadcast_dim1_broadcasting(%a : vector<2x1xf32>,
+func.func @fold_extract_broadcast_21_to_124(%a : vector<2x1xf32>,
   %idx : index, %idx1 : index, %idx2 : index) -> f32 {
   %b = vector.broadcast %a : vector<2x1xf32> to vector<1x2x4xf32>
   %r = vector.extract %b[%idx, %idx1, %idx2] : f32 from vector<1x2x4xf32>
@@ -799,6 +819,20 @@ func.func @fold_extract_broadcast_dim1_broadcasting(%a : vector<2x1xf32>,
 
 // -----
 
+// CHECK-LABEL: fold_extract_broadcast_21_to_224
+//  CHECK-SAME:   %[[A:.*]]: vector<2x1xf32>
+//  CHECK-SAME:   %[[IDX:.*]]: index, %[[IDX1:.*]]: index, %[[IDX2:.*]]: index
+//       CHECK:   %[[R:.*]] = vector.extract %[[A]][%[[IDX1]], 0] : f32 from vector<2x1xf32>
+//       CHECK:   return %[[R]] : f32
+func.func @fold_extract_broadcast_21_to_224(%a : vector<2x1xf32>,
+  %idx : index, %idx1 : index, %idx2 : index) -> f32 {
+  %b = vector.broadcast %a : vector<2x1xf32> to vector<2x2x4xf32>
+  %r = vector.extract %b[%idx, %idx1, %idx2] : f32 from vector<2x2x4xf32>
+  return %r : f32
+}
+
+// -----
+
 // CHECK-LABEL: fold_extract_broadcast_to_lower_rank
 //  CHECK-SAME:   %[[A:.*]]: vector<2x4xf32>
 //  CHECK-SAME:   %[[IDX0:.*]]: index, %[[IDX1:.*]]: index
@@ -1559,7 +1593,7 @@ func.func @negative_store_to_load_tensor_memref(
     %arg0 : tensor<?x?xf32>,
     %arg1 : memref<?x?xf32>,
     %v0 : vector<4x2xf32>
-  ) -> vector<4x2xf32> 
+  ) -> vector<4x2xf32>
 {
   %c0 = arith.constant 0 : index
   %cf0 = arith.constant 0.0 : f32
@@ -1616,7 +1650,7 @@ func.func @negative_store_to_load_tensor_broadcast_out_of_bounds(%arg0 : tensor<
 //       CHECK:   vector.transfer_read
 func.func @negative_store_to_load_tensor_broadcast_masked(
     %arg0 : tensor<?x?xf32>, %v0 : vector<4x2xf32>, %mask : vector<4x2xi1>)
-  -> vector<4x2x6xf32> 
+  -> vector<4x2x6xf32>
 {
   %c0 = arith.constant 0 : index
   %cf0 = arith.constant 0.0 : f32