[mlir][tensor] Support parallel_insert_slice in MergeConsecutiveInsertExtractSlicePatterns.cpp

mlir/lib/Dialect/Tensor/Transforms/MergeConsecutiveInsertExtractSlicePatterns.cpp

LogicalResult matchAndRewrite(ExtractSliceOp nextOp,
     rewriter.replaceOpWithNewOp<ExtractSliceOp>(nextOp, nextOp.getType(),
                                                 prevOp.getSource(), newOffsets,
                                                 newSizes, newStrides);
     return success();
   }
 };
 
 /// Merges consecutive tensor.insert_slice ops into one.
-struct MergeConsecutiveInsertSlice : public OpRewritePattern<InsertSliceOp> {
-  using OpRewritePattern::OpRewritePattern;
+template <typename OpTy>
+struct MergeConsecutiveInsertSlice : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(InsertSliceOp nextOp,
+  LogicalResult matchAndRewrite(OpTy nextOp,
                                 PatternRewriter &rewriter) const override {
-    auto prevOp = nextOp.getSource().getDefiningOp<InsertSliceOp>();
+    auto prevOp = nextOp.getSource().template getDefiningOp<InsertSliceOp>();
     if (!prevOp)
       return failure();
 
     if (!prevOp.hasUnitStride() || !nextOp.hasUnitStride())
       return failure();
 
     // The first insert_slice op should be rank reducing to make sure we cover
     // the full source tensor to be inserted in the second insert_slice op.
     SliceVerificationResult result =
         isRankReducedType(prevOp.getDestType(), prevOp.getSourceType());
     if (result != SliceVerificationResult::Success)
       return failure();
 
     // Dynamic dimensions can pass rank reducing check in the above, e.g,
     // inserting <?xf32> into <1x?x1xf32>. For such cases we cannot be certain
     // the dynamic size covers the full tensor.
     if (!prevOp.getSourceType().hasStaticShape() ||
         !prevOp.getDestType().hasStaticShape())
       return failure();
 
-    rewriter.replaceOpWithNewOp<InsertSliceOp>(
+    rewriter.replaceOpWithNewOp<OpTy>(
         nextOp, prevOp.getSource(), nextOp.getDest(), nextOp.getMixedOffsets(),
         nextOp.getMixedSizes(), nextOp.getMixedStrides());
     return success();
   }
 };
 } // namespace
 
 void mlir::tensor::populateMergeConsecutiveInsertExtractSlicePatterns(
     RewritePatternSet &patterns) {
-  patterns.add<MergeConsecutiveExtractSlice, MergeConsecutiveInsertSlice>(
+  patterns.add<MergeConsecutiveExtractSlice,
+               MergeConsecutiveInsertSlice<InsertSliceOp>,
+               MergeConsecutiveInsertSlice<ParallelInsertSliceOp>>(
       patterns.getContext());
 }

mlir/test/Dialect/Tensor/fold-consecutive-insert-extract-slice.mlir

func.func @insert_slice_rank_reducing_dynamic_shape(
     %dst: tensor<128x128x128x128xf32>, %mid: tensor<1x?x1xf32>, %src: tensor<?xf32>, %offset: index, %size: index) -> tensor<128x128x128x128xf32> {
   %0 = tensor.insert_slice %src into %mid[0, 0, 0] [1, %size, 1] [1, 1, 1] : tensor<?xf32> into tensor<1x?x1xf32>
   %1 = tensor.insert_slice %0 into %dst[6, 7, 8, %offset] [1, 1, %size, 1] [1, 1, 1, 1] : tensor<1x?x1xf32> into tensor<128x128x128x128xf32>
   return %1: tensor<128x128x128x128xf32>
 }
 
 //   CHECK-LABEL: func.func @insert_slice_rank_reducing_dynamic_shape
 // CHECK-COUNT-2:   tensor.insert_slice
+
+// -----
+
+// CHECK-LABEL: func.func @parallel_insert_slice
+//   CHECK-NOT:   tensor.insert_slice
+//       CHECK:   tensor.parallel_insert_slice %{{.*}} into %{{.*}}[%{{.*}}, %{{.*}}] [1, 1] [1, 1] : tensor<f32> into tensor<1x2xf32>
+func.func @parallel_insert_slice(%t0: tensor<1x2xf32>, %t1: tensor<f32>, %t2: tensor<1x1xf32>) -> tensor<1x2xf32> {
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %r = scf.foreach_thread (%arg2, %arg3) in (%c1, %c2) shared_outs(%arg4 = %t0) -> (tensor<1x2xf32>) {
+    %inserted_slice = tensor.insert_slice %t1 into %t2[0, 0] [1, 1] [1, 1] : tensor<f32> into tensor<1x1xf32>
+    scf.foreach_thread.perform_concurrently {
+      tensor.parallel_insert_slice %inserted_slice into %arg4[%arg2, %arg3] [1, 1] [1, 1] : tensor<1x1xf32> into tensor<1x2xf32>
+    }
+  }
+  return %r : tensor<1x2xf32>
+}