Decouple OpPassManager from the the MLIRContext (NFC)

[mehdi_amini]

Summary

This is allowing to build an OpPassManager from a StringRef instead of an
Identifier, which enables building pipelines without an MLIRContext.
An identifier is still cached on-demand on the OpPassManager for efficiency
during the IR traversal.

[rriddle]

mlir/include/mlir/Pass/PassManager.h

311

nit: Add a comment.

mlir/lib/Pass/Pass.cpp

405

When is this one going to be used?

[mehdi_amini]

Fix build

[mehdi_amini]

Landed in c0b6bc070e78cbd20bc4351704f52d85192e8804

mlir/lib/Pass/Pass.cpp

405

It is used when building the pipeline: merging and nesting new OpPassManager.
The other one is used in the "run" path.

mlir/include/mlir/Pass/PassManager.h

 // OpPassManager
 //===----------------------------------------------------------------------===//
 
 /// This class represents a pass manager that runs passes on a specific
 /// operation type. This class is not constructed directly, but nested within
 /// other OpPassManagers or the top-level PassManager.
 class OpPassManager {
 public:
-  OpPassManager(Identifier name, MLIRContext *context, bool verifyPasses);
+  OpPassManager(Identifier name, bool verifyPasses);
+  OpPassManager(StringRef name, bool verifyPasses);
   OpPassManager(OpPassManager &&rhs);
   OpPassManager(const OpPassManager &rhs);
   ~OpPassManager();
   OpPassManager &operator=(const OpPassManager &rhs);
 
   /// Iterator over the passes in this pass manager.
   using pass_iterator =
       llvm::pointee_iterator<MutableArrayRef<std::unique_ptr<Pass>>::iterator>;
iterator_range<const_pass_iterator> getPasses() const {
     return {begin(), end()};
   }
 
   /// Nest a new operation pass manager for the given operation kind under this
   /// pass manager.
   OpPassManager &nest(Identifier nestedName);
   OpPassManager &nest(StringRef nestedName);
   template <typename OpT> OpPassManager &nest() {
-    return nest(Identifier::get(OpT::getOperationName(), getContext()));
+    return nest(OpT::getOperationName());
   }
 
   /// Add the given pass to this pass manager. If this pass has a concrete
   /// operation type, it must be the same type as this pass manager.
   void addPass(std::unique_ptr<Pass> pass);
 
   /// Add the given pass to a nested pass manager for the given operation kind
   /// `OpT`.
   template <typename OpT> void addNestedPass(std::unique_ptr<Pass> pass) {
     nest<OpT>().addPass(std::move(pass));
   }
 
   /// Returns the number of passes held by this manager.
   size_t size() const;
 
-  /// Return an instance of the context.
-  MLIRContext *getContext() const;
+  /// Return the operation name that this pass manager operates on.
+  Identifier getOpName(MLIRContext &context) const;
 
   /// Return the operation name that this pass manager operates on.
-  Identifier getOpName() const;
+  StringRef getOpName() const;
 
   /// Returns the internal implementation instance.
   detail::OpPassManagerImpl &getImpl();
 
   /// Prints out the passes of the pass manager as the textual representation
   /// of pipelines.
   /// Note: The quality of the string representation depends entirely on the
   /// the correctness of per-pass overrides of Pass::printAsTextualPipeline.
public:
   // If verifyPasses is true, the verifier is run after each pass.
   PassManager(MLIRContext *ctx, bool verifyPasses = true);
   ~PassManager();
 
   /// Run the passes within this manager on the provided module.
   LLVM_NODISCARD
   LogicalResult run(ModuleOp module);
 
+  /// 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
   /// smallest pipeline.
   void enableCrashReproducerGeneration(StringRef outputFile,
                                        bool genLocalReproducer = false);
 
private:
 
   /// Run the pass manager with crash recover enabled.
   LogicalResult runWithCrashRecovery(ModuleOp module, AnalysisManager am);
   /// Run the given passes with crash recover enabled.
   LogicalResult
   runWithCrashRecovery(MutableArrayRef<std::unique_ptr<Pass>> passes,
                        ModuleOp module, AnalysisManager am);
 
+  MLIRContext *context;
+
   /// Flag that specifies if pass statistics should be dumped.
   Optional<PassDisplayMode> passStatisticsMode;
 
   /// A manager for pass instrumentations.
   std::unique_ptr<PassInstrumentor> instrumentor;
 
   /// An optional filename to use when generating a crash reproducer if valid.
   Optional<std::string> crashReproducerFileName;

mlir/lib/Pass/Pass.cpp

 
 //===----------------------------------------------------------------------===//
 // OpPassManagerImpl
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace detail {
 struct OpPassManagerImpl {
-  OpPassManagerImpl(Identifier name, MLIRContext *ctx, bool verifyPasses)
-      : name(name), context(ctx), verifyPasses(verifyPasses) {}
+  OpPassManagerImpl(Identifier identifier, bool verifyPasses)
+      : name(identifier), identifier(identifier), verifyPasses(verifyPasses) {}
+  OpPassManagerImpl(StringRef name, bool verifyPasses)
+      : name(name), verifyPasses(verifyPasses) {}
 
   /// Merge the passes of this pass manager into the one provided.
   void mergeInto(OpPassManagerImpl &rhs);
 
   /// Nest a new operation pass manager for the given operation kind under this
   /// pass manager.
   OpPassManager &nest(Identifier nestedName);
-  OpPassManager &nest(StringRef nestedName) {
-    return nest(Identifier::get(nestedName, getContext()));
-  }
+  OpPassManager &nest(StringRef nestedName);
 
   /// Add the given pass to this pass manager. If this pass has a concrete
   /// operation type, it must be the same type as this pass manager.
   void addPass(std::unique_ptr<Pass> pass);
 
   /// Coalesce adjacent AdaptorPasses into one large adaptor. This runs
   /// recursively through the pipeline graph.
   void coalesceAdjacentAdaptorPasses();
 
   /// Split all of AdaptorPasses such that each adaptor only contains one leaf
   /// pass.
   void splitAdaptorPasses();
 
-  /// Return an instance of the context.
-  MLIRContext *getContext() const { return context; }
+  Identifier getOpName(MLIRContext &context) {
+    if (!identifier)
+      identifier = Identifier::get(name, &context);
+    return *identifier;
+  }
 
   /// The name of the operation that passes of this pass manager operate on.
-  Identifier name;
+  StringRef name;
 
-  /// The current context for this pass manager
-  MLIRContext *context;
+  /// The cached identifier (internalized in the context) for the name of the
+  /// operation that passes of this pass manager operate on.
+  Optional<Identifier> identifier;
 
   /// Flag that specifies if the IR should be verified after each pass has run.
   bool verifyPasses : 1;
 
   /// The set of passes to run as part of this pass manager.
   std::vector<std::unique_ptr<Pass>> passes;
 };
 } // end namespace detail
 } // end namespace mlir
 
 void OpPassManagerImpl::mergeInto(OpPassManagerImpl &rhs) {
   assert(name == rhs.name && "merging unrelated pass managers");
   for (auto &pass : passes)
     rhs.passes.push_back(std::move(pass));
   passes.clear();
 }
 
 OpPassManager &OpPassManagerImpl::nest(Identifier nestedName) {
-  OpPassManager nested(nestedName, getContext(), verifyPasses);
+  OpPassManager nested(nestedName, verifyPasses);
+  auto *adaptor = new OpToOpPassAdaptor(std::move(nested));
+  addPass(std::unique_ptr<Pass>(adaptor));
+  return adaptor->getPassManagers().front();
+}
+
+OpPassManager &OpPassManagerImpl::nest(StringRef nestedName) {
+  OpPassManager nested(nestedName, verifyPasses);
   auto *adaptor = new OpToOpPassAdaptor(std::move(nested));
   addPass(std::unique_ptr<Pass>(adaptor));
   return adaptor->getPassManagers().front();
 }
 
 void OpPassManagerImpl::addPass(std::unique_ptr<Pass> pass) {
   // If this pass runs on a different operation than this pass manager, then
   // implicitly nest a pass manager for this operation.
   auto passOpName = pass->getOpName();
-  if (passOpName && passOpName != name.strref())
+  if (passOpName && passOpName != name)
     return nest(*passOpName).addPass(std::move(pass));
 
   passes.emplace_back(std::move(pass));
   if (verifyPasses)
     passes.emplace_back(std::make_unique<VerifierPass>());
 }
 
 void OpPassManagerImpl::coalesceAdjacentAdaptorPasses() {
for (std::unique_ptr<Pass> &pass : oldPasses) {
     }
   }
 }
 
 //===----------------------------------------------------------------------===//
 // OpPassManager
 //===----------------------------------------------------------------------===//
 
