[mlir] make MLIRContext available to getDependentDialects

[ftynse]

Summary

The dependent dialect mechanism allows the pass manager to load upfront the
dialects containing the entities that may be produced by the passes. Until now,
it only allowed the pass to provide a fixed list of dialects. Make a refrence
to the MLIRContext object available in getDependentDialects API to make this
more flexible. For example, an interface-based pass may query if interface
implementations are provided for certain ops and only require the dialects to
be loaded when they are, thus potentially reducing the number of dialects to
load. This is particularly useful for external, out-of-tree, interface
implementations that need to be plugged into passes from a different tree.
Since most of the passes are unlikely to use this functionality, the original
context-less getDependentDialects overload is preserved as is.

[mehdi_amini]

I don't quite get the use case, but more importantly this makes the getDependentDialects "contextual": the context can change between the time getDependentDialects is called and the moment the pass pipeline runs.

[ftynse]

See the child patch for the use case.

[ftynse]

Ping.

We've hit another case that is solved by this: bufferization. For example, bufferization of tensor.generate produces scf.for. The logic is provided by the external interface implementation. The pass itself doesn't know about the exact logic and only uses the interface, but would have to still list the scf dialect as dependent. With this, the pass would be able to query the known interface implementations for the list of dependent passes via the context. I expect the same to be true for many other interface-based passes: transform, tiling, canonicalizer to some extent (now we have to load the dialects canonicalization depends on in parent dialect), etc.

but more importantly this makes the getDependentDialects "contextual"

Yes, and I think this is actually the desired behavior for interface-driven passes.

I know there is some explicit expectation of passes being somehow agnostic to what is registered/loaded in the context. However, this is more of a "thou shalt not" recommendation, since practically one can easily write a pass conditioned on something being available in the context by simply querying it in the pass entry function.

the context can change between the time getDependentDialects is called and the moment the pass pipeline runs.

Technically, it can. Practically, this seems highly unlikely. There are two dozen lines between the call to getDependentDialects and the passes actually running https://github.com/llvm/llvm-project/blob/6c3990acfbb933a61e2c74332bb252121c06bd14/mlir/lib/Pass/Pass.cpp#L791-L817, and they don't return control back to the caller. So there seem to be two possibilities: (1) another thread modifies the context between these two calls, but modifying the context concurrently sounds like a terrible idea to start with; (2) one of the getDependentDialects modifies the context, which is solvable by sprinkling some const on it or providing some immutable wrapper.

[mehdi_amini]

In D128528#3636290, @ftynse wrote:

Ping.

I stil have the same concerns, despite the use cases you're trying to enable. I haven't had time to try to collect all the invariants this can potentially break and I don't think I'll be able to get to it before I'm back from my current vacation (on 7/18): this is tricky.

[ftynse]

Okay, I'll ping you again then.

Leaving this here so I don't forget: we have always had a pass with registration-dependent behavior -- canonicalizer. Depending on whether the op is registered (unregistered ops can be created via Operation::create), it will or will not canonicalize it. This works because canonicalization and folding are privileged and have opaque hooks that the context is aware of. I suspect that if we were to write the canonicalizer today using interfaces, we would have hit the same issue. What I'd like is to generalize this sort of passes, with appropriate safety mechanisms, in a way that doesn't require privileged hooks in the context.

[mehdi_amini]

I don’t quite get what is privileged about the canonicalizer here? At least I don’t believe it comes from hooks in the context, the way it is handled there is by saying that canonicalization patterns introduce static dialect dependencies: you necessarily load dialect B when loading dialect A if a canonicalization pattern from dialect A is emitting an entity from dialect B.
But this is more of a policy than anything and we could move canonicalization patterns to use an interface without any change to how we manage this.

[ftynse]

There are several privileged things that may be a bit conflated. For what you are asking specifically, we have these static dialect dependencies because of canonicalization. We don't introduce them because of other passes AFAIK. This makes the canonicalization privileged in a sense. We could also add these static dependencies because of other passes (e.g., tensor bufferizes to scf so it should depend on scf in the case above), but it doesn't sound like it would scale well. I am looking for a mechanism to generalize and modularize this.

Side note: don't feel pressured to answer, we can look at code and specific options separately.

[mehdi_amini]

In D128528#3638245, @ftynse wrote:

Depending on whether the op is registered

Something I wanted to point out is that "registered" here has a different meaning: an op is "registered" after a dialect is loaded, not just merely registered.
Part of my concerns here (which I need to think more about, it is a bit fuzzy for me right now) is about acting on "registered but not loaded" here.

(unregistered ops can be created via Operation::create),

Just a nit: but unless the dialect is loaded, this isn't possible by default (there is a flag in the context to allow it explicitly, but that's not really how systems are intended to be setup). If the dialect is loaded, then yeah you can do that, but dialects are supposed to register their operation on load so I think it is "work as intended".

[ftynse]

We discussed offline and found a different solution that loads dialects together with the extension.

mlir/include/mlir/Pass/Pass.h

