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

[Maksim Levental]

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

>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 1/2] [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

>From ba7568eab77c102fd258303e196df1aad0d2ec52 Mon Sep 17 00:00:00 2001
From: max <maksim.levental at gmail.com>
Date: Sun, 12 Nov 2023 22:40:30 -0600
Subject: [PATCH 2/2] call `registerPass` directly to exercise code.

---
 mlir/unittests/Pass/PassManagerTest.cpp | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/mlir/unittests/Pass/PassManagerTest.cpp b/mlir/unittests/Pass/PassManagerTest.cpp
index ceb7a22d8c1226e..2de11a9877c1567 100644
--- a/mlir/unittests/Pass/PassManagerTest.cpp
+++ b/mlir/unittests/Pass/PassManagerTest.cpp
@@ -186,6 +186,9 @@ struct ReRegisterPass
     : public PassWrapper<ReRegisterPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ReRegisterPass)
 
+  StringRef getArgument() const final {
+    return "test-pass-reregistration-pass";
+  }
   ReRegisterPass(std::string annotation) : annotation(annotation) {}
 
   std::string annotation;
@@ -213,14 +216,25 @@ TEST(PassManagerTest, PassReRegistration) {
   OwningOpRef<ModuleOp> module =
       parseSourceString<ModuleOp>(moduleStr, &context);
 
-  // Instantiate and run pass in first configuration.
+  // Register in first configuration.
+  registerPass([]() -> std::unique_ptr<mlir::Pass> {
+    return std::make_unique<ReRegisterPass>("custom.first");
+  });
+
+  // Run pass to verify first configuration is effective.
   auto pm = PassManager::on<ModuleOp>(&context);
-  pm.addPass(std::make_unique<ReRegisterPass>("custom.first"));
+  (void)parsePassPipeline("test-pass-reregistration-pass", pm, llvm::errs());
   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"));
+  // Register a "reconfiguration" of the pass, i.e., with a different
+  // annotation.
+  registerPass([]() -> std::unique_ptr<mlir::Pass> {
+    return std::make_unique<ReRegisterPass>("custom.second");
+  });
+
+  // Run pass to verify second configuration is effective.
+  (void)parsePassPipeline("test-pass-reregistration-pass", pm, llvm::errs());
   EXPECT_TRUE(succeeded(pm.run(module.get())));
   module->walk(
       [](Operation *op) { EXPECT_TRUE(op->hasAttr("custom.second")); });