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

[rkayaith]

Summary

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.

[rkayaith]

fix for doc change

[rriddle]

mlir/include/mlir/Pass/PassManager.h

211–220

Can you add a utility template class that takes the operation type? It'd be nice to preserve the simplicity when the anchor type is known (common for real pipelines), e.g.:

template <typename OpType>
class PassManagerOn : public PassManager ...

That would remove the need for most of the "dialect.opname"/"Op::getOperationName()"/"getName" calls when building a known passmanager.

[rkayaith]

add PassManagerOn class

[rkayaith]

mlir/include/mlir/Pass/PassManager.h

211–220

I added this + ModulePassManager as an alias, though I think there's a potential footgun since it's possible to change the anchor after the pass manager is created, e.g. https://github.com/llvm/llvm-project/blob/ecab1bc0dcdc04ec863f7aa3eaa5daad8232ba65/mlir/lib/Pass/PassCrashRecovery.cpp#L477-L480

[rriddle]

mlir/include/mlir/Pass/PassManager.h

211–220

Good point there. We could have a run method in PassManagerOn that verifies that the anchor is actually correct. We could also just have a templated construction method that returns a pass manager? Just making sure we don't lose the niceties of the API when we don't want "any" anchor.

466–467

I'd just drop any special usings for ModuleOp, it helps keep the API free from builtin deps.

[rkayaith]

I plan to rebase this + switch to a constructor method (static PassManager::on<T>(MLIRContext*) -> PassManager)

[rkayaith]

Switch from PassManagerOn<T> to constructor method

flang/lib/Frontend/FrontendActions.cpp

