[RFC] Dynamic Pass Pipeline

Today the pass pipeline is fairly static: the user creates a PassManager and populate it with the nested structure that defines the passes to execute.
The new features implemented in ⚙ D86392 Implement a new kind of Pass: dynamic pass pipeline aim to include in a pass pipeline a different kind of pass: Dynamic Passes.
Instead of inheriting from PassWrapper and defining runOnOperation() to perform a transformation, a Dynamic Pass inherits from DynamicPassWrapper and instead of performing a transformation (even though it can), it returns an optional OpPassManager * that defines a sub-pipeline to schedule on the currently visited operation.

This feature can be used for example to implement a sub-pipeline that would run only on an operation with specific attributes. Another example would be to compute a cost model and dynamic schedule a pipeline based on the result of this analysis.

Dumb question, but how is this different from having a pass that just creates a OpPassManager and runs passes internally to itself? Sharing the outer pass manager thread pool?

This is a very good question!

Not much different!
An alternative to this feature would be to allow to create an OpPassManager and run it, but that’s not possible today. It is mostly justified by the PassManager infrastructure integration: all options (verify, print-ir-after-all, …), instrumentation, statistics, and thread-pool sharing (as you mentioned).

Independently of this, I thought there was no way to run other passes from within a pass because you aren’t allowed to mutate anything outside of the op the pass is called on (for eg. outside of a func op). If you did, you’d have inner passes running in parallel on different functions with the outer pass also mutating the same functions concurrently – leading to race conditions. It would only work with -disable-pass-threading. But I don’t know if this has changed.

If your pass is operating on a Module (for example a nested module), you could create a PassManager on this nested module your pass is scheduled on. Otherwise I’ve seen people trying to do what you describe but with constructing the pass so that it would filter on the current operation, nothing really pretty there.

It seems more nicely layered to me if these are passed in as an argument to OpPassManager constructor so that they can be properly “inherited” (much as an OpBuilder can take a listener to inherit from a parent OpBuilder). However, looking at the code, that is a lot more work than your patch :slight_smile:

It looks like the key part of your patch is the OpToOpPassAdaptor::runPipeline call here:

  if (pass->isDynamic()) {
    OpPassManager *dynamicPipeline = pass->buildDynamicPipeline();
    if (dynamicPipeline)
      OpToOpPassAdaptor::runPipeline(dynamicPipeline->getPasses(), op, am);
  } else {
    pass->runOnOperation();
  }

Rather than this “return the OpPassManager” thing, which is quite rigid, how about just exposing a way for passes to call back into the OpToOpPassAdaptor::runPipeline. Consider, say, a pass that wants to run the same OpPassManager multiple times until it reaches a fixedpoint or other goal. E.g. you could have repeat-pipeline{repeat-count=5, inner-pipeline=...} which would call OpToOpPassAdaptor::runPipeline 5 times. Or inliner{custom-simplification-pipeline=...}.

I can look into this.

Technically you can do all of this: just re-enqueue a copy of yourself as a “continuation” in the end of the pipeline :slight_smile:

Update the revision, wanna take another look?

Just looked. I don’t understand why we even need the concept of a “dynamic pass”. A callback to run sub-pipelines seems like a general utility that should be available to all passes.

1 Like

Absolutely, I also realized this shortly after introducing the callback, and just updated the code! Thanks

We basically have a pass that does this in IREE today. Having something more nicely integrated would be great. One question I have is how will this interact with dialect registration? That’s something we’re running up against in our current version as the global registry goes away. Now the pass has to declare all the dialects the pipeline might use, which means it’s not super dynamic :-/ I’m hoping this new infrastructure to make dynamic pass pipelines first-class can help here. If one of the passes in a pipeline uses a dialect that is not otherwise loaded, how can that be handled? By the time runOnOperation is called, it’s too late to load new dialects (because we’re multithreaded now). But we should be able to do it while building the pipeline. Can this provide a hook so that the passmanager can be created as part of this pass being created rather than when it’s run?

I see it as orthogonal: the dynamic pass will need to expose through getDependentDialects the Dialects to load in the Context.

getDependentDialects is the intended mechanism. How you get the information there is a pass implementation problem: for instance a dynamic pass that would load dynamically a backend would need to have this contract with their backend somehow to retrieve the dependent dialects.
This can be implemented also with a global static registry if it makes sense in your context, but that’s all part of your specific registration mechanism.

