Implement a new kind of Pass: dynamic pass pipeline

[mehdi_amini]

Summary

Instead of performing a transformation, such pass yields a new pass pipeline
to run 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.

Discussion: https://llvm.discourse.group/t/rfc-dynamic-pass-pipeline/1637

[bondhugula]

mlir/include/mlir/Pass/Pass.h

169

Who owns / has to free this OpPassManager?

mlir/lib/Pass/Pass.cpp

368

You are leaking memory here?

[bondhugula]

mlir/include/mlir/Pass/Pass.h

166–167

Rephrase/reorder this a bit to start with "Builds an OpPassManager containing .... This is equivalent ..."?

[mehdi_amini]

mlir/include/mlir/Pass/Pass.h

169

I wrote it initially as returning a reference and changed it to pointer to allow for returning nullptr, I didn't think it could be confusing from an ownership question (also I'd return a unique_ptr if I was transferring ownership). I'll update the doc!

The intent is that the pass can build the pipeline once, or caching multiple versions and return the one to apply.

[bondhugula]

mlir/include/mlir/Pass/Pass.h

169

Thanks. Then, it would be good to document that nullptr could be returned. Is nullptr is equivalent to an empty pipeline?

mlir/lib/Pass/Pass.cpp

365–366

The comment is grammatically broken at "... either ..." and also ".. retried ..."

[bondhugula]

You may also want to add an [MLIR] tag to the title to not confuse LLVM proper folks.

[mehdi_amini]

In D86392#2232401, @bondhugula wrote:

You may also want to add an [MLIR] tag to the title to not confuse LLVM proper folks.

I think this is noise in my git log and emails: it is likely better solved with tooling.

[bondhugula]

In D86392#2232403, @mehdi_amini wrote:

In D86392#2232401, @bondhugula wrote:

You may also want to add an [MLIR] tag to the title to not confuse LLVM proper folks.

I think this is noise in my git log and emails: it is likely better solved with tooling.

Sure - sounds good. But I had a couple of other questions above - hope you didn't miss them.

[mehdi_amini]

Use a callback instead to drive the pipeline run

[mehdi_amini]

mlir/include/mlir/Pass/Pass.h

169

Yes.

[mehdi_amini]

Rename: renameDynamicPipeline -> runDynamicPipeline

[mehdi_amini]

Remove the special pass and make the callback universally available

[silvas]

mlir/include/mlir/Pass/Pass.h

45

Is it possible to allow this to take an op or set of ops to run on?

For example, an inliner doing SCC visitation might want to runPipeline(pipeline, currentSCC)

[mehdi_amini]

mlir/include/mlir/Pass/Pass.h

45

I would imagine an inlined working on SCC to have to handle SCC somehow: the pass infra does not allow that just now so I'm not sure what would you do with this extra argument?

[silvas]

mlir/include/mlir/Pass/Pass.h

45

Yes, it has to handle the SCC, but it will want to run a simplification pipeline on the current SCC before moving to the parent SCC. An inliner is a ModulePass in MLIR, and so the current runPipeline would just re-run the pipeline on the entire module. But the inliner only wants to run the pipeline on some subset of functions in the module.

[mehdi_amini]

Allow to run the pipeline on any op nested under the current operation

[mehdi_amini]

mlir/include/mlir/Pass/Pass.h

45

Right but it would have to reschedule itself the function passes individually on the functions in the SCC.
As long as the PM in MLIR does not have a concept of SCCs, I'm not sure how to express this otherwise?

[silvas]

Now that this functionality has evolved, is there a better name to describe this functionality than "dynamic pipeline"? It's not associated with pipelines in any specific way, and it is not associated with dynamism per se (e.g. inliner could use a statically hardcoded simplification pipeline). Perhaps "recursive pass invocation"?

Also, I assume this needs doc updates?

mlir/include/mlir/Pass/Pass.h

45

Correct, all I meant was that the inliner needed a hook to manually run passes on operations in some effectively arbitrary way. Not that the PM needed special knowledge.

The new code looks good.

mlir/lib/Pass/Pass.cpp

344

change to assert or add test?

355

should this be an assert? if not it needs a test.

mlir/test/Pass/dynamic_pipeline_nested.mlir

1 ↗

