[mlir] properly fix concurrent transform interpreter pass base

[ftynse]

Summary

The original implementation of the transform interpreter pass base was
cloning the entire transform IR in presence of PDL-related operations to
avoid concurrency issues when running the pass with the same transform
IR on multiple operations of the payload IR. The root cause of those
issues is the transform.pdl_match operation that was moving the PDL
pattern definition operation into a new module, consumed by the PDL
interpreter and leading to a race. Clone the pattern operation instead.
This avoids the race as well as the cost for transform IR that doesn't
use PDL.

Depends on D142729.

mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp

 //===- TransformInterfaces.cpp - Transform Dialect Interfaces -------------===//
 //
 // 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/Transform/IR/TransformInterfaces.h"
-#include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Operation.h"
-#include "mlir/IR/OwningOpRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 
 #define DEBUG_TYPE "transform-dialect"
 #define DEBUG_PRINT_AFTER_ALL "transform-dialect-print-top-level-after-all"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "] ")
if (!transform->hasTrait<PossibleTopLevelTransformOpTrait>() ||
       transform->getNumOperands() != 0) {
     transform->emitError()
         << "expected transform to start at the top-level transform op";
     llvm::report_fatal_error("could not run transforms",
                              /*gen_crash_diag=*/false);
   }
 #endif // NDEBUG
 
-  // If the transform dialect may use PDL which may modify the IR, clone it
-  // before use to avoid concurrent modification in case this is being called
-  // from pass instances running concurrently with a shared transform script.
-  auto *pdlDialect =
-      transform->getContext()->getLoadedDialect<pdl::PDLDialect>();
-  bool hasPDL = transform
-                    .walk([pdlDialect](Operation *op) {
-                      if (op->getDialect() == pdlDialect)
-                        return WalkResult::interrupt();
-                      return WalkResult::advance();
-                    })
-                    .wasInterrupted();
-
-  OwningOpRef<TransformOpInterface> owningCopy;
-  if (hasPDL) {
-    owningCopy = OwningOpRef<TransformOpInterface>(transform->clone());
-    transform = owningCopy.get();
-  }
-
   TransformState state(transform->getParentRegion(), payloadRoot, extraMapping,
                        options);
   return state.applyTransform(transform).checkAndReport();
 }
 
 //===----------------------------------------------------------------------===//
 // Generated interface implementation.
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.cpp.inc"

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

 //===- TransformDialect.cpp - Transform dialect operations ----------------===//
 //
 // 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/Transform/IR/TransformOps.h"
 #include "mlir/Dialect/PDL/IR/PDLOps.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/Dialect/Transform/IR/TransformUtils.h"
-#include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Rewrite/FrozenRewritePatternSet.h"
 #include "mlir/Rewrite/PatternApplicator.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Debug.h"
 
