[Mlir-commits] [mlir] [mlir][pass] Allow reregistration of pass with same typeids (PR #72067)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Maksim Levental (makslevental)

<details>
<summary>Changes</summary>

I would like to be able to "dynamically" reregister the same pass (i.e., same `TypeID`) but with different "configuration" (e.g., different constructor args).

---
Full diff: https://github.com/llvm/llvm-project/pull/72067.diff


3 Files Affected:

- (modified) mlir/lib/Pass/PassRegistry.cpp (+1-1) 
- (modified) mlir/unittests/Pass/CMakeLists.txt (+1) 
- (modified) mlir/unittests/Pass/PassManagerTest.cpp (+45) 


``````````diff
diff --git a/mlir/lib/Pass/PassRegistry.cpp b/mlir/lib/Pass/PassRegistry.cpp
index b0c314369190a40..f05dbc22cbaf103 100644
--- a/mlir/lib/Pass/PassRegistry.cpp
+++ b/mlir/lib/Pass/PassRegistry.cpp
@@ -129,7 +129,7 @@ void mlir::registerPass(const PassAllocatorFunction &function) {
                              "' pass that does not override `getArgument()`");
   StringRef description = pass->getDescription();
   PassInfo passInfo(arg, description, function);
-  passRegistry->try_emplace(arg, passInfo);
+  passRegistry->insert_or_assign(arg, passInfo);
 
   // Verify that the registered pass has the same ID as any registered to this
   // arg before it.
diff --git a/mlir/unittests/Pass/CMakeLists.txt b/mlir/unittests/Pass/CMakeLists.txt
index 65f07741238659d..b711fd8f72713a2 100644
--- a/mlir/unittests/Pass/CMakeLists.txt
+++ b/mlir/unittests/Pass/CMakeLists.txt
@@ -6,4 +6,5 @@ add_mlir_unittest(MLIRPassTests
 target_link_libraries(MLIRPassTests
   PRIVATE
   MLIRFuncDialect
+  MLIRParser
   MLIRPass)
diff --git a/mlir/unittests/Pass/PassManagerTest.cpp b/mlir/unittests/Pass/PassManagerTest.cpp
index 9a30f64eaabc293..ceb7a22d8c1226e 100644
--- a/mlir/unittests/Pass/PassManagerTest.cpp
+++ b/mlir/unittests/Pass/PassManagerTest.cpp
@@ -11,6 +11,7 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Diagnostics.h"
+#include "mlir/Parser/Parser.h"
 #include "mlir/Pass/Pass.h"
 #include "gtest/gtest.h"
 
@@ -181,4 +182,48 @@ TEST(PassManagerTest, PassInitialization) {
   EXPECT_TRUE(succeeded(pm.run(module.get())));
 }
 
+struct ReRegisterPass
+    : public PassWrapper<ReRegisterPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ReRegisterPass)
+
+  ReRegisterPass(std::string annotation) : annotation(annotation) {}
+
+  std::string annotation;
+
+  void runOnOperation() override {
+    ModuleOp op = this->getOperation();
+    Builder builder(op);
+    op->walk([this, &builder](Operation *op) {
+      op->setAttr(annotation, builder.getUnitAttr());
+    });
+  }
+};
+
+TEST(PassManagerTest, PassReRegistration) {
+  MLIRContext context;
+  context.allowUnregisteredDialects();
+
+  std::string moduleStr = R"mlir(
+    module {
+      "custom.op1"() : () -> ()
+      "custom.op2"() : () -> ()
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module =
+      parseSourceString<ModuleOp>(moduleStr, &context);
+
+  // Instantiate and run pass in first configuration.
+  auto pm = PassManager::on<ModuleOp>(&context);
+  pm.addPass(std::make_unique<ReRegisterPass>("custom.first"));
+  EXPECT_TRUE(succeeded(pm.run(module.get())));
+  module->walk([](Operation *op) { EXPECT_TRUE(op->hasAttr("custom.first")); });
+
+  // Adding a "reconfiguration" of the pass, i.e., with a different annotation.
+  pm.addPass(std::make_unique<ReRegisterPass>("custom.second"));
+  EXPECT_TRUE(succeeded(pm.run(module.get())));
+  module->walk(
+      [](Operation *op) { EXPECT_TRUE(op->hasAttr("custom.second")); });
+}
+
 } // namespace

``````````

</details>


https://github.com/llvm/llvm-project/pull/72067