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

[Maksim Levental]

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

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).

>From 2b5612fc494313666551ce3447e95d8e5997a500 Mon Sep 17 00:00:00 2001
From: max <maksim.levental at gmail.com>
Date: Sun, 12 Nov 2023 14:32:17 -0600
Subject: [PATCH] [mlir][pass] Allow reregistration of pass with same typeids

---
 mlir/lib/Pass/PassRegistry.cpp          |  2 +-
 mlir/unittests/Pass/CMakeLists.txt      |  1 +
 mlir/unittests/Pass/PassManagerTest.cpp | 45 +++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)

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