(On Diff #288737)

use - instead of _ in test names.

5 ↗

(On Diff #288737)

I find the FileCheck lines here hard to follow. Do we have to actually check all the inner_mod1 stuff and "Dump Before" lines? (and any mention of CSE itself). Can we just do

NESTED: @foo
NESTED: @bar
NONESTED-NOT: @foo 
NONESTED-NOT: @bar

[mehdi_amini]

In D86392#2247937, @silvas wrote:

Now that this functionality has evolved, is there a better name to describe this functionality than "dynamic pipeline"? It's not associated with pipelines in any specific way, and it is not associated with dynamism per se (e.g. inliner could use a statically hardcoded simplification pipeline). Perhaps "recursive pass invocation"?

I still see this as dynamic: from the pass pipeline point of view there is a "pass" which will execute "dynamically" sub-pass-pipelines. This feature enables dynamism in the pass pipeline where it does not exist today.

Also, I assume this needs doc updates?

Right, I'll document the final version :)

mlir/lib/Pass/Pass.cpp

355

Could be an assert, but seems more friendly to have this always diagnosed and gracefully fail: it is easy to get wrong in a pass I suspect.

mlir/test/Pass/dynamic_pipeline_nested.mlir

5 ↗

(On Diff #288737)

I don't think so: @foo and @bar will be in the output always (we print the IR...).

I'm checking at which granularity we apply the pipeline by checking where the CSE pass applies: will it run on the inner_mod1 module, where the dynamic pass is scheduled? Or will it run on each of the nested functions individually?

[mehdi_amini]

Add comments

mlir/lib/Pass/Pass.cpp

344

Made is a report_fatal_error (so that the error message is richer than an assert)

355

Added a test for this one now.

mlir/test/Pass/dynamic_pipeline_nested.mlir

5 ↗

(On Diff #288737)

I increase the doc in the test, let me know if it more clear.

[silvas]

In D86392#2248436, @mehdi_amini wrote:

In D86392#2247937, @silvas wrote:

Now that this functionality has evolved, is there a better name to describe this functionality than "dynamic pipeline"? It's not associated with pipelines in any specific way, and it is not associated with dynamism per se (e.g. inliner could use a statically hardcoded simplification pipeline). Perhaps "recursive pass invocation"?

I still see this as dynamic: from the pass pipeline point of view there is a "pass" which will execute "dynamically" sub-pass-pipelines. This feature enables dynamism in the pass pipeline where it does not exist today.

I mainly don't like the association with "pipeline", since "pipeline" doesn't really mean anything specific currently, except our concept of registering a "pass pipeline" (which is purely a concept for registration / textual parsing of pipelines). In fact, the name of the method is runPipeline which is confusing IMO. It runs an OpPassManager, and we don't really refer to an OpPassManager as a "pipeline" in any public API or documentation AFAICT. We seem to use "pipeline" mostly when talking about a textual description.

Thoughts? Maybe OpPassManager needs commentary that it can be informally be referred to as "a pipeline" and maybe should be renamed to OpPassPipeline? In fact, can OpPassManager just be a Pass, and then we can call this function runPass?

[mehdi_amini]

In D86392#2250121, @silvas wrote:

In D86392#2248436, @mehdi_amini wrote:

In D86392#2247937, @silvas wrote:

Now that this functionality has evolved, is there a better name to describe this functionality than "dynamic pipeline"? It's not associated with pipelines in any specific way, and it is not associated with dynamism per se (e.g. inliner could use a statically hardcoded simplification pipeline). Perhaps "recursive pass invocation"?

I still see this as dynamic: from the pass pipeline point of view there is a "pass" which will execute "dynamically" sub-pass-pipelines. This feature enables dynamism in the pass pipeline where it does not exist today.

I mainly don't like the association with "pipeline", since "pipeline" doesn't really mean anything specific currently,

I thought "pass pipeline" was a fairly common concept, I'm using it as "sequence of passes to apply on some IR". Anecdotally a quick Google query on "pass pipeline compiler" seems to find mostly reference in the same direction (doc from LLVM, Rust, Haskell, Numba).

except our concept of registering a "pass pipeline" (which is purely a concept for registration / textual parsing of pipelines)

Our pipeline aren't tied to parsing a textual representation: https://mlir.llvm.org/docs/PassManagement/#pass-pipeline-registration ; note that the "textual" part is in a separate section after this one.

It runs an OpPassManager, and we don't really refer to an OpPassManager as a "pipeline" in any public API or documentation AFAICT.

https://mlir.llvm.org/docs/PassManagement/#pass-manager "The PassManager class is used to configure and run a pipeline. The OpPassManager class is used to schedule passes to run at a specific level of nesting."

Thoughts? Maybe OpPassManager needs commentary that it can be informally be referred to as "a pipeline" and maybe should be renamed to OpPassPipeline? In fact, can OpPassManager just be a Pass, and then we can call this function runPass?

runPass is misleading to me because of the singular here: I see a difference between one pass and a sequence of passes. It could be runPasses maybe?

[silvas]

I still consider this feature naming (such as test case naming) very confusing. It's actually very simple to explain: "it allows a pass to call other passes". The name "dynamic pipeline" is very confusing -- like I'll need to re-explain to every new MLIR user that encounters the feature what the feature really is. "user: How can I call a pass from within my pass? mlir community: oh, use our dynamic pipeline feature. user: what is that? mlir community: it's a way to call a pass from within another pass; user: oh".

Maybe "dynamic pass invocation"?

runPass is misleading to me because of the singular here: I see a difference between one pass and a sequence of passes. It could be runPasses maybe?

runPasses SGTM.

runPass singular does make sense when an OpPassManager is itself a pass that simply runs a sequence of inner passes (such as in the llvm new PM). That was one of my major learnings from working on the LLVM new PM: there is actually no need for a thing to "manage" passes (which *was* necessary in the legacy PM design). After properly factoring things, the only remnant of the previous "management" component is 1) Adaptors that convert from (in MLIR speak) one operation type to another (module pass vs function pass etc.) 2) Passes that aggregate a sequence of other passes (there's no "management"; just "run a sequence") and 3) analysis manager. That is, we could have func(cse,canonicalize) really just be syntax sugar for sequence-of-passes{op-name=std.func passes=cse,canonicalize}.

[mehdi_amini]

In D86392#2250855, @silvas wrote:

I still consider this feature naming (such as test case naming) very confusing.

But it seems like you also never encountered "pipeline" which is surprising to me as it seems like a very common terminology (and we've been using it in MLIR "forever").

runPass singular does make sense when an OpPassManager is itself a pass that simply runs a sequence of inner passes (such as in the llvm new PM).

The LLVM PM entry point is run, isn't it?

That was one of my major learnings from working on the LLVM new PM: there is actually no need for a thing to "manage" passes (which *was* necessary in the legacy PM design).

I don't share your view on the new PM apparently. The new PM is a departure in that there is no need to manage *pass dependencies* statically, but that does not mean that passes aren't managed entirely: there is a piece of code that traverse the IR and schedule the execution of these passes on an IR unit, this management hasn't really changed.

After properly factoring things, the only remnant of the previous "management" component is 1) Adaptors that convert from (in MLIR speak) one operation type to another (module pass vs function pass etc.) 2) Passes that aggregate a sequence of other passes (there's no "management"; just "run a sequence") and 3) analysis manager. That is, we could have func(cse,canonicalize) really just be syntax sugar for sequence-of-passes{op-name=std.func passes=cse,canonicalize}.

I don't really get your point: I still see a difference between the management and the passes themselves: the 1, 2, and 3 is exactly this.

[silvas]

In D86392#2250881, @mehdi_amini wrote:

In D86392#2250855, @silvas wrote:

I still consider this feature naming (such as test case naming) very confusing.

But it seems like you also never encountered "pipeline" which is surprising to me as it seems like a very common terminology (and we've been using it in MLIR "forever").

No, I'm well aware of that terminology. As you mentioned it is very common. It just doesn't really have a precise meaning we can refer to in the code, and is only a specific case of what this feature allows.

The below discussion need not block this review, since we don't currently implement OpPassManager as a pass.

runPass singular does make sense when an OpPassManager is itself a pass that simply runs a sequence of inner passes (such as in the llvm new PM).

The LLVM PM entry point is run, isn't it?

Yes, "pass" has a method "run".

That was one of my major learnings from working on the LLVM new PM: there is actually no need for a thing to "manage" passes (which *was* necessary in the legacy PM design).

I don't share your view on the new PM apparently. The new PM is a departure in that there is no need to manage *pass dependencies* statically, but that does not mean that passes aren't managed entirely: there is a piece of code that traverse the IR and schedule the execution of these passes on an IR unit, this management hasn't really changed.

There isn't actually a piece of code. It's distributed across many different pieces of code. For example, CGSCC "pass manager" is completely outside the core pass management library (it happens to be a template specialization of PassManager, but arguably would be cleaner if it wasn't, but Chandler wanted to avoid a virtual call or something; such is modern C++)

It's like if I have a GUI library and have a GUI element that stacks other GUI elements vertically. One wouldn't call it a "WidgetManager". It would be named it based on what it does, which is stack widgets vertically (and in this case for "pass management", it's to run a sequence of passes). It's the same object oriented object recursion pattern. A "WidgetManager" to me implies a single object that has global visibility of all widgets (think like a FontManager that has the responsibility of caching/recalculating all glyphs for all fonts in an application).

I don't really get your point: I still see a difference between the management and the passes themselves: the 1, 2, and 3 is exactly this.

I'm just speaking from experience here -- I've had an intern ask "well, what 'management' does it actually do? it just runs a sequence of passes. It doesn't 'manage' anything." (referring to the LLVM new PM "PassManager")

[mehdi_amini]

I'm just speaking from experience here -- I've had an intern ask "well, what 'management' does it actually do? it just runs a sequence of passes. It doesn't 'manage' anything." (referring to the LLVM new PM "PassManager")

Yeah well, that's what a "pass manager" is in the compiler literature AFAIK, I can't do much for folks who come from a non-compiler background and have a different association with the terminology "manager" or "management".

[silvas]

mlir/test/lib/Transforms/TestDynamicPipeline.cpp

25

This pass needs a getDependentDialects for any passes it might run.

[mehdi_amini]

mlir/test/lib/Transforms/TestDynamicPipeline.cpp

25

True: so it can't run for any pass that require a new dialect (like a conversion pass), but since it is just used in unit-testing we don't really care right?

[silvas]

mlir/test/lib/Transforms/TestDynamicPipeline.cpp

25

I think it's an important design aspect that we should at least have an answer to, and exhibit it with this test. We run into a situation analogous to this in IREE (and I suspect anyone using this feature will too), so it would be good to address it.

Can we do the parsePipeline in the constructor, and then implement getDependentDialects by recursively calling getDependentDialects?

[mehdi_amini]

mlir/test/lib/Transforms/TestDynamicPipeline.cpp

25

I think it's an important design aspect that we should at least have an answer to,

I have an answer, it just does not please IREE ;)

and exhibit it with this test.

I don't consider it a bug though.

We run into a situation analogous to this in IREE (and I suspect anyone using this feature will too), so it would be good to address it.

I don't quite get the problem actually, I can trivially fix it here, it'll just be ugly because taking a string it totally "unconstrained".

[silvas]

mlir/test/lib/Transforms/TestDynamicPipeline.cpp

25

I don't quite get the problem actually, I can trivially fix it here, it'll just be ugly because taking a string it totally "unconstrained".

I think that parsing the pipeline in the constructor and using that to implement getDependentDialects is not too ugly?

[mehdi_amini]

Rebase and add a getDependentDialects() in the Test pass

[rriddle]

mlir/test/lib/Transforms/TestDynamicPipeline.cpp

25

Parsing a pipeline is not something I would encourage in a constructor, at least not in any production compiler. That relies on a pass registry, and I would expect that real toolchains have a structured way of populating these "dynamic" pass pipelines.

[rriddle]

mlir/include/mlir/Pass/Pass.h

167

nit: invoke -> invoked at any time in a pass to dynamically schedule more passes.

mlir/include/mlir/Pass/PassManager.h

39

Is this used in this file?

mlir/lib/Pass/Pass.cpp

348

nit: Return this directly, it converts to failure.

354

Hmmm, this will prevent us from running multiple dynamic pipelines concurrently. Such a thing would happen if we tried to map the current inliner use case to this framework. Nesting an analysis manager is not thread-safe ATM.

812

If it isn't a direct parent, I would expect that we properly nest each parent. Otherwise, we break the tree structure of the analysis manager.

mlir/test/lib/Transforms/TestDynamicPipeline.cpp

27

Why do we need this here? This is just a test, I would remove it. There is no value to having this here.

[mehdi_amini]

Rebase on D87072

[GMNGeoffrey]

mlir/test/lib/Transforms/TestDynamicPipeline.cpp

56–60

Now that https://github.com/llvm/llvm-project/commit/c0b6bc070e78cbd20bc4351704f52d85192e8804 is in, how about moving this to the constructor so it can be used in getDepedentDialects?

[mehdi_amini]

mlir/test/lib/Transforms/TestDynamicPipeline.cpp

56–60

There is a subtlety: it is built with the currently visited operation here, which isn't fixed for this pass and so not available in the ctor.

Also... why? This sole purpose of this pass is to test the codepath in the pass manager.

[silvas]

LGTM from my side

[silvas]

mlir/test/lib/Transforms/TestDynamicPipeline.cpp

25

So how about we just don't use a textual pipeline in this example and just hardcode a set of passes to run? What value is being added to this test from using a textual pipeline?

[mehdi_amini]

mlir/test/lib/Transforms/TestDynamicPipeline.cpp

25

The only reason for this pass to exist is to test the variation on the dynamic dispatch: the options on this pass really don't make sense (there is even an option run-on-parent that should always make it fail!)

Right now the tests are relying on being able slightly tweak the pipeline I'm passing in, I could probably get around it if you feel it is important. I may have to do add an option --use-pipeline=1 / --use-pipeline=2 to switch between two predefined one.

mlir/include/mlir/Pass/Pass.h

 
 namespace mlir {
 namespace detail {
 class OpToOpPassAdaptor;
 
 /// The state for a single execution of a pass. This provides a unified
 /// interface for accessing and initializing necessary state for pass execution.
 struct PassExecutionState {
-  PassExecutionState(Operation *ir, AnalysisManager analysisManager)
-      : irAndPassFailed(ir, false), analysisManager(analysisManager) {}
+  PassExecutionState(Operation *ir, AnalysisManager analysisManager,
+                     function_ref<LogicalResult(OpPassManager &, Operation *)>
+                         pipelineExecutor)
+      : irAndPassFailed(ir, false), analysisManager(analysisManager),
+        pipelineExecutor(pipelineExecutor) {}
 
   /// The current operation being transformed and a bool for if the pass
   /// signaled a failure.
   llvm::PointerIntPair<Operation *, 1, bool> irAndPassFailed;
 
   /// The analysis manager for the operation.
   AnalysisManager analysisManager;
 
   /// The set of preserved analyses for the current execution.
   detail::PreservedAnalyses preservedAnalyses;
+
+  /// This is a callback in the PassManager that allows to schedule dynamic
+  /// pipelines that will be rooted at the provided operation.
+  function_ref<LogicalResult(OpPassManager &, Operation *)> pipelineExecutor;
 };
 } // namespace detail
 
 /// The abstract base pass class. This class contains information describing the
 /// derived pass object, e.g its kind and abstract TypeID.
 class Pass {
 public:
   virtual ~Pass() = default;
protected:
   }
 
   /// Return the MLIR context for the current function being transformed.
   MLIRContext &getContext() { return *getOperation()->getContext(); }
 
   /// The polymorphic API that runs the pass over the currently held operation.
   virtual void runOnOperation() = 0;
 
+  /// Schedule an arbitrary pass pipeline on the provided operation.
+  /// This can be invoke any time in a pass to dynamic schedule more passes.
+  /// The provided operation must be the current one or one nested below.
+  LogicalResult runPipeline(OpPassManager &pipeline, Operation *op) {
+    return passState->pipelineExecutor(pipeline, op);
+  }
+
   /// A clone method to create a copy of this pass.
   std::unique_ptr<Pass> clone() const {
     auto newInst = clonePass();
     newInst->copyOptionValuesFrom(this);
     return newInst;
   }
 
   /// Return the current operation being transformed.

mlir/include/mlir/Pass/PassManager.h

 class ModuleOp;
 class Operation;
 class Pass;
 class PassInstrumentation;
 class PassInstrumentor;
 
 namespace detail {
 struct OpPassManagerImpl;
+struct PassExecutionState;
 } // end namespace detail
 
 //===----------------------------------------------------------------------===//
 // OpPassManager
 //===----------------------------------------------------------------------===//
 
 /// This class represents a pass manager that runs passes on a specific
 /// operation type. This class is not constructed directly, but nested within
public:
   void getDependentDialects(DialectRegistry &dialects) const;
 
 private:
   /// A pointer to an internal implementation instance.
   std::unique_ptr<detail::OpPassManagerImpl> impl;
 
   /// Allow access to the constructor.
   friend class PassManager;
+  friend class Pass;
 
   /// Allow access.
   friend detail::OpPassManagerImpl;
 };
 
 //===----------------------------------------------------------------------===//
 // PassManager
 //===----------------------------------------------------------------------===//

mlir/lib/Pass/Pass.cpp

LogicalResult OpToOpPassAdaptor::run(Pass *pass, Operation *op,
   if (!op->getName().getAbstractOperation())
     return op->emitOpError()
            << "trying to schedule a pass on an unregistered operation";
   if (!op->getName().getAbstractOperation()->hasProperty(
           OperationProperty::IsolatedFromAbove))
     return op->emitOpError() << "trying to schedule a pass on an operation not "
                                 "marked as 'IsolatedFromAbove'";
 
-  pass->passState.emplace(op, am);
+  // Initialize the pass state with a callback for the pass to dynamically
+  // execute a pipeline on the currently visited operation.
+  pass->passState.emplace(
+      op, am, [&](OpPassManager &pipeline, Operation *root) {
+        if (!op->isAncestor(root)) {
+          root->emitOpError()
+              << "Trying to schedule a dynamic pipeline on an "
+                 "operation that isn't "
+                 "nested under the current operation the pass is process";
+          return failure();
+        }
+        AnalysisManager nestedAm = am.nest(root);
+        return OpToOpPassAdaptor::runPipeline(pipeline.getPasses(), root,
+                                              nestedAm);
+      });
 
   // Instrument before the pass has run.
   PassInstrumentor *pi = am.getPassInstrumentor();
   if (pi)
     pi->runBeforePass(pass, op);
 
   // Invoke the virtual runOnOperation method.
   pass->runOnOperation();
 
   // Invalidate any non preserved analyses.
   am.invalidate(pass->passState->preservedAnalyses);
 
   // Instrument after the pass has run.
   bool passFailed = pass->passState->irAndPassFailed.getInt();
   if (pi) {
     if (passFailed)
       pi->runAfterPassFailed(pass, op);
     else
       pi->runAfterPass(pass, op);
PassInstrumentor *AnalysisManager::getPassInstrumentor() const {
   ParentPointerT curParent = parent;
   while (auto *parentAM = curParent.dyn_cast<const AnalysisManager *>())
     curParent = parentAM->parent;
   return curParent.get<const ModuleAnalysisManager *>()->getPassInstrumentor();
 }
 
 /// Get an analysis manager for the given child operation.
 AnalysisManager AnalysisManager::nest(Operation *op) {
-  assert(op->getParentOp() == impl->getOperation() &&
-         "'op' has a different parent operation");
   auto it = impl->childAnalyses.find(op);
   if (it == impl->childAnalyses.end())
     it = impl->childAnalyses
              .try_emplace(op, std::make_unique<NestedAnalysisMap>(op))
              .first;
   return {this, it->second.get()};
 }
 

mlir/test/Pass/dynamic-pipeline-fail-on-parent.mlir

+// RUN: mlir-opt %s -pass-pipeline='module(test-dynamic-pipeline{op-name=inner_mod1 run-on-parent=1 dynamic-pipeline=test-patterns})'  -split-input-file -verify-diagnostics
+
+// Verify that we fail to schedule a dynamic pipeline on the parent operation.
+
+// expected-error @+1 {{'module' op Trying to schedule a dynamic pipeline on an operation that isn't nested under the current operation}}
+module {
+module @inner_mod1 {
+  "test.symbol"() {sym_name = "foo"} : () -> ()
+  func @bar()
+}
+}

mlir/test/Pass/dynamic-pipeline-nested.mlir

+// RUN: mlir-opt %s -pass-pipeline='module(test-dynamic-pipeline{op-name=inner_mod1 dynamic-pipeline=cse})'  --mlir-disable-threading  -print-ir-before-all 2>&1 | FileCheck %s --check-prefix=NOTNESTED --check-prefix=CHECK
+// RUN: mlir-opt %s -pass-pipeline='module(test-dynamic-pipeline{op-name=inner_mod1 run-on-nested-operations=1 dynamic-pipeline=cse})'  --mlir-disable-threading  -print-ir-before-all 2>&1 | FileCheck %s --check-prefix=NESTED --check-prefix=CHECK
+
+
+// Verify that we can schedule a dynamic pipeline on a nested operation
+
+func @f() {
+  return
+}
+
+// CHECK: IR Dump Before
+// CHECK-SAME: TestDynamicPipelinePass
+// CHECK-NEXT: module @inner_mod1
+module @inner_mod1 {
+// We use the print-ir-after-all dumps to check the granularity of the
+// scheduling: if we are nesting we expect to see to individual "Dump Before
+// CSE" output: one for each of the function. If we don't nest, then we expect
+// the CSE pass to run on the `inner_mod1` module directly.
+
+// CHECK: Dump Before CSE
+// NOTNESTED-NEXT: @inner_mod1
+// NESTED-NEXT: @foo
+  func @foo()
+// Only in the nested case we have a second run of the pass here.
+// NESTED: Dump Before CSE
+// NESTED-NEXT: @baz
+  func @baz()
+}

mlir/test/Pass/dynamic-pipeline.mlir

+// RUN: mlir-opt %s -pass-pipeline='module(test-dynamic-pipeline{op-name=inner_mod1, dynamic-pipeline=func(cse,canonicalize)})'  --mlir-disable-threading  -print-ir-before-all 2>&1 | FileCheck %s --check-prefix=MOD1 --check-prefix=MOD1-ONLY --check-prefix=CHECK
+// RUN: mlir-opt %s -pass-pipeline='module(test-dynamic-pipeline{op-name=inner_mod2, dynamic-pipeline=func(cse,canonicalize)})'  --mlir-disable-threading  -print-ir-before-all 2>&1 | FileCheck %s --check-prefix=MOD2 --check-prefix=MOD2-ONLY --check-prefix=CHECK
+// RUN: mlir-opt %s -pass-pipeline='module(test-dynamic-pipeline{op-name=inner_mod1,inner_mod2, dynamic-pipeline=func(cse,canonicalize)})'  --mlir-disable-threading  -print-ir-before-all 2>&1 | FileCheck %s --check-prefix=MOD1 --check-prefix=MOD2 --check-prefix=CHECK
+// RUN: mlir-opt %s -pass-pipeline='module(test-dynamic-pipeline{dynamic-pipeline=func(cse,canonicalize)})'  --mlir-disable-threading  -print-ir-before-all 2>&1 | FileCheck %s --check-prefix=MOD1 --check-prefix=MOD2 --check-prefix=CHECK
+
+
+func @f() {
+  return
+}
+
+// CHECK: IR Dump Before
+// CHECK-SAME: TestDynamicPipelinePass
+// CHECK-NEXT: module @inner_mod1
+// MOD2-ONLY: dynamic-pipeline skip op name: inner_mod1
+module @inner_mod1 {
+// MOD1: Dump Before CSE
+// MOD1-NEXT: @foo
+// MOD1: Dump Before Canonicalizer
+// MOD1-NEXT: @foo
+  func @foo() {
+    return
+  }
+// MOD1: Dump Before CSE
+// MOD1-NEXT: @baz
+// MOD1: Dump Before Canonicalizer
+// MOD1-NEXT: @baz
+  func @baz() {
+    return
+  }
+}
+
+// CHECK: IR Dump Before
+// CHECK-SAME: TestDynamicPipelinePass
+// CHECK-NEXT: module @inner_mod2
+// MOD1-ONLY: dynamic-pipeline skip op name: inner_mod2
+module @inner_mod2 {
+// MOD2: Dump Before CSE
+// MOD2-NEXT: @foo
+// MOD2: Dump Before Canonicalizer
+// MOD2-NEXT: @foo
+  func @foo() {
+    return
+  }
+}

mlir/test/lib/Transforms/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestTransforms
   TestAllReduceLowering.cpp
   TestBufferPlacement.cpp
   TestExpandTanh.cpp
   TestCallGraph.cpp
   TestConstantFold.cpp
   TestConvertCallOp.cpp
   TestConvertGPUKernelToCubin.cpp
   TestConvertGPUKernelToHsaco.cpp
   TestDominance.cpp
+  TestDynamicPipeline.cpp
   TestLoopFusion.cpp
   TestGpuMemoryPromotion.cpp
   TestGpuParallelLoopMapping.cpp
   TestInlining.cpp
   TestLinalgHoisting.cpp
   TestLinalgTransforms.cpp
   TestLiveness.cpp
   TestLoopMapping.cpp

mlir/test/lib/Transforms/TestDynamicPipeline.cpp

+//===------ TestDynamicPipeline.cpp --- dynamic pipeline test pass --------===//
+//
+// 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 a pass to test the dynamic pipeline feature.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/LoopUtils.h"
+#include "mlir/Transforms/Passes.h"
+
+using namespace mlir;
+
+namespace {
+
+class TestDynamicPipelinePass
+    : public PassWrapper<TestDynamicPipelinePass, OperationPass<>> {
+public:
+  void getDependentDialects(DialectRegistry &registry) const override {
+    MLIRContext ctx;
+    OpPassManager pm(Identifier::get(ModuleOp::getOperationName(), &ctx),
+                     false);
+    parsePassPipeline(pipeline, pm, llvm::errs());
+    pm.getDependentDialects(registry);
+  }
+
+  TestDynamicPipelinePass(){};
+  TestDynamicPipelinePass(const TestDynamicPipelinePass &) {}
+
+  void runOnOperation() override {
+    llvm::errs() << "Dynamic execute '" << pipeline << "' on "
+                 << getOperation()->getName() << "\n";
+    if (pipeline.empty()) {
+      llvm::errs() << "Empty pipeline\n";
+      return;
+    }
+    auto symbolOp = dyn_cast<SymbolOpInterface>(getOperation());
+    if (!symbolOp) {
+      getOperation()->emitWarning()
+          << "Ignoring because not implementing SymbolOpInterface\n";
+      return;
+    }
+
+    auto opName = symbolOp.getName();
+    if (!opNames.empty() && !llvm::is_contained(opNames, opName)) {
+      llvm::errs() << "dynamic-pipeline skip op name: " << opName << "\n";
+      return;
+    }
+    if (!pm) {
+      pm = std::make_unique<OpPassManager>(getOperation()->getName(), false);
+      parsePassPipeline(pipeline, *pm, llvm::errs());
+    }
+    // Check that running on the parent operation always immediately fails.
+    if (runOnParent) {
+      if (getOperation()->getParentOp())
+        if (!failed(runPipeline(*pm, getOperation()->getParentOp())))
+          signalPassFailure();
+      return;
+    }
+
+    if (runOnNestedOp) {
+      llvm::errs() << "Run on nested op\n";
+      getOperation()->walk([&](Operation *op) {
+        if (op == getOperation())
+          return;
+        llvm::errs() << "Run on " << *op << "\n";
+        // Run on the current operation
+        if (failed(runPipeline(*pm, op)))
+          signalPassFailure();
+      });
+    } else {
+      // Run on the current operation
+      if (failed(runPipeline(*pm, getOperation())))
+        signalPassFailure();
+    }
+  }
+
+  std::unique_ptr<OpPassManager> pm;
+
+  Option<bool> runOnNestedOp{
+      *this, "run-on-nested-operations",
+      llvm::cl::desc("This will apply the pipeline on nested operations under "
+                     "the visited operation.")};
+  Option<bool> runOnParent{
+      *this, "run-on-parent",
+      llvm::cl::desc("This will apply the pipeline on the parent operation if "
+                     "it exist, this is expected to fail.")};
+  Option<std::string> pipeline{
+      *this, "dynamic-pipeline",
+      llvm::cl::desc("The pipeline description that "
+                     "will run on the filtered function.")};
+  ListOption<std::string> opNames{
+      *this, "op-name", llvm::cl::MiscFlags::CommaSeparated,
+      llvm::cl::desc("List of function name to apply the pipeline to")};
+};
+} // end namespace
+
+namespace mlir {
+void registerTestDynamicPipelinePass() {
+  PassRegistration<TestDynamicPipelinePass>(
+      "test-dynamic-pipeline", "Tests the dynamic pipeline feature by applying "
+                               "a pipeline on a selected set of functions");
+}
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestAllReduceLoweringPass();
 void registerTestBufferPlacementPreparationPass();
 void registerTestCallGraphPass();
 void registerTestConstantFold();
 void registerTestConvertGPUKernelToCubinPass();
 void registerTestConvertGPUKernelToHsacoPass();
 void registerTestDominancePass();
 void registerTestDialect(DialectRegistry &);
+void registerTestDynamicPipelinePass();
 void registerTestExpandTanhPass();
 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
 void registerTestGpuParallelLoopMappingPass();
 void registerTestInterfaces();
 void registerTestLinalgHoisting();
 void registerTestLinalgTransforms();
 void registerTestLivenessPass();

 #if MLIR_CUDA_CONVERSIONS_ENABLED
   registerTestConvertGPUKernelToCubinPass();
 #endif
 #if MLIR_ROCM_CONVERSIONS_ENABLED
   registerTestConvertGPUKernelToHsacoPass();
 #endif
   registerTestBufferPlacementPreparationPass();
   registerTestDominancePass();
+  registerTestDynamicPipelinePass();
   registerTestFunc();
   registerTestExpandTanhPass();
   registerTestGpuMemoryPromotionPass();
   registerTestInterfaces();
   registerTestLinalgHoisting();
   registerTestLinalgTransforms();
   registerTestLivenessPass();
   registerTestLoopFusion();