[mlir][Linalg] Support tensor.parallel_insert_slice in transform.insert_slice_to_copy

[ftynse]

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

3222–3224

Nit: llvm::is_one_of has a variadic template argument.

3231–3237

Nit: unfinished comment .

3232

Nit: if constexpr for more efficient code.

3262

Nit: targetOp should never be null, so just dyn_cast should suffice here.

[nicolasvasilache]

Address comments.

[nicolasvasilache]

Missed a spot

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

def InsertSliceToCopyOp :
 
   let assemblyFormat = "$target attr-dict `:` functional-type(operands, results) ";
 
   let builders = [
     OpBuilder<(ins "Value":$target)>,
   ];
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
-        ::mlir::tensor::InsertSliceOp target,
+        ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 #endif // LINALG_TRANSFORM_OPS

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/TilingInterface.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
+#include <type_traits>
 
 using namespace mlir;
 using namespace mlir::linalg;
 using namespace mlir::transform;
 
 #define DEBUG_TYPE "linalg-transforms"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
 #define DBGSNL() (llvm::dbgs() << "\n")
transform::HoistRedundantTensorSubsetsOp::applyToOne(
   });
   results.push_back(target);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // InsertSliceToCopyOp
 //===----------------------------------------------------------------------===//
-
-DiagnosedSilenceableFailure transform::InsertSliceToCopyOp::applyToOne(
-    tensor::InsertSliceOp target, transform::ApplyToEachResultList &results,
+template <typename OpTy>
+DiagnosedSilenceableFailure doit(RewriterBase &rewriter, OpTy target,
+                                 transform::ApplyToEachResultList &results,
                                  transform::TransformState &state) {
-  if (auto copySource = target.getSource().getDefiningOp<linalg::CopyOp>()) {
+  static_assert(llvm::is_one_of<OpTy, tensor::InsertSliceOp, tensor::ParallelInsertSliceOp>) &&
+                    "wrong op type");
+
+  if (auto copySource =
+          target.getSource().template getDefiningOp<linalg::CopyOp>()) {
     results.push_back(copySource);
     return DiagnosedSilenceableFailure::success();
   }
 
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(target->getContext(), &listener);
-  rewriter.setInsertionPoint(target);
+  // If we are inside an InParallel region, temporarily set the insertion point
+  // outside: only tensor.parallel_insert_slice ops are allowed in there.
+  if (constexpr std::is_same_v<OpTy, tensor::ParallelInsertSliceOp>) {
+    rewriter.setInsertionPoint(
+        target->template getParentOfType<scf::InParallelOp>());
+  }
+
   Value extracted = rewriter.create<tensor::ExtractSliceOp>(
       target.getLoc(), target.getDest(), target.getMixedOffsets(),
       target.getMixedSizes(), target.getMixedStrides());
   Value copied = rewriter
                      .create<linalg::CopyOp>(target.getLoc(),
                                              target.getSource(), extracted)
                      .getResult(0);
-  rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
+  // Reset the insertion point.
+  rewriter.setInsertionPoint(target);
+  rewriter.replaceOpWithNewOp<OpTy>(
       target, copied, target.getDest(), target.getMixedOffsets(),
       target.getMixedSizes(), target.getMixedStrides());
 
   results.push_back(copied.getDefiningOp());
   return DiagnosedSilenceableFailure::success();
 }
 
+DiagnosedSilenceableFailure transform::InsertSliceToCopyOp::applyToOne(
+    Operation *targetOp, transform::ApplyToEachResultList &results,
+    transform::TransformState &state) {
+
+  TrackingListener listener(state, *this);
+  IRRewriter rewriter(targetOp->getContext(), &listener);
+  rewriter.setInsertionPoint(targetOp);
+  if (auto target = dyn_cast<tensor::InsertSliceOp>(targetOp))
+    return doit(rewriter, target, results, state);
+  if (auto target = dyn_cast<tensor::ParallelInsertSliceOp>(targetOp))
+    return doit(rewriter, target, results, state);
+
+  DiagnosedSilenceableFailure diag =
+      emitSilenceableError()
+      << "only InsertSliceOp and ParallelInsertSliceOp ops are supported";
+  diag.attachNote(targetOp->getLoc()) << "target op";
+  return diag;
+}
+
 //===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// Registers new ops and declares PDL as dependent dialect since the
 /// additional ops are using PDL types for operands and results.
 class LinalgTransformDialectExtension

mlir/test/Dialect/Linalg/transform-op-insert-slice-to-copy.mlir

-// RUN: mlir-opt -test-transform-dialect-interpreter %s --split-input-file | FileCheck %s
+// RUN: mlir-opt -test-transform-dialect-interpreter %s --split-input-file --allow-unregistered-dialect | FileCheck %s
 
 // CHECK-LABEL: func @insert_slice_to_copy
     // CHECK-SAME: %[[I:.*]]: tensor<2x3xf32>
     // CHECK-SAME: %[[O:.*]]: tensor<?x?xf32>, 
     // CHECK-SAME: %[[OFF0:[0-9a-zA-Z]+]]: index,
     // CHECK-SAME: %[[OFF1:[0-9a-zA-Z]+]]: index,
     // CHECK-SAME: %[[SZ0:[0-9a-zA-Z]+]]: index,
     // CHECK-SAME: %[[SZ1:[0-9a-zA-Z]+]]: index,

 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1 = transform.structured.insert_slice_to_copy %0 : (!transform.any_op) -> !transform.any_op
   transform.cast %1 : !transform.any_op to !transform.op<"linalg.copy">
 }
 
+// -----
+
+// CHECK-LABEL: func @parallel_insert_slice_to_copy
+func.func @parallel_insert_slice_to_copy(%out : tensor<?x?xf32>, %sz0: index, %sz1: index) {
+  %0 = scf.forall (%arg0, %arg1) in (27, 8) shared_outs(%arg2 = %out) -> (tensor<?x?xf32>) {
+    %t = "make_me_a_tensor"() : () -> (tensor<?x?xf32> )
+
+    //      CHECK: tensor.extract_slice
+    //      CHECK: linalg.copy
+    //      CHECK: scf.forall.in_parallel
+    //      CHECK:   tensor.parallel_insert_slice
+    scf.forall.in_parallel {
+      tensor.parallel_insert_slice %t into %arg2[0, 0] [%sz0, %sz1] [1, 1] 
+        : tensor<?x?xf32> into tensor<?x?xf32>
+    }
+  }
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %arg1
+    : (!transform.any_op) -> !transform.any_op
+  %1 = transform.structured.insert_slice_to_copy %0
+    : (!transform.any_op) -> !transform.any_op
+  transform.cast %1 : !transform.any_op to !transform.op<"linalg.copy">
+}