bool CodeGenAction::beginSourceFileAction() {
   const llvm::DataLayout &dl = tm->createDataLayout();
   setMLIRDataLayout(*mlirModule, dl);
 
   // Create a parse tree and lower it to FIR
   Fortran::parser::Program &parseTree{*ci.getParsing().parseTree()};
   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>());
 
   if (mlir::failed(pm.run(*mlirModule))) {
     unsigned diagID = ci.getDiagnostics().getCustomDiagID(
         clang::DiagnosticsEngine::Error,
         "verification of lowering to FIR failed");
     ci.getDiagnostics().Report(diagID);
void CodeGenAction::generateLLVMIR() {
   CompilerInstance &ci = this->getInstance();
   auto opts = ci.getInvocation().getCodeGenOpts();
   llvm::OptimizationLevel level = mapToLevel(opts);
 
   fir::support::loadDialects(*mlirCtx);
   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);
 
   // Create the pass pipeline
   fir::createMLIRToLLVMPassPipeline(pm, level);
   mlir::applyPassManagerCLOptions(pm);
 

flang/tools/bbc/bbc.cpp

if (!outputName.size())
     outputName = llvm::sys::path::stem(inputFilename).str().append(".mlir");
   llvm::raw_fd_ostream out(outputName, ec);
   if (ec)
     return mlir::emitError(mlir::UnknownLoc::get(&ctx),
                            "could not open output file ")
            << 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()) {
     // run the command-line specified pipeline
     (void)passPipeline.addToPipeline(pm, [&](const llvm::Twine &msg) {
       mlir::emitError(mlir::UnknownLoc::get(&ctx)) << msg;
       return mlir::failure();
     });

flang/tools/tco/tco.cpp

compileFIR(const mlir::PassPipelineCLParser &passPipeline) {
 
   std::error_code ec;
   ToolOutputFile out(outputFilename, ec, sys::fs::OF_None);
 
   // run passes
   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) {
     // parse the input and pretty-print it back out
     // -emit-fir intentionally disables all the passes
   } else if (passPipeline.hasAnyOccurrences()) {
     auto errorHandler = [&](const Twine &msg) {
       mlir::emitError(mlir::UnknownLoc::get(pm.getContext())) << msg;

mlir/docs/PassManagement.md

`builtin.module`
   `spirv.module`
     `spirv.func`
 ```
 
 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>());
 
 // Nest a pass manager that operates on `spirv.module` operations nested
 // directly under the top-level module.
 OpPassManager &nestedModulePM = pm.nest<spirv::ModuleOp>();
 nestedModulePM.addPass(std::make_unique<MySPIRVModulePass>());

mlir/docs/Tutorials/Toy/Ch-3.md

 }
 ```
 
 We also need to update our main file, `toyc.cpp`, to add an optimization
 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());
 ```
 
 Finally, we can run `toyc-ch3 test/Examples/Toy/Ch3/transpose_transpose.toy 
 -emit=mlir -opt` and observe our pattern in action:
 
 ```mlir
 toy.func @transpose_transpose(%arg0: tensor<*xf64>) -> tensor<*xf64> {

mlir/examples/toy/Ch3/toyc.cpp

int dumpMLIR() {
 
   mlir::OwningOpRef<mlir::ModuleOp> module;
   llvm::SourceMgr sourceMgr;
   mlir::SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
   if (int error = loadMLIR(sourceMgr, context, module))
     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);
 
     // Add a run of the canonicalizer to optimize the mlir module.
     pm.addNestedPass<mlir::toy::FuncOp>(mlir::createCanonicalizerPass());
     if (mlir::failed(pm.run(*module)))
       return 4;
   }

mlir/examples/toy/Ch4/toyc.cpp

int dumpMLIR() {
 
   mlir::OwningOpRef<mlir::ModuleOp> module;
   llvm::SourceMgr sourceMgr;
   mlir::SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
   if (int error = loadMLIR(sourceMgr, context, module))
     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);
 
     // Inline all functions into main and then delete them.
     pm.addPass(mlir::createInlinerPass());
 
     // Now that there is only one function, we can infer the shapes of each of
     // the operations.

mlir/examples/toy/Ch5/toyc.cpp

int dumpMLIR() {
   context.getOrLoadDialect<mlir::toy::ToyDialect>();
 
   mlir::OwningOpRef<mlir::ModuleOp> module;
   llvm::SourceMgr sourceMgr;
   mlir::SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
   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);
 
   // Check to see what granularity of MLIR we are compiling to.
   bool isLoweringToAffine = emitAction >= Action::DumpMLIRAffine;
 
   if (enableOpt || isLoweringToAffine) {
     // Inline all functions into main and then delete them.

mlir/examples/toy/Ch6/toyc.cpp

int loadMLIR(mlir::MLIRContext &context,
   return 0;
 }
 
 int loadAndProcessMLIR(mlir::MLIRContext &context,
                        mlir::OwningOpRef<mlir::ModuleOp> &module) {
   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);
 
   // Check to see what granularity of MLIR we are compiling to.
   bool isLoweringToAffine = emitAction >= Action::DumpMLIRAffine;
   bool isLoweringToLLVM = emitAction >= Action::DumpMLIRLLVM;
 
   if (enableOpt || isLoweringToAffine) {

mlir/examples/toy/Ch7/toyc.cpp

int loadMLIR(mlir::MLIRContext &context,
   return 0;
 }
 
 int loadAndProcessMLIR(mlir::MLIRContext &context,
                        mlir::OwningOpRef<mlir::ModuleOp> &module) {
   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);
 
   // Check to see what granularity of MLIR we are compiling to.
   bool isLoweringToAffine = emitAction >= Action::DumpMLIRAffine;
   bool isLoweringToLLVM = emitAction >= Action::DumpMLIRLLVM;
 
   if (enableOpt || isLoweringToAffine) {

mlir/include/mlir/IR/OperationSupport.h

public:
   }
 
   /// Return the name of the dialect this operation is registered to.
   StringRef getDialectNamespace() const;
 
   /// 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(); }
 
   /// Return the name of this operation as a StringAttr.
   StringAttr getIdentifier() const { return getImpl()->getName(); }
 
   void print(raw_ostream &os) const;
   void dump() const;

mlir/include/mlir/Pass/PassManager.h

enum class PassDisplayMode {
 
   // In this mode the results are displayed in a nested pipeline view that
   // mirrors the internal pass pipeline that is being executed in the pass
   // manager.
   Pipeline,
 };
 
 /// The main pass manager and pipeline builder.
 class PassManager : public OpPassManager {
 public:
   /// 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.
   LogicalResult run(Operation *op);
 
   /// Return an instance of the context.
   MLIRContext *getContext() const { return context; }
 
private:
   /// A manager for pass instrumentations.
   std::unique_ptr<PassInstrumentor> instrumentor;
 
   /// An optional crash reproducer generator, if this pass manager is setup to
   /// generate reproducers.
   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;
 
   /// A flag that indicates if the IR should be verified in between passes.
   bool verifyPasses : 1;
 };
 
 /// Register a set of useful command-line options that can be used to configure
 /// a pass manager. The values of these options can be applied via the
 /// 'applyPassManagerCLOptions' method below.
 void registerPassManagerCLOptions();
 
 /// Apply any values provided to the pass manager options that were registered
 /// with 'registerPassManagerOptions'.
 void applyPassManagerCLOptions(PassManager &pm);
 
 /// Apply any values provided to the timing manager options that were registered
 /// with `registerDefaultTimingManagerOptions`. This is a handy helper function
 /// if you do not want to bother creating your own timing manager and passing it
 /// to the pass manager.
 void applyDefaultTimingPassManagerCLOptions(PassManager &pm);
 
 } // namespace mlir
 
 #endif // MLIR_PASS_PASSMANAGER_H

mlir/lib/Pass/Pass.cpp

void OpToOpPassAdaptor::runOnOperationAsyncImpl(bool verifyPasses) {
   if (failed(failableParallelForEach(context, opInfos, processFn)))
     signalPassFailure();
 }
 
 //===----------------------------------------------------------------------===//
 // 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;
 
 void PassManager::enableVerifier(bool enabled) { verifyPasses = enabled; }
 
 /// Run the passes within this manager on the provided operation.
 LogicalResult PassManager::run(Operation *op) {
   MLIRContext *context = getContext();

mlir/lib/Rewrite/FrozenRewritePatternSet.cpp

convertPDLToPDLInterp(ModuleOp pdlModule,
   auto simplifyFn = [](Operation *op) {
     // TODO: Add folding here if ever necessary.
     if (isOpTriviallyDead(op))
       op->erase();
   };
   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.
   pdlPipeline.enableVerifier(false);
 #endif
   pdlPipeline.addPass(createPDLToPDLInterpPass(configMap));
   if (failed(pdlPipeline.run(pdlModule)))
     return failure();

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

performActions(raw_ostream &os, bool verifyDiagnostics, bool verifyPasses,
   OwningOpRef<Operation *> op =
       parseSourceFileForTool(sourceMgr, config, implicitModule);
   context->enableMultithreading(wasThreadingEnabled);
   if (!op)
     return failure();
   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);
   if (failed(reproOptions.apply(pm)) || failed(passManagerSetupFn(pm)))
     return failure();
 
   // Run the pipeline.
   if (failed(pm.run(*op)))

mlir/unittests/ExecutionEngine/Invoke.cpp

LLVMInitializer() {
     llvm::InitializeNativeTarget();
     llvm::InitializeNativeTargetAsmPrinter();
   }
 } initializer;
 
 /// 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());
   pm.addPass(mlir::createReconcileUnrealizedCastsPass());
   return pm.run(module);
 }
 
 TEST(MLIRExecutionEngine, SKIP_WITHOUT_JIT(AddInteger)) {

mlir/unittests/Pass/PassManagerTest.cpp

for (StringRef name : {"secret", "not_secret"}) {
     auto func = func::FuncOp::create(
         builder.getUnknownLoc(), name,
         builder.getFunctionType(std::nullopt, std::nullopt));
     func.setPrivate();
     module->push_back(func);
   }
 
   // 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));
 
   // Verify that each function got annotated with expected attributes.
   for (func::FuncOp func : module->getOps<func::FuncOp>()) {
     ASSERT_TRUE(func->getAttr("isFunc").isa<BoolAttr>());
     EXPECT_TRUE(func->getAttr("isFunc").cast<BoolAttr>().getValue());
TEST(PassManagerTest, InvalidPass) {
   // Register a diagnostic handler to capture the diagnostic so that we can
   // check it later.
   std::unique_ptr<Diagnostic> diagnostic;
   context.getDiagEngine().registerHandler([&](Diagnostic &diag) {
     diagnostic = std::make_unique<Diagnostic>(std::move(diag));
   });
 
   // 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));
   ASSERT_TRUE(diagnostic.get() != nullptr);
   EXPECT_EQ(
       diagnostic->str(),
       "'invalid_op' op trying to schedule a pass on an unregistered operation");
 
   // Check that clearing the pass manager effectively removed the pass.
   pm.clear();
   result = pm.run(module.get());
   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