public:
   /// Returns the derived pass name.
   virtual StringRef getName() const = 0;
 
   /// Register dependent dialects for the current pass.
   /// A pass is expected to register the dialects it will create entities for
   /// (Operations, Types, Attributes), other than dialect that exists in the
   /// input. For example, a pass that converts from Linalg to Affine would
   /// register the Affine dialect but does not need to register Linalg.
+  /// The context may be used to identify the dialects that may be produced by
+  /// the pass based on the presence of some registerable objects, such as
+  /// interfaces and dialect extensions.
   virtual void getDependentDialects(DialectRegistry &registry) const {}
+  virtual void getDependentDialects(DialectRegistry &registry,
+                                    MLIRContext &context) const {
+    getDependentDialects(registry);
+  }
 
   /// Return the command line argument used when registering this pass. Return
   /// an empty string if one does not exist.
   virtual StringRef getArgument() const { return ""; }
 
   /// Return the command line description used when registering this pass.
   /// Return an empty string if one does not exist.
   virtual StringRef getDescription() const { return ""; }

mlir/include/mlir/Pass/PassManager.h

public:
   void dump();
 
   /// Merge the pass statistics of this class into 'other'.
   void mergeStatisticsInto(OpPassManager &other);
 
   /// Register dependent dialects for the current pass manager.
   /// This is forwarding to every pass in this PassManager, see the
   /// documentation for the same method on the Pass class.
-  void getDependentDialects(DialectRegistry &dialects) const;
+  void getDependentDialects(DialectRegistry &dialects,
+                            MLIRContext &context) const;
 
   /// Enable or disable the implicit nesting on this particular PassManager.
   /// This will also apply to any newly nested PassManager built from this
   /// instance.
   void setNesting(Nesting nesting);
 
   /// Return the current nesting mode.
   Nesting getNesting();

mlir/lib/Pass/Pass.cpp

 
 void OpPassManager::dump() {
   llvm::errs() << "Pass Manager with " << impl->passes.size() << " passes: ";
   ::printAsTextualPipeline(impl->passes, llvm::errs());
   llvm::errs() << "\n";
 }
 
 static void registerDialectsForPipeline(const OpPassManager &pm,
-                                        DialectRegistry &dialects) {
+                                        DialectRegistry &dialects,
+                                        MLIRContext &context) {
   for (const Pass &pass : pm.getPasses())
-    pass.getDependentDialects(dialects);
+    pass.getDependentDialects(dialects, context);
 }
 
-void OpPassManager::getDependentDialects(DialectRegistry &dialects) const {
-  registerDialectsForPipeline(*this, dialects);
+void OpPassManager::getDependentDialects(DialectRegistry &dialects,
+                                         MLIRContext &context) const {
+  registerDialectsForPipeline(*this, dialects, context);
 }
 
 void OpPassManager::setNesting(Nesting nesting) { impl->nesting = nesting; }
 
 OpPassManager::Nesting OpPassManager::getNesting() { return impl->nesting; }
 
 LogicalResult OpPassManager::initialize(MLIRContext *context,
                                         unsigned newInitGeneration) {
static OpPassManager *findPassManagerFor(MutableArrayRef<OpPassManager> mgrs,
   });
   return it == mgrs.end() ? nullptr : &*it;
 }
 
 OpToOpPassAdaptor::OpToOpPassAdaptor(OpPassManager &&mgr) {
   mgrs.emplace_back(std::move(mgr));
 }
 