-OpPassManager::OpPassManager(Identifier name, MLIRContext *context,
-                             bool verifyPasses)
-    : impl(new OpPassManagerImpl(name, context, verifyPasses)) {}
+OpPassManager::OpPassManager(Identifier name, bool verifyPasses)
+    : impl(new OpPassManagerImpl(name, verifyPasses)) {}
+OpPassManager::OpPassManager(StringRef name, bool verifyPasses)
+    : impl(new OpPassManagerImpl(name, verifyPasses)) {}
 OpPassManager::OpPassManager(OpPassManager &&rhs) : impl(std::move(rhs.impl)) {}
 OpPassManager::OpPassManager(const OpPassManager &rhs) { *this = rhs; }
 OpPassManager &OpPassManager::operator=(const OpPassManager &rhs) {
-  impl.reset(new OpPassManagerImpl(rhs.impl->name, rhs.impl->getContext(),
-                                   rhs.impl->verifyPasses));
+  impl.reset(new OpPassManagerImpl(rhs.impl->name, rhs.impl->verifyPasses));
   for (auto &pass : rhs.impl->passes)
     impl->passes.emplace_back(pass->clone());
   return *this;
 }
 
 OpPassManager::~OpPassManager() {}
 
 OpPassManager::pass_iterator OpPassManager::begin() {

 }
 
 /// Returns the number of passes held by this manager.
 size_t OpPassManager::size() const { return impl->passes.size(); }
 
 /// Returns the internal implementation instance.
 OpPassManagerImpl &OpPassManager::getImpl() { return *impl; }
 
