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

Mehdi Amini llvmlistbot at llvm.org
Tue Nov 14 19:09:12 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:

Right now the pass registration contract with the user is that they provide a `PassAllocatorFunction` callback : `using PassAllocatorFunction = std::function<std::unique_ptr<Pass>()>;`

>From there the system will create the pass and check that the pass name  isn't registered yet, and if it is that the TypeID matches.

To forbid this we would have to change the API here, maybe expose exclusively a `template<typename PassTy> registerPass();` and nothing else?


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


More information about the Mlir-commits mailing list