[mlir][Pass] Remove the restriction that PassManager can only run on ModuleOp

[rriddle]

Summary

This was a somewhat important restriction in the past when ModuleOp was distinctly the top-level container operation, as well as before the pass manager had support for running nested pass managers natively. With these two issues fading away, there isn't really a good reason to enforce that a ModuleOp is the thing running within a pass manager. As such, this revision removes the restriction and allows for users to pass in the name of the operation that the pass manager will be scheduled on.

The only remaining dependency on BuiltinOps from Pass after this revision is due to FunctionPass, which will be resolved in a followup revision.

Depends On D92448

[mehdi_amini]

Sweet!

mlir/include/mlir/Pass/AnalysisManager.h

 //===- AnalysisManager.h - Analysis Management Infrastructure ---*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_PASS_ANALYSISMANAGER_H
 #define MLIR_PASS_ANALYSISMANAGER_H
 
-#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Operation.h"
 #include "mlir/Pass/PassInstrumentation.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/TypeName.h"
 
 namespace mlir {
 //===----------------------------------------------------------------------===//
private:
   template <typename AnalysisT, typename OpT>
   AnalysisT &getAnalysisImpl(PassInstrumentor *pi, OpT op) {
     TypeID id = TypeID::get<AnalysisT>();
 
     typename ConceptMap::iterator it;
     bool wasInserted;
     std::tie(it, wasInserted) = analyses.try_emplace(id);
 
-    // If we don't have a cached analysis for this function, compute it directly
-    // and add it to the cache.
+    // If we don't have a cached analysis for this operation, compute it
+    // directly and add it to the cache.
     if (wasInserted) {
       if (pi)
         pi->runBeforeAnalysis(getAnalysisName<AnalysisT>(), id, ir);
 
       it->second = std::make_unique<AnalysisModel<AnalysisT>>(op);
 
       if (pi)
         pi->runAfterAnalysis(getAnalysisName<AnalysisT>(), id, ir);
private:
 
   /// A reference to the impl analysis map within the parent analysis manager.
   detail::NestedAnalysisMap *impl;
 
   /// Allow access to the constructor.
   friend class ModuleAnalysisManager;
 };
 
-/// An analysis manager class specifically for the top-level module operation.
-/// This class contains the memory allocations for all nested analysis managers,
-/// and provides an anchor point. This is necessary because AnalysisManager is
+/// An analysis manager class specifically for the top-level operation. This
+/// class contains the memory allocations for all nested analysis managers, and
+/// provides an anchor point. This is necessary because AnalysisManager is
 /// designed to be a thin wrapper around an existing analysis map instance.
 class ModuleAnalysisManager {
 public:
-  ModuleAnalysisManager(ModuleOp module, PassInstrumentor *passInstrumentor)
-      : analyses(module), passInstrumentor(passInstrumentor) {}
+  ModuleAnalysisManager(Operation *op, PassInstrumentor *passInstrumentor)
+      : analyses(op), passInstrumentor(passInstrumentor) {}
   ModuleAnalysisManager(const ModuleAnalysisManager &) = delete;
   ModuleAnalysisManager &operator=(const ModuleAnalysisManager &) = delete;
 
   /// Returns a pass instrumentation object for the current module. This value
   /// may be null.
   PassInstrumentor *getPassInstrumentor() const { return passInstrumentor; }
 
   /// Returns an analysis manager for the current top-level module.

mlir/include/mlir/Pass/PassManager.h

 namespace llvm {
 class Any;
 } // end namespace llvm
 
 namespace mlir {
 class AnalysisManager;
 class Identifier;
 class MLIRContext;
-class ModuleOp;
 class Operation;
 class Pass;
 class PassInstrumentation;
 class PassInstrumentor;
 
 namespace detail {
 struct OpPassManagerImpl;
 struct PassExecutionState;
enum class PassDisplayMode {
   // 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:
-  PassManager(MLIRContext *ctx, Nesting nesting = Nesting::Explicit);
+  /// Create a new pass manager under the given context with a specific nesting
+  /// style. The created pass manager can schedule operations that match
+  /// `operationName`.
+  PassManager(MLIRContext *ctx, Nesting nesting = Nesting::Explicit,
+              StringRef operationName = "module");
+  PassManager(MLIRContext *ctx, StringRef operationName)
+      : PassManager(ctx, Nesting::Explicit, operationName) {}
   ~PassManager();
 
-  /// Run the passes within this manager on the provided module.
+  /// 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.
   LLVM_NODISCARD
-  LogicalResult run(ModuleOp module);
+  LogicalResult run(Operation *op);
 
   /// Return an instance of the context.
   MLIRContext *getContext() const { return context; }
 
   /// Enable support for the pass manager to generate a reproducer on the event
   /// of a crash or a pass failure. `outputFile` is a .mlir filename used to
   /// write the generated reproducer. If `genLocalReproducer` is true, the pass
   /// manager will attempt to generate a local reproducer that contains the
public:
   void
   enableStatistics(PassDisplayMode displayMode = PassDisplayMode::Pipeline);
 
 private:
   /// Dump the statistics of the passes within this pass manager.
   void dumpStatistics();
 
   /// Run the pass manager with crash recover enabled.
-  LogicalResult runWithCrashRecovery(ModuleOp module, AnalysisManager am);
+  LogicalResult runWithCrashRecovery(Operation *op, AnalysisManager am);
   /// Run the given passes with crash recover enabled.
   LogicalResult
   runWithCrashRecovery(MutableArrayRef<std::unique_ptr<Pass>> passes,
-                       ModuleOp module, AnalysisManager am);
+                       Operation *op, AnalysisManager am);
 
   /// Context this PassManager was initialized with.
   MLIRContext *context;
 
   /// Flag that specifies if pass statistics should be dumped.
   Optional<PassDisplayMode> passStatisticsMode;
 
   /// A manager for pass instrumentations.

mlir/lib/Pass/IRPrinting.cpp

 //===- IRPrinting.cpp -----------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
-#include "mlir/IR/BuiltinOps.h"
 #include "mlir/Pass/PassManager.h"
 #include "llvm/Support/Format.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/SHA1.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
private:
   /// being operated on in a pass. This field is only used when the
   /// configuration asked for change detection.
   DenseMap<Pass *, OperationFingerPrint> beforePassFingerPrints;
 };
 } // end anonymous namespace
 
 static void printIR(Operation *op, bool printModuleScope, raw_ostream &out,
                     OpPrintingFlags flags) {
-  // Check to see if we are printing the top-level module.
-  auto module = dyn_cast<ModuleOp>(op);
-  if (module && !op->getBlock())
-    return module.print(out << "\n", flags);
-
   // Otherwise, check to see if we are not printing at module scope.
   if (!printModuleScope)
-    return op->print(out << "\n", flags.useLocalScope());
+    return op->print(out << "\n",
+                     op->getBlock() ? flags.useLocalScope() : flags);
 
   // Otherwise, we are printing at module scope.
   out << " ('" << op->getName() << "' operation";
   if (auto symbolName =
           op->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName()))
     out << ": @" << symbolName.getValue();
   out << ")\n";
 
-  // Find the top-level module operation.
+  // Find the top-level operation.
   auto *topLevelOp = op;
   while (auto *parentOp = topLevelOp->getParentOp())
     topLevelOp = parentOp;
-
-  // Check to see if the top-level operation is actually a module in the case of
-  // invalid-ir.
-  if (auto module = dyn_cast<ModuleOp>(topLevelOp))
-    module.print(out, flags);
-  else
-    topLevelOp->print(out, flags);
 }
 
 /// Instrumentation hooks.
 void IRPrinterInstrumentation::runBeforePass(Pass *pass, Operation *op) {
   if (isa<OpToOpPassAdaptor>(pass))
     return;
   // If the config asked to detect changes, record the current fingerprint.
   if (config->shouldPrintAfterOnlyOnChange())

mlir/lib/Pass/Pass.cpp

 //===- Pass.cpp - Pass infrastructure implementation ----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements common pass infrastructure.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Pass/Pass.h"
 #include "PassDetail.h"
-#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Support/FileUtilities.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/CommandLine.h"
void OpToOpPassAdaptor::runOnOperationAsyncImpl(bool verifyPasses) {
   AnalysisManager am = getAnalysisManager();
 
   // Create the async executors if they haven't been created, or if the main
   // pipeline has changed.
   if (asyncExecutors.empty() || hasSizeMismatch(asyncExecutors.front(), mgrs))
     asyncExecutors.assign(llvm::hardware_concurrency().compute_thread_count(),
                           mgrs);
 
-  // Run a prepass over the module to collect the operations to execute over.
-  // This ensures that an analysis manager exists for each operation, as well as
-  // providing a queue of operations to execute over.
+  // Run a prepass over the operation to collect the nested operations to
+  // execute over. This ensures that an analysis manager exists for each
+  // operation, as well as providing a queue of operations to execute over.
   std::vector<std::pair<Operation *, AnalysisManager>> opAMPairs;
   for (auto &region : getOperation()->getRegions()) {
     for (auto &block : region) {
       for (auto &op : block) {
         // Add this operation iff the name matches any of the pass managers.
         if (findPassManagerFor(mgrs, op.getName().getIdentifier(),
                                getContext()))
           opAMPairs.emplace_back(&op, am.nest(&op));

 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class contains all of the context for generating a recovery reproducer.
 /// Each recovery context is registered globally to allow for generating
 /// reproducers when a signal is raised, such as a segfault.
 struct RecoveryReproducerContext {
   RecoveryReproducerContext(MutableArrayRef<std::unique_ptr<Pass>> passes,
-                            ModuleOp module, StringRef filename,
+                            Operation *op, StringRef filename,
                             bool disableThreads, bool verifyPasses);
   ~RecoveryReproducerContext();
 
   /// Generate a reproducer with the current context.
   LogicalResult generate(std::string &error);
 
 private:
   /// This function is invoked in the event of a crash.
   static void crashHandler(void *);
 
   /// Register a signal handler to run in the event of a crash.
   static void registerSignalHandler();
 
   /// The textual description of the currently executing pipeline.
   std::string pipeline;
 
-  /// The MLIR module representing the IR before the crash.
-  OwningModuleRef module;
+  /// The MLIR operation representing the IR before the crash.
+  Operation *preCrashOperation;
 
   /// The filename to use when generating the reproducer.
   StringRef filename;
 
   /// Various pass manager and context flags.
   bool disableThreads;
   bool verifyPasses;
 

 } // end anonymous namespace
 
 llvm::ManagedStatic<llvm::sys::SmartMutex<true>>
     RecoveryReproducerContext::reproducerMutex;
 llvm::ManagedStatic<llvm::SmallSetVector<RecoveryReproducerContext *, 1>>
     RecoveryReproducerContext::reproducerSet;
 
 RecoveryReproducerContext::RecoveryReproducerContext(
-    MutableArrayRef<std::unique_ptr<Pass>> passes, ModuleOp module,
+    MutableArrayRef<std::unique_ptr<Pass>> passes, Operation *op,
     StringRef filename, bool disableThreads, bool verifyPasses)
-    : module(module.clone()), filename(filename),
+    : preCrashOperation(op->clone()), filename(filename),
       disableThreads(disableThreads), verifyPasses(verifyPasses) {
   // Grab the textual pipeline being executed..
   {
     llvm::raw_string_ostream pipelineOS(pipeline);
     ::printAsTextualPipeline(passes, pipelineOS);
   }
 
   // Make sure that the handler is registered, and update the current context.
   llvm::sys::SmartScopedLock<true> producerLock(*reproducerMutex);
   if (reproducerSet->empty())
     llvm::CrashRecoveryContext::Enable();
   registerSignalHandler();
   reproducerSet->insert(this);
 }
 
 RecoveryReproducerContext::~RecoveryReproducerContext() {
+  // Erase the cloned preCrash IR that we cached.
+  preCrashOperation->erase();
+
   llvm::sys::SmartScopedLock<true> producerLock(*reproducerMutex);
   reproducerSet->remove(this);
   if (reproducerSet->empty())
     llvm::CrashRecoveryContext::Disable();
 }
 
 LogicalResult RecoveryReproducerContext::generate(std::string &error) {
   std::unique_ptr<llvm::ToolOutputFile> outputFile =
       mlir::openOutputFile(filename, &error);
   if (!outputFile)
     return failure();
   auto &outputOS = outputFile->os();
 
   // Output the current pass manager configuration.
   outputOS << "// configuration: -pass-pipeline='" << pipeline << "'";
   if (disableThreads)
     outputOS << " -mlir-disable-threading";
 
   // TODO: Should this also be configured with a pass manager flag?
   outputOS << "\n// note: verifyPasses=" << (verifyPasses ? "true" : "false")
            << "\n";
 
   // Output the .mlir module.
-  module->print(outputOS);
+  preCrashOperation->print(outputOS);
   outputFile->keep();
   return success();
 }
 
 void RecoveryReproducerContext::crashHandler(void *) {
   // Walk the current stack of contexts and generate a reproducer for each one.
   // We can't know for certain which one was the cause, so we need to generate
   // a reproducer for all of them.
   std::string ignored;
   for (RecoveryReproducerContext *context : *reproducerSet)
     context->generate(ignored);
 }
 
 void RecoveryReproducerContext::registerSignalHandler() {
   // Ensure that the handler is only registered once.
   static bool registered =
       (llvm::sys::AddSignalHandler(crashHandler, nullptr), false);
   (void)registered;
 }
 
 /// Run the pass manager with crash recover enabled.
-LogicalResult PassManager::runWithCrashRecovery(ModuleOp module,
+LogicalResult PassManager::runWithCrashRecovery(Operation *op,
                                                 AnalysisManager am) {
   // If this isn't a local producer, run all of the passes in recovery mode.
   if (!localReproducer)
-    return runWithCrashRecovery(impl->passes, module, am);
+    return runWithCrashRecovery(impl->passes, op, am);
 
   // Split the passes within adaptors to ensure that each pass can be run in
   // isolation.
   impl->splitAdaptorPasses();
 
   // If this is a local producer, run each of the passes individually.
   MutableArrayRef<std::unique_ptr<Pass>> passes = impl->passes;
   for (std::unique_ptr<Pass> &pass : passes)
-    if (failed(runWithCrashRecovery(pass, module, am)))
+    if (failed(runWithCrashRecovery(pass, op, am)))
       return failure();
   return success();
 }
 
 /// Run the given passes with crash recover enabled.
 LogicalResult
 PassManager::runWithCrashRecovery(MutableArrayRef<std::unique_ptr<Pass>> passes,
-                                  ModuleOp module, AnalysisManager am) {
-  RecoveryReproducerContext context(passes, module, *crashReproducerFileName,
+                                  Operation *op, AnalysisManager am) {
+  RecoveryReproducerContext context(passes, op, *crashReproducerFileName,
                                     !getContext()->isMultithreadingEnabled(),
                                     verifyPasses);
 
   // Safely invoke the passes within a recovery context.
   LogicalResult passManagerResult = failure();
   llvm::CrashRecoveryContext recoveryContext;
   recoveryContext.RunSafelyOnThread([&] {
     for (std::unique_ptr<Pass> &pass : passes)
-      if (failed(OpToOpPassAdaptor::run(pass.get(), module, am, verifyPasses)))
+      if (failed(OpToOpPassAdaptor::run(pass.get(), op, am, verifyPasses)))
         return;
     passManagerResult = success();
   });
   if (succeeded(passManagerResult))
     return success();
 
   std::string error;
   if (failed(context.generate(error)))
-    return module.emitError("<MLIR-PassManager-Crash-Reproducer>: ") << error;
-  return module.emitError()
+    return op->emitError("<MLIR-PassManager-Crash-Reproducer>: ") << error;
+  return op->emitError()
          << "A failure has been detected while processing the MLIR module, a "
             "reproducer has been generated in '"
          << *crashReproducerFileName << "'";
 }
 
 //===----------------------------------------------------------------------===//
 // PassManager
 //===----------------------------------------------------------------------===//
 
-PassManager::PassManager(MLIRContext *ctx, Nesting nesting)
-    : OpPassManager(Identifier::get(ModuleOp::getOperationName(), ctx),
-                    nesting),
-      context(ctx), passTiming(false), localReproducer(false),
-      verifyPasses(true) {}
+PassManager::PassManager(MLIRContext *ctx, Nesting nesting,
+                         StringRef operationName)
+    : OpPassManager(Identifier::get(operationName, ctx), nesting), context(ctx),
+      passTiming(false), localReproducer(false), verifyPasses(true) {}
 
 PassManager::~PassManager() {}
 
 void PassManager::enableVerifier(bool enabled) { verifyPasses = enabled; }
 
-/// Run the passes within this manager on the provided module.
-LogicalResult PassManager::run(ModuleOp module) {
+/// Run the passes within this manager on the provided operation.
+LogicalResult PassManager::run(Operation *op) {
+  MLIRContext *context = getContext();
+  assert(op->getName().getIdentifier() == getOpName(*context) &&
+         "operation has a different name than the PassManager");
+
   // Before running, make sure to coalesce any adjacent pass adaptors in the
   // pipeline.
   getImpl().coalesceAdjacentAdaptorPasses();
 
   // Register all dialects for the current pipeline.
   DialectRegistry dependentDialects;
   getDependentDialects(dependentDialects);
-  dependentDialects.loadAll(module.getContext());
+  dependentDialects.loadAll(context);
 
-  // Construct an analysis manager for the pipeline.
-  ModuleAnalysisManager am(module, instrumentor.get());
+  // Construct a top level analysis manager for the pipeline.
+  ModuleAnalysisManager am(op, instrumentor.get());
 
   // Notify the context that we start running a pipeline for book keeping.
-  module.getContext()->enterMultiThreadedExecution();
+  context->enterMultiThreadedExecution();
 
   // If reproducer generation is enabled, run the pass manager with crash
   // handling enabled.
-  LogicalResult result = crashReproducerFileName
-                             ? runWithCrashRecovery(module, am)
-                             : OpToOpPassAdaptor::runPipeline(
-                                   getPasses(), module, am, verifyPasses);
+  LogicalResult result =
+      crashReproducerFileName
+          ? runWithCrashRecovery(op, am)
+          : OpToOpPassAdaptor::runPipeline(getPasses(), op, am, verifyPasses);
 
   // Notify the context that the run is done.
-  module.getContext()->exitMultiThreadedExecution();
+  context->exitMultiThreadedExecution();
 
   // Dump all of the pass statistics if necessary.
   if (passStatisticsMode)
     dumpStatistics();
   return result;
 }
 
 /// Enable support for the pass manager to generate a reproducer on the event

mlir/lib/Pass/PassManagerOptions.cpp

struct PassManagerOptions {
   llvm::cl::opt<bool> printAfterChange{
       "print-ir-after-change",
       llvm::cl::desc(
           "When printing the IR after a pass, only print if the IR changed"),
       llvm::cl::init(false)};
   llvm::cl::opt<bool> printModuleScope{
       "print-ir-module-scope",
       llvm::cl::desc("When printing IR for print-ir-[before|after]{-all} "
-                     "always print the top-level module operation"),
+                     "always print the top-level operation"),
       llvm::cl::init(false)};
 
   /// Add an IR printing instrumentation if enabled by any 'print-ir' flags.
   void addPrinterInstrumentation(PassManager &pm);
 
   //===--------------------------------------------------------------------===//
   // Pass Timing
   //===--------------------------------------------------------------------===//