-/// Return an instance of the context.
-MLIRContext *OpPassManager::getContext() const { return impl->getContext(); }
+/// Return the operation name that this pass manager operates on.
+StringRef OpPassManager::getOpName() const { return impl->name; }
 
 /// Return the operation name that this pass manager operates on.
-Identifier OpPassManager::getOpName() const { return impl->name; }
+Identifier OpPassManager::getOpName(MLIRContext &context) const {
+  return impl->getOpName(context);
+}
 
 /// Prints out the given passes as the textual representation of a pipeline.
 static void printAsTextualPipeline(ArrayRef<std::unique_ptr<Pass>> passes,
                                    raw_ostream &os) {
   // Filter out passes that are not part of the public pipeline.
   auto filteredPasses =
       llvm::make_filter_range(passes, [](const std::unique_ptr<Pass> &pass) {
         return !isa<VerifierPass>(pass);
if (failed(run(&pass, op, am)))
       return failure();
 
   return success();
 }
 
 /// Find an operation pass manager that can operate on an operation of the given
 /// type, or nullptr if one does not exist.
 static OpPassManager *findPassManagerFor(MutableArrayRef<OpPassManager> mgrs,
-                                         Identifier name) {
+                                         StringRef name) {
   auto it = llvm::find_if(
       mgrs, [&](OpPassManager &mgr) { return mgr.getOpName() == name; });
   return it == mgrs.end() ? nullptr : &*it;
 }
 
+/// Find an operation pass manager that can operate on an operation of the given
+/// type, or nullptr if one does not exist.
+static OpPassManager *findPassManagerFor(MutableArrayRef<OpPassManager> mgrs,
+                                         Identifier name,
+                                         MLIRContext &context) {
+  auto it = llvm::find_if(
+      mgrs, [&](OpPassManager &mgr) { return mgr.getOpName(context) == name; });
+  return it == mgrs.end() ? nullptr : &*it;
+}
+
 OpToOpPassAdaptor::OpToOpPassAdaptor(OpPassManager &&mgr) {
   mgrs.emplace_back(std::move(mgr));
 }
 
 void OpToOpPassAdaptor::getDependentDialects(DialectRegistry &dialects) const {
   for (auto &pm : mgrs)
     pm.getDependentDialects(dialects);
 }
if (auto *existingPM = findPassManagerFor(rhs.mgrs, pm.getOpName())) {
       rhs.mgrs.emplace_back(std::move(pm));
     }
   }
   mgrs.clear();
 
   // After coalescing, sort the pass managers within rhs by name.
   llvm::array_pod_sort(rhs.mgrs.begin(), rhs.mgrs.end(),
                        [](const OpPassManager *lhs, const OpPassManager *rhs) {
-                         return lhs->getOpName().strref().compare(
-                             rhs->getOpName().strref());
+                         return lhs->getOpName().compare(rhs->getOpName());
                        });
 }
 
 /// Returns the adaptor pass name.
 std::string OpToOpPassAdaptor::getAdaptorName() {
   std::string name = "Pipeline Collection : [";
   llvm::raw_string_ostream os(name);
   llvm::interleaveComma(getPassManagers(), os, [&](OpPassManager &pm) {
PassInstrumentation::PipelineParentInfo parentInfo = {llvm::get_threadid(),
                                                         this};
   auto *instrumentor = am.getPassInstrumentor();
   for (auto &region : getOperation()->getRegions()) {
     for (auto &block : region) {
       for (auto &op : block) {
         auto *mgr = findPassManagerFor(mgrs, op.getName().getIdentifier());
         if (!mgr)
           continue;
+        Identifier opName = mgr->getOpName(*getOperation()->getContext());
 
         // Run the held pipeline over the current operation.
         if (instrumentor)
-          instrumentor->runBeforePipeline(mgr->getOpName(), parentInfo);
+          instrumentor->runBeforePipeline(opName, parentInfo);
         auto result = runPipeline(mgr->getPasses(), &op, am.nest(&op));
         if (instrumentor)
-          instrumentor->runAfterPipeline(mgr->getOpName(), parentInfo);
+          instrumentor->runAfterPipeline(opName, parentInfo);
 
         if (failed(result))
           return signalPassFailure();
       }
     }
   }
 }
 
llvm::parallelForEach(
           diagHandler.setOrderIDForThread(nextID);
 
           // Get the pass manager for this operation and execute it.
           auto &it = opAMPairs[nextID];
           auto *pm =
               findPassManagerFor(pms, it.first->getName().getIdentifier());
           assert(pm && "expected valid pass manager for operation");
 
+          Identifier opName = pm->getOpName(*getOperation()->getContext());
           if (instrumentor)
-            instrumentor->runBeforePipeline(pm->getOpName(), parentInfo);
+            instrumentor->runBeforePipeline(opName, parentInfo);
           auto pipelineResult =
               runPipeline(pm->getPasses(), it.first, it.second);
           if (instrumentor)
-            instrumentor->runAfterPipeline(pm->getOpName(), parentInfo);
+            instrumentor->runAfterPipeline(opName, parentInfo);
 
           // Drop this thread from being tracked by the diagnostic handler.
           // After this task has finished, the thread may be used outside of
           // this pass manager context meaning that we don't want to track
           // diagnostics from it anymore.
           diagHandler.eraseOrderIDForThread();
 
           // Handle a failed pipeline result.
return module.emitError()
          << *crashReproducerFileName << "'";
 }
 
 //===----------------------------------------------------------------------===//
 // PassManager
 //===----------------------------------------------------------------------===//
 
 PassManager::PassManager(MLIRContext *ctx, bool verifyPasses)
-    : OpPassManager(Identifier::get(ModuleOp::getOperationName(), ctx), ctx,
+    : OpPassManager(Identifier::get(ModuleOp::getOperationName(), ctx),
                     verifyPasses),
-      passTiming(false), localReproducer(false) {}
+      context(ctx), passTiming(false), localReproducer(false) {}
 
 PassManager::~PassManager() {}
 
 /// Run the passes within this manager on the provided module.
 LogicalResult PassManager::run(ModuleOp module) {
   // Before running, make sure to coalesce any adjacent pass adaptors in the
   // pipeline.
   getImpl().coalesceAdjacentAdaptorPasses();

mlir/lib/Pass/PassStatistics.cpp

if (auto *adaptor = dyn_cast<OpToOpPassAdaptor>(pass)) {
       auto mgrs = adaptor->getPassManagers();
       if (mgrs.size() > 1) {
         printPassEntry(os, indent, adaptor->getAdaptorName());
         indent += 2;
       }
 
       // Print each of the children passes.
       for (OpPassManager &mgr : mgrs) {
-        auto name = ("'" + mgr.getOpName().strref() + "' Pipeline").str();
+        auto name = ("'" + mgr.getOpName() + "' Pipeline").str();
         printPassEntry(os, indent, name);
         for (Pass &pass : mgr.getPasses())
           printPass(indent + 2, &pass);
       }
       return;
     }
 
     // Otherwise, we print the statistics for this pass.