[mlir][Pass] Make PassManager default to op-agnostic

rkayaith · rkayaith · commit 94a309284d72 · 2023-01-25T15:38:19.000-05:00
Currently `PassManager` defaults to being anchored on `builtin.module`. Switching the default makes `PassManager` consistent with `OpPassManager` and avoids the implicit dependency on `builtin.module`. Specifying the anchor op type isn't strictly necessary when using explicit nesting (existing pipelines will continue to work), but I've updated most call sites to specify the anchor since it allows for better error-checking during pipeline construction. Reviewed By: rriddle Differential Revision: https://reviews.llvm.org/D137731
diff --git a/flang/lib/Frontend/FrontendActions.cpp b/flang/lib/Frontend/FrontendActions.cpp
@@ -184,7 +184,8 @@ bool CodeGenAction::beginSourceFileAction() {
   lb.lower(parseTree, ci.getInvocation().getSemanticsContext());
 
   // run the default passes.
-  mlir::PassManager pm(mlirCtx.get(), mlir::OpPassManager::Nesting::Implicit);
+  mlir::PassManager pm((*mlirModule)->getName(),
+                       mlir::OpPassManager::Nesting::Implicit);
   pm.enableVerifier(/*verifyPasses=*/true);
   pm.addPass(std::make_unique<Fortran::lower::VerifierPass>());
 
@@ -535,7 +536,8 @@ void CodeGenAction::generateLLVMIR() {
   fir::support::registerLLVMTranslation(*mlirCtx);
 
   // Set-up the MLIR pass manager
-  mlir::PassManager pm(mlirCtx.get(), mlir::OpPassManager::Nesting::Implicit);
+  mlir::PassManager pm((*mlirModule)->getName(),
+                       mlir::OpPassManager::Nesting::Implicit);
 
   pm.addPass(std::make_unique<Fortran::lower::VerifierPass>());
   pm.enableVerifier(/*verifyPasses=*/true);
diff --git a/flang/tools/bbc/bbc.cpp b/flang/tools/bbc/bbc.cpp
@@ -249,7 +249,8 @@ static mlir::LogicalResult convertFortranSourceToMLIR(
            << outputName;
 
   // Otherwise run the default passes.
-  mlir::PassManager pm(&ctx, mlir::OpPassManager::Nesting::Implicit);
+  mlir::PassManager pm(mlirModule->getName(),
+                       mlir::OpPassManager::Nesting::Implicit);
   pm.enableVerifier(/*verifyPasses=*/true);
   mlir::applyPassManagerCLOptions(pm);
   if (passPipeline.hasAnyOccurrences()) {
diff --git a/flang/tools/tco/tco.cpp b/flang/tools/tco/tco.cpp
@@ -103,7 +103,8 @@ compileFIR(const mlir::PassPipelineCLParser &passPipeline) {
   fir::KindMapping kindMap{&context};
   fir::setTargetTriple(*owningRef, targetTriple);
   fir::setKindMapping(*owningRef, kindMap);
-  mlir::PassManager pm(&context, mlir::OpPassManager::Nesting::Implicit);
+  mlir::PassManager pm((*owningRef)->getName(),
+                       mlir::OpPassManager::Nesting::Implicit);
   pm.enableVerifier(/*verifyPasses=*/true);
   mlir::applyPassManagerCLOptions(pm);
   if (emitFir) {
diff --git a/mlir/docs/PassManagement.md b/mlir/docs/PassManagement.md
@@ -399,11 +399,8 @@ Below is an example of constructing a pipeline that operates on the above
 structure:
 
 ```c++
-// Create a top-level `PassManager` class. If an operation type is not
-// explicitly specific, the default is the builtin `module` operation.
-PassManager pm(ctx);
-// Note: We could also create the above `PassManager` this way.
-PassManager pm(ctx, /*operationName=*/"builtin.module");
+// Create a top-level `PassManager` class.
+auto pm = PassManager::on<ModuleOp>(ctx);
 
 // Add a pass on the top-level module operation.
 pm.addPass(std::make_unique<MyModulePass>());
diff --git a/mlir/docs/Tutorials/Toy/Ch-3.md b/mlir/docs/Tutorials/Toy/Ch-3.md
@@ -124,7 +124,7 @@ pipeline. In MLIR, the optimizations are run through a `PassManager` in a
 similar way to LLVM:
 
 ```c++
-  mlir::PassManager pm(module.getContext());
+  mlir::PassManager pm(module->getName());
   pm.addNestedPass<mlir::toy::FuncOp>(mlir::createCanonicalizerPass());
 ```
 
diff --git a/mlir/examples/toy/Ch3/toyc.cpp b/mlir/examples/toy/Ch3/toyc.cpp
@@ -113,7 +113,7 @@ int dumpMLIR() {
     return error;
 
   if (enableOpt) {
-    mlir::PassManager pm(&context);
+    mlir::PassManager pm(module.get()->getName());
     // Apply any generic pass manager command line options and run the pipeline.
     applyPassManagerCLOptions(pm);
 
diff --git a/mlir/examples/toy/Ch4/toyc.cpp b/mlir/examples/toy/Ch4/toyc.cpp
@@ -114,7 +114,7 @@ int dumpMLIR() {
     return error;
 
   if (enableOpt) {
-    mlir::PassManager pm(&context);
+    mlir::PassManager pm(module.get()->getName());
     // Apply any generic pass manager command line options and run the pipeline.
     applyPassManagerCLOptions(pm);
 
diff --git a/mlir/examples/toy/Ch5/toyc.cpp b/mlir/examples/toy/Ch5/toyc.cpp
@@ -117,7 +117,7 @@ int dumpMLIR() {
   if (int error = loadMLIR(sourceMgr, context, module))
     return error;
 
-  mlir::PassManager pm(&context);
+  mlir::PassManager pm(module.get()->getName());
   // Apply any generic pass manager command line options and run the pipeline.
   applyPassManagerCLOptions(pm);
 
diff --git a/mlir/examples/toy/Ch6/toyc.cpp b/mlir/examples/toy/Ch6/toyc.cpp
@@ -132,7 +132,7 @@ int loadAndProcessMLIR(mlir::MLIRContext &context,
   if (int error = loadMLIR(context, module))
     return error;
 
-  mlir::PassManager pm(&context);
+  mlir::PassManager pm(module.get()->getName());
   // Apply any generic pass manager command line options and run the pipeline.
   applyPassManagerCLOptions(pm);
 
diff --git a/mlir/examples/toy/Ch7/toyc.cpp b/mlir/examples/toy/Ch7/toyc.cpp
@@ -132,7 +132,7 @@ int loadAndProcessMLIR(mlir::MLIRContext &context,
   if (int error = loadMLIR(context, module))
     return error;
 
-  mlir::PassManager pm(&context);
+  mlir::PassManager pm(module.get()->getName());
   // Apply any generic pass manager command line options and run the pipeline.
   applyPassManagerCLOptions(pm);
 
diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
@@ -322,6 +322,9 @@ class OperationName {
   /// Return the operation name with dialect name stripped, if it has one.
   StringRef stripDialect() const { return getStringRef().split('.').second; }
 
+  /// Return the context this operation is associated with.
+  MLIRContext *getContext() { return getIdentifier().getContext(); }
+
   /// Return the name of this operation. This always succeeds.
   StringRef getStringRef() const { return getIdentifier(); }
 
diff --git a/mlir/include/mlir/Pass/PassManager.h b/mlir/include/mlir/Pass/PassManager.h
@@ -213,14 +213,20 @@ class PassManager : public OpPassManager {
   /// Create a new pass manager under the given context with a specific nesting
   /// style. The created pass manager can schedule operations that match
   /// `operationName`.
-  /// FIXME: We should make the specification of `builtin.module` explicit here,
-  /// so that we can have top-level op-agnostic pass managers.
-  PassManager(MLIRContext *ctx, Nesting nesting = Nesting::Explicit,
-              StringRef operationName = "builtin.module");
-  PassManager(MLIRContext *ctx, StringRef operationName)
-      : PassManager(ctx, Nesting::Explicit, operationName) {}
+  PassManager(MLIRContext *ctx,
+              StringRef operationName = PassManager::getAnyOpAnchorName(),
+              Nesting nesting = Nesting::Explicit);
+  PassManager(OperationName operationName, Nesting nesting = Nesting::Explicit);
   ~PassManager();
 
+  /// Create a new pass manager under the given context with a specific nesting
+  /// style. The created pass manager can schedule operations that match
+  /// `OperationTy`.
+  template <typename OperationTy>
+  static PassManager on(MLIRContext *ctx, Nesting nesting = Nesting::Explicit) {
+    return PassManager(ctx, OperationTy::getOperationName(), nesting);
+  }
+
   /// Run the passes within this manager on the provided operation. The
   /// specified operation must have the same name as the one provided the pass
   /// manager on construction.
@@ -438,7 +444,8 @@ class PassManager : public OpPassManager {
   std::unique_ptr<detail::PassCrashReproducerGenerator> crashReproGenerator;
 
   /// A hash key used to detect when reinitialization is necessary.
-  llvm::hash_code initializationKey;
+  llvm::hash_code initializationKey =
+      DenseMapInfo<llvm::hash_code>::getTombstoneKey();
 
   /// Flag that specifies if pass timing is enabled.
   bool passTiming : 1;
diff --git a/mlir/lib/Pass/Pass.cpp b/mlir/lib/Pass/Pass.cpp
@@ -769,11 +769,15 @@ void OpToOpPassAdaptor::runOnOperationAsyncImpl(bool verifyPasses) {
 // PassManager
 //===----------------------------------------------------------------------===//
 
-PassManager::PassManager(MLIRContext *ctx, Nesting nesting,
-                         StringRef operationName)
-    : OpPassManager(OperationName(operationName, ctx), nesting), context(ctx),
-      initializationKey(DenseMapInfo<llvm::hash_code>::getTombstoneKey()),
-      passTiming(false), verifyPasses(true) {}
+PassManager::PassManager(MLIRContext *ctx, StringRef operationName,
+                         Nesting nesting)
+    : OpPassManager(operationName, nesting), context(ctx), passTiming(false),
+      verifyPasses(true) {}
+
+PassManager::PassManager(OperationName operationName, Nesting nesting)
+    : OpPassManager(operationName, nesting),
+      context(operationName.getContext()), passTiming(false),
+      verifyPasses(true) {}
 
 PassManager::~PassManager() = default;
 
diff --git a/mlir/lib/Rewrite/FrozenRewritePatternSet.cpp b/mlir/lib/Rewrite/FrozenRewritePatternSet.cpp
@@ -34,7 +34,7 @@ convertPDLToPDLInterp(ModuleOp pdlModule,
   pdlModule.getBody()->walk(simplifyFn);
 
   /// Lower the PDL pattern module to the interpreter dialect.
-  PassManager pdlPipeline(pdlModule.getContext());
+  PassManager pdlPipeline(pdlModule->getName());
 #ifdef NDEBUG
   // We don't want to incur the hit of running the verifier when in release
   // mode.
diff --git a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
@@ -79,8 +79,7 @@ performActions(raw_ostream &os, bool verifyDiagnostics, bool verifyPasses,
   parserTiming.stop();
 
   // Prepare the pass manager, applying command-line and reproducer options.
-  PassManager pm(context, OpPassManager::Nesting::Implicit,
-                 op.get()->getName().getStringRef());
+  PassManager pm(op.get()->getName(), PassManager::Nesting::Implicit);
   pm.enableVerifier(verifyPasses);
   applyPassManagerCLOptions(pm);
   pm.enableTiming(timing);
diff --git a/mlir/unittests/ExecutionEngine/Invoke.cpp b/mlir/unittests/ExecutionEngine/Invoke.cpp
@@ -52,7 +52,7 @@ static struct LLVMInitializer {
 /// Simple conversion pipeline for the purpose of testing sources written in
 /// dialects lowering to LLVM Dialect.
 static LogicalResult lowerToLLVMDialect(ModuleOp module) {
-  PassManager pm(module.getContext());
+  PassManager pm(module->getName());
   pm.addPass(mlir::createMemRefToLLVMConversionPass());
   pm.addNestedPass<func::FuncOp>(mlir::createArithToLLVMConversionPass());
   pm.addPass(mlir::createConvertFuncToLLVMPass());
diff --git a/mlir/unittests/Pass/PassManagerTest.cpp b/mlir/unittests/Pass/PassManagerTest.cpp
@@ -68,7 +68,7 @@ TEST(PassManagerTest, OpSpecificAnalysis) {
   }
 
   // Instantiate and run our pass.
-  PassManager pm(&context);
+  auto pm = PassManager::on<ModuleOp>(&context);
   pm.addNestedPass<func::FuncOp>(std::make_unique<AnnotateFunctionPass>());
   LogicalResult result = pm.run(module.get());
   EXPECT_TRUE(succeeded(result));
@@ -123,7 +123,7 @@ TEST(PassManagerTest, InvalidPass) {
   });
 
   // Instantiate and run our pass.
-  PassManager pm(&context);
+  auto pm = PassManager::on<ModuleOp>(&context);
   pm.nest("invalid_op").addPass(std::make_unique<InvalidPass>());
   LogicalResult result = pm.run(module.get());
   EXPECT_TRUE(failed(result));
@@ -138,7 +138,10 @@ TEST(PassManagerTest, InvalidPass) {
   EXPECT_TRUE(succeeded(result));
 
   // Check that adding the pass at the top-level triggers a fatal error.
-  ASSERT_DEATH(pm.addPass(std::make_unique<InvalidPass>()), "");
+  ASSERT_DEATH(pm.addPass(std::make_unique<InvalidPass>()),
+               "Can't add pass 'Invalid Pass' restricted to 'invalid_op' on a "
+               "PassManager intended to run on 'builtin.module', did you "
+               "intend to nest?");
 }
 
 } // namespace
