[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

212–221

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

212–221

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

212–221

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.

459–460

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::ModulePassManager pm(mlirCtx.get(),
+                             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::ModulePassManager pm(mlirCtx.get(),
+                             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::ModulePassManager pm(&ctx, 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::ModulePassManager pm(&context, 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.
+// Create a top-level `PassManager` class.
+ModulePassManager pm(ctx);
+// Note: `ModulePassManager` is just a utility class, and we could also create
+// the above `PassManager` by specifying the anchor name:
 PassManager pm(ctx, /*operationName=*/"builtin.module");
 
 // 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>();

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 impl->name; }
 
   void print(raw_ostream &os) const;
   void dump() const;

mlir/include/mlir/Pass/PassManager.h

 
 namespace llvm {
 class Any;
 } // namespace llvm
 
 namespace mlir {
 class AnalysisManager;
 class MLIRContext;
+class ModuleOp;
 class Operation;
 class Pass;
 class PassInstrumentation;
 class PassInstrumentor;
 
 namespace detail {
 struct OpPassManagerImpl;
 class OpToOpPassAdaptor;
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();
 
   /// 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;
 };
 
+/// Utility class for constructing pass managers anchored on a specific
+/// operation type.
+template <typename OpType>
+class PassManagerOn : public PassManager {
+public:
+  PassManagerOn(MLIRContext *ctx, Nesting nesting = Nesting::Explicit)
+      : PassManager(ctx, OpType::getOperationName(), nesting) {}
+};
+
+/// A pass manager anchored on `builtin.module`
+using ModulePassManager = PassManagerOn<ModuleOp>;
+
 /// 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);

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

static LogicalResult performActions(raw_ostream &os, bool verifyDiagnostics,
   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

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);
+  ModulePassManager pm(&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