LogicalResult PatternApplicatorExtension::findAllMatches(
     StringRef patternName, Operation *root,
     SmallVectorImpl<Operation *> &results) {
   auto it = compiledPatterns.find(patternName);
   if (it == compiledPatterns.end()) {
     auto patternOp = patterns.lookup<pdl::PatternOp>(patternName);
     if (!patternOp)
       return failure();
 
+    // Copy the pattern operation into a new module that is compiled and
+    // consumed by the PDL interpreter.
     OwningOpRef<ModuleOp> pdlModuleOp = ModuleOp::create(patternOp.getLoc());
-    patternOp->moveBefore(pdlModuleOp->getBody(),
-                          pdlModuleOp->getBody()->end());
+    auto builder = OpBuilder::atBlockEnd(pdlModuleOp->getBody());
+    builder.clone(*patternOp);
     PDLPatternModule patternModule(std::move(pdlModuleOp));
 
     // Merge in the hooks owned by the dialect. Make a copy as they may be
     // also used by the following operations.
     auto *dialect =
         root->getContext()->getLoadedDialect<transform::TransformDialect>();
     for (const auto &[name, constraintFn] : dialect->getPDLConstraintHooks())
       patternModule.registerConstraintFunction(name, constraintFn);

 
 //===----------------------------------------------------------------------===//
 // WithPDLPatternsOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::WithPDLPatternsOp::apply(transform::TransformResults &results,
                                     transform::TransformState &state) {
-  OwningOpRef<ModuleOp> pdlModuleOp =
-      ModuleOp::create(getOperation()->getLoc());
   TransformOpInterface transformOp = nullptr;
   for (Operation &nested : getBody().front()) {
     if (!isa<pdl::PatternOp>(nested)) {
       transformOp = cast<TransformOpInterface>(nested);
       break;
     }
   }
 

mlir/test/Dialect/Transform/test-interpreter-external-concurrent-source.mlir

+// RUN: mlir-opt %s
+// No need to check anything else than parsing here, this is being used by another test as data.
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !transform.any_op):
+  pdl.pattern @func_return : benefit(1) {
+    %0 = pdl.operation "func.return"
+    pdl.rewrite %0 with "transform.dialect"
+  }
+
+  sequence %arg0 : !transform.any_op failures(propagate) {
+  ^bb1(%arg1: !transform.any_op):
+    %0 = pdl_match @func_return in %arg1 : (!transform.any_op) -> !transform.op<"func.return">
+    test_print_remark_at_operand %0, "matched" : !transform.op<"func.return">
+  }
+}

mlir/test/Dialect/Transform/test-interpreter-external-concurrent.mlir

+// RUN: mlir-opt %s --pass-pipeline="builtin.module(func.func(test-transform-dialect-interpreter{transform-file-name=%p/test-interpreter-external-concurrent-source.mlir}))" \
+// RUN:             --verify-diagnostics
+
+// Exercising the pass on multiple functions of different lengths that may be
+// processed concurrently. This should expose potential races.
+
+func.func @f1() {
+  // expected-remark @below {{matched}}
+  return
+}
+
+func.func @f2() {
+  // expected-remark @below {{matched}}
+  return
+}
+
+func.func @f3() {
+  call @f2() : () -> ()
+  call @f2() : () -> ()
+  call @f5() : () -> ()
+  call @f7() : () -> ()
+  call @f5() : () -> ()
+  call @f5() : () -> ()
+  // expected-remark @below {{matched}}
+  return
+}
+
+func.func @f4() {
+  call @f3() : () -> ()
+  call @f3() : () -> ()
+  // expected-remark @below {{matched}}
+  return
+}
+
+func.func @f5() {
+  call @f7() : () -> ()
+  call @f7() : () -> ()
+  call @f7() : () -> ()
+  call @f7() : () -> ()
+  call @f1() : () -> ()
+  call @f1() : () -> ()
+  call @f7() : () -> ()
+  call @f7() : () -> ()
+  call @f7() : () -> ()
+  call @f7() : () -> ()
+  // expected-remark @below {{matched}}
+  return
+}
+
+func.func @f6() {
+  // expected-remark @below {{matched}}
+  return
+}
+
+func.func @f7() {
+  // expected-remark @below {{matched}}
+  return
+}

mlir/test/lib/Dialect/Transform/TestTransformDialectInterpreter.cpp

 
 using namespace mlir;
 
 namespace {
 /// Simple pass that applies transform dialect ops directly contained in a
 /// module.
 
 template <typename Derived>
-class ModulePassWrapper : public PassWrapper<Derived, OperationPass<ModuleOp>> {
-};
+class OpPassWrapper : public PassWrapper<Derived, OperationPass<>> {};
 
 class TestTransformDialectInterpreterPass
     : public transform::TransformInterpreterPassBase<
-          TestTransformDialectInterpreterPass, ModulePassWrapper> {
+          TestTransformDialectInterpreterPass, OpPassWrapper> {
 public:
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
       TestTransformDialectInterpreterPass)
 
   TestTransformDialectInterpreterPass() = default;
   TestTransformDialectInterpreterPass(
       const TestTransformDialectInterpreterPass &pass)
       : TransformInterpreterPassBase(pass) {}