Dynamic Pass Scheduling and Dialect Loading

MLIR
1

Hi,

Please bare with me, my understanding of the interals of dialect loading might not be fully correct here, but we’ve hit a snag with dynamic pass scheduling that I’ll try to descrive below.

We have a pass that internally builds pass manager and runs it via the Pass::runPipeline method. We need to do this before we want to avoid running a pipeline on operations that have a certain attribute.

Normally when running a pipeline from a PM via the PassManager::run method the first thing that happens is that the PM loads all dependent dialects into the MLIRContext. This doesn’t happen when a PM is run from within a pass via Pass::runPipeline since it is being called in a multi-threaded context. Because of this some dialects that are required by the dynamically scheduled passes aren’t loaded and we hit an assertion.

I can think of two possible solutions to this:

  1. List the dependent dialects of all dynamically scheduled passes as dependent dialects of the pass that contains the PM. This isn’t ideal since it requires us to maintain a separate list of all dialects that might be required. If new passes are added and they have dependencies then the list will need to be updated.
  2. Register all dependent dialects of all dynamically scheduled passes in the MLIRContext upfront and call MLIRContext::loadAllAvailableDialects(). This has the same issue as 1. we need to maintain list of dependent dialects, its just now that they are in some C++ code that registers them with the context instead of in the tablgen for the pass, we also need to call the loadAllAvailableDialects() API.
  3. Build the pipeline for the pass externally before the top level PM is run and pass it into the constructor for the pass that dynamically runs the passes. If there is a way to extract the dependent dialects from all passes in a PM they could then be extracted from the inner PM and added to the top level PM. This would remove the need to manually maintain a list of dependencies as in 1 and 2 but I’m not sure if this is actually possible.

Has anyone come across this issue before and if so is there a solution that avoids the pitfalls of 1 and 2?

Ta,
Jack.

2

This is the way to go. However to address the drawback you mention, in general you first build the various pipeline you want to run, and then call their getDependentDialects() from the getDependentDialects() implementation of the pass invoking runPipeline(). That way there is no redundant list to maintain.

Here is an example: llvm-project/mlir/lib/Transforms/CompositePass.cpp at 13261e856aa54e6c86b500c6d0bb329f08ad5eb7 · llvm/llvm-project · GitHub

3

Grand, thanks Mehdi that pass looks like what we need.

4

Implementing the getDependentDialects hook as suggested in the example solved the problem, thank you for your help.