-void OpToOpPassAdaptor::getDependentDialects(DialectRegistry &dialects) const {
+void OpToOpPassAdaptor::getDependentDialects(DialectRegistry &dialects,
+                                             MLIRContext &context) const {
   for (auto &pm : mgrs)
-    pm.getDependentDialects(dialects);
+    pm.getDependentDialects(dialects, context);
 }
 
 LogicalResult OpToOpPassAdaptor::tryMergeInto(MLIRContext *ctx,
                                               OpToOpPassAdaptor &rhs) {
   // Functor used to check if a pass manager is generic, i.e. op-agnostic.
   auto isGenericPM = [&](OpPassManager &pm) { return !pm.getOpName(); };
 
   // Functor used to detect if the given generic pass manager will have a

 LogicalResult PassManager::run(Operation *op) {
   MLIRContext *context = getContext();
   assert(op->getName() == getOpName(*context) &&
          "operation has a different name than the PassManager or is from a "
          "different context");
 
   // Register all dialects for the current pipeline.
   DialectRegistry dependentDialects;
-  getDependentDialects(dependentDialects);
+  getDependentDialects(dependentDialects, *context);
   context->appendDialectRegistry(dependentDialects);
   for (StringRef name : dependentDialects.getDialectNames())
     context->getOrLoadDialect(name);
 
   // Before running, make sure to finalize the pipeline pass list.
   if (failed(getImpl().finalizePassList(context)))
     return failure();
 

mlir/lib/Pass/PassDetail.h

public:
   /// merge.
   LogicalResult tryMergeInto(MLIRContext *ctx, OpToOpPassAdaptor &rhs);
 
   /// Returns the pass managers held by this adaptor.
   MutableArrayRef<OpPassManager> getPassManagers() { return mgrs; }
 
   /// Populate the set of dependent dialects for the passes in the current
   /// adaptor.
-  void getDependentDialects(DialectRegistry &dialects) const override;
+  void getDependentDialects(DialectRegistry &dialects,
+                            MLIRContext &context) const override;
 
   /// Return the async pass managers held by this parallel adaptor.
   MutableArrayRef<SmallVector<OpPassManager, 1>> getParallelPassManagers() {
     return asyncExecutors;
   }
 
   /// Returns the adaptor pass name.
   std::string getAdaptorName();

mlir/test/lib/Pass/TestDynamicPipeline.cpp

 public:
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestDynamicPipelinePass)
 
   StringRef getArgument() const final { return "test-dynamic-pipeline"; }
   StringRef getDescription() const final {
     return "Tests the dynamic pipeline feature by applying "
            "a pipeline on a selected set of functions";
   }
-  void getDependentDialects(DialectRegistry &registry) const override {
+  void getDependentDialects(DialectRegistry &registry,
+                            MLIRContext &context) const override {
     OpPassManager pm(ModuleOp::getOperationName(),
                      OpPassManager::Nesting::Implicit);
     (void)parsePassPipeline(pipeline, pm, llvm::errs());
-    pm.getDependentDialects(registry);
+    pm.getDependentDialects(registry, context);
   }
 
   TestDynamicPipelinePass() = default;
   TestDynamicPipelinePass(const TestDynamicPipelinePass &) {}
 
   void runOnOperation() override {
     Operation *currentOp = getOperation();
 

mlir/unittests/Pass/PassManagerTest.cpp

for (func::FuncOp func : module->getOps<func::FuncOp>()) {
     EXPECT_TRUE(func->getAttr("isFunc").cast<BoolAttr>().getValue());
 
     bool isSecret = func.getName() == "secret";
     ASSERT_TRUE(func->getAttr("isSecret").isa<BoolAttr>());
     EXPECT_EQ(func->getAttr("isSecret").cast<BoolAttr>().getValue(), isSecret);
   }
 }
 
-namespace {
 struct InvalidPass : Pass {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InvalidPass)
 
   InvalidPass() : Pass(TypeID::get<InvalidPass>(), StringRef("invalid_op")) {}
   StringRef getName() const override { return "Invalid Pass"; }
   void runOnOperation() override {}
   bool canScheduleOn(RegisteredOperationName opName) const override {
     return true;
   }
 
   /// A clone method to create a copy of this pass.
   std::unique_ptr<Pass> clonePass() const override {
     return std::make_unique<InvalidPass>(
         *static_cast<const InvalidPass *>(this));
   }
 };
-} // namespace
 
 TEST(PassManagerTest, InvalidPass) {
   MLIRContext context;
   context.allowUnregisteredDialects();
 
-  // Create a module
+  // Create a module.
   OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));
 
-  // Add a single "invalid_op" operation
+  // Add a single "invalid_op" operation.
   OpBuilder builder(&module->getBodyRegion());
   OperationState state(UnknownLoc::get(&context), "invalid_op");
   builder.insert(Operation::create(state));
 
   // 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) {
TEST(PassManagerTest, InvalidPass) {
   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>()), "");
 }
 
+/// A dummy pass that lists as dependent all dialects registered with the
+/// context.
+struct ContextuallyDependentPass
+    : PassWrapper<ContextuallyDependentPass, OperationPass<>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ContextuallyDependentPass)
+
+  StringRef getName() const override {
+    return "test-contextually-dependent-pass";
+  }
+
+  void runOnOperation() override {}
+  void getDependentDialects(DialectRegistry &registry,
+                            MLIRContext &context) const override {
+    context.getDialectRegistry().appendTo(registry);
+  }
+};
+
+TEST(PassManagerTest, ContextuallyDependentPass) {
+  // Create an empty module.
+  MLIRContext context;
+  OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));
+
+  // Run the pass using the pass manager. This loads all dialects declared as
+  // dependent by the pass. The pass declares as dependent every registered
+  // dialect. Since the Func dialect is not registered, it should not be loaded
+  // after the pass manager completes.
+  PassManager pm(&context);
+  pm.addPass(std::make_unique<ContextuallyDependentPass>());
+  LogicalResult result = pm.run(module.get());
+  ASSERT_TRUE(succeeded(result));
+  EXPECT_EQ(context.getLoadedDialect<func::FuncDialect>(), nullptr);
+
+  // Register the Func dialect with the context and run the same pass manager
+  // again. Now the dialect should have been loaded by the pass manager.
+  DialectRegistry registry;
+  registry.insert<func::FuncDialect>();
+  context.appendDialectRegistry(registry);
+  result = pm.run(module.get());
+  ASSERT_TRUE(succeeded(result));
+  EXPECT_NE(context.getLoadedDialect<func::FuncDialect>(), nullptr);
+}
+
 } // namespace