[mlir][bufferize] Support parallel_insert_slice in EmptyTensorElimination

mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp

 ///
 /// Starting from an InsertSliceOp, an tensor::EmptyOp at the end of the insert
 /// source's reverse use-def chain is eliminated if:
 /// * On the reverse use-def chain path from the InsertSliceOp to the
 ///   tensor::EmptyOp, all ops were decided to bufferize inplace and the buffer
 ///   relation is "equivalent" (TODO: can be relaxed if needed).
 /// * The reverse use-def chain has exactly one end, which is the
 ///   tensor::EmptyOp.
-LogicalResult
-mlir::bufferization::insertSliceAnchoredEmptyTensorEliminationStep(
+template <typename OpTy>
+static LogicalResult insertSliceLikeAnchoredEmptyTensorEliminationStep(
     RewriterBase &rewriter, Operation *op, AnalysisState &state) {
   return eliminateEmptyTensors(
       rewriter, op, state,
       /*anchorMatchFunc=*/
       [&](OpOperand &operand, SmallVector<Value> &neededValues) {
-        auto insertSliceOp =
-            dyn_cast<tensor::InsertSliceOp>(operand.getOwner());
+        auto insertSliceOp = dyn_cast<OpTy>(operand.getOwner());
         if (!insertSliceOp)
           return false;
         if (&operand != &insertSliceOp->getOpOperand(0) /*source*/)
           return false;
 
         // Collect all values that are needed to construct the replacement op.
         neededValues.append(insertSliceOp.getOffsets().begin(),
                             insertSliceOp.getOffsets().end());
         neededValues.append(insertSliceOp.getSizes().begin(),
                             insertSliceOp.getSizes().end());
         neededValues.append(insertSliceOp.getStrides().begin(),
                             insertSliceOp.getStrides().end());
         neededValues.push_back(insertSliceOp.getDest());
 
         return true;
       },
       /*rewriteFunc=*/
       [](OpBuilder &b, Location loc, OpOperand &operand) {
-        auto insertOp = cast<tensor::InsertSliceOp>(operand.getOwner());
+        auto insertOp = cast<OpTy>(operand.getOwner());
         auto extractOp = b.create<tensor::ExtractSliceOp>(
             loc, insertOp.getSourceType(), insertOp.getDest(),
             insertOp.getMixedOffsets(), insertOp.getMixedSizes(),
             insertOp.getMixedStrides());
         return extractOp.getResult();
       });
 }
 
+LogicalResult
+mlir::bufferization::insertSliceAnchoredEmptyTensorEliminationStep(
+    RewriterBase &rewriter, Operation *op, AnalysisState &state) {
+  if (failed(insertSliceLikeAnchoredEmptyTensorEliminationStep<
+             tensor::InsertSliceOp>(rewriter, op, state)))
+    return failure();
+  if (failed(insertSliceLikeAnchoredEmptyTensorEliminationStep<
+             tensor::ParallelInsertSliceOp>(rewriter, op, state)))
+    return failure();
+  return success();
+}
+
 namespace {
 struct EmptyTensorElimination
     : public bufferization::impl::EmptyTensorEliminationBase<
           EmptyTensorElimination> {
   EmptyTensorElimination() = default;
 
   void runOnOperation() override;
 

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir

func.func @shape_mismatch(%t: tensor<5x6x128xf32>) -> tensor<5x6x128xf32> {
   %0 = tensor.empty() : tensor<128xf32>
   %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<128xf32>) -> tensor<128xf32>
   %2 = tensor.expand_shape %1 [[0, 1, 2]]
       : tensor<128xf32> into tensor<1x1x128xf32>
   %3 = tensor.insert_slice %2 into %t[2, 3, 0][1, 1, 128][1, 1, 1]
       : tensor<1x1x128xf32> into tensor<5x6x128xf32>
   return %3 : tensor<5x6x128xf32>
 }
+
+// -----
+
+//      CHECK: func @parallel_insert_slice(
+// CHECK-SAME:   %[[FUNC_ARG:[0-9a-zA-Z]*]]: memref<?xf32>
+// CHECK-SAME:   %[[sz:[0-9a-zA-Z]*]]: index
+func.func @parallel_insert_slice(
+  %t: tensor<?xf32> {bufferization.buffer_layout = affine_map<(d0) -> (d0)>, bufferization.writable = true},
+  %sz: index)
+    -> (tensor<?xf32>)
+{
+  %f0 = arith.constant 0.0: f32
+  %c512 = arith.constant 512 : index
+
+  %r1 = scf.foreach_thread (%iv) in (%c512) shared_outs(%o = %t) -> (tensor<?xf32>) {
+    // tensor.empty itself does not alloc but forwards to the insert_slice.
+    // EmptyTensorOpElimination replaces the tensor.empty with an inplace
+    // extract_slice.
+    // CHECK: %[[T_SUBVIEW:.*]] =  memref.subview %[[FUNC_ARG]][42] [%[[sz]]] [1]
+    %a = tensor.empty(%sz) : tensor<?xf32>
+
+    // CHECK: linalg.fill ins({{.*}} : f32) outs(%[[T_SUBVIEW]] : memref<?xf32
+    %f = linalg.fill ins(%f0 : f32) outs(%a : tensor<?xf32>) -> tensor<?xf32>
+
+    // Self-copy canonicalizes away later.
+    scf.foreach_thread.perform_concurrently {
+      tensor.parallel_insert_slice %f into %o[42][%sz][1]: tensor<?xf32> into tensor<?xf32>
+    }
+  }
+
+  return %r1: tensor<?xf32>
+}