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

Mehdi Amini llvmlistbot at llvm.org
Sun Nov 12 20:13:41 PST 2023


================
@@ -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"));
----------------
joker-eph wrote:

> You can't but mlir-opt isn't the only way that passes are run right?

I actually use it as a litmus test for "is this a correct setup".

> No but without the patch the second addPass -> registerPass is a no-op because of try_emplace which doesn't update the value (the pass info) associated with the key (the pass arg) if the key has already been stored in passRegistry.

I understand that, but I have yet to understand why it is a problem?
If you write a test, I would expect it in general to fail without the other changes, otherwise it's not really testing the change.

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


More information about the Mlir-commits mailing list