I don’t get the question here? Nothing prevents you from assembling everything as pass construction already, the getDependentDialects() method isn’t static and will be called on your already-built pass instance.

Also if you know at pass construction the exact pass pipeline to always run, it isn’t clear to me why you need a dynamic pass in the first place? You can just insert the end pipeline instead.

+1, this is a really elegant design, Mehdi! Any time a pass calls other passes, it needs to be able to somehow implement getDependentDialects.

IREE"s TranslateExecutablesPass will need to implement getDependentDialects, by either:

  1. building the translation pipelines for the matched backends and calling getDependentDialects on that. This would require removing the (unused) targetOp argument to that function.
  2. if we want to retain the ability for a targetOp-dependent pass pipeline (such as having the HAL::ExecutableOp embedding some sort of custom per-op code generation options), then the TargetBackend interface will need to be extended to have a getDependentDialects method.

My point is that if a standard use case is running another pipeline, it would be nice to facilitate this a bit better.

Yes (2) is what I implemented in https://github.com/google/iree/pull/3036 and it’s a bit gross. It requires that the target know all possible dialects that might be used by any of the passes (and so kind of breaks the nice chaining of getDependentDialects). (1) is exactly what I’m asking about. Does this become possible? I think the targetOp is only passed on building the pipeline to allow snooping on the context because of the current limitations of dynamic pass pipelines, so we should be OK dropping it. To build a passmanager you need a context, which we don’t have with getDependentDialects, but it looks like maybe OpPassManager is becoming more exposed, so we wouldn’t need to do that anymore (it has a comment saying not to construct it directly right now). Even building an OpPassManager in the pass constructor or getDependentDialects runs into issues:

assertion failed at mlir/include/mlir/Pass/Pass.h:149 in detail::PassExecutionState &mlir::Pass::getPassState(): passState && “pass state was never initialized”

If you remove the targetOp dependence, then you can do something like:

class MyTargetBackend : TargetBackend {
  MyTargetBackend() {
     pipeline = createPipeline();
  }
  void getDependentDialectsForTarget(DialectRegistry &registry) { return pipeline.getDependentDialects(registry); }
  OpPassManager pipeline;
};

Then instead of targetBackend->buildTranslationPassPipeline you would have targetBackend->getTranslationPipeline.

I think if we do this, it should work cleanly.

If we later need targetOp dependent behavior, we can accomplish this with the new dynamic pass invocation support. That is, we push down the sniffing of targetOp from “build*Pipeline” down into some pass’s runOnOperation. (we could have a pass whose only purpose is to do such sniffing and dispatching to different pipelines)

Yeah this is what I tried (or several variations of approximately the same thing). You can’t do it because you don’t have an MLIRContext and you can’t create a PassManager without one. By the time you have a context, it’s too late to load the dialects.

Is there a fundamental reason for this, or is it just a limitation of the current implementation?

If we view creating the TargetBackend’s as part of creating the pass pipeline, then I don’t see any fundamental conflict. Is there a situation. Are there times when we create a TargetBackend in the absence of creating a pass pipeline? If so, then we will need to split that from the part of the TargetBackend that depends on the MLIRContext. Such as a method targetBackend->getCompilerInterface(mlirContext) -> std::unique_ptr<TargetBackendCompiler> that returns a TargetBackendCompiler interface that contains the MLIRContext-dependent pieces.

The issue is that creating a OpPassManager requires an OperationName, which is uniqued in the MLIRContext and on construction of a pass we don’t have an MLIRContext available. This doesn’t really have anything to do with IREE’s TargetBackend or TranslateExecutables specifically. In order for a Pass to truly run another pass using the pass pipeline and have its dependent dialects recursively compose, the pass needs access to the MLIRContext in the constructor.

Why do you need to create an OpPassManager to answer the getDependentDialects query? (maybe setting up a meeting would be easier to get through all this?)

Because the only source of truth about what dialects are needed is the getDependentDialects on the pipeline that is about to be run. Consider if some pass deep inside the pipeline changes to use a new dialect as an intermediate step of its lowering – that one pass updating its own getDependentDialects should Just Work and not require duplicating that information in all places that use that pass.

1 Like