Implement a new kind of Pass: dynamic pass pipeline

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

Recommit after fixing an ASAN issue: the callback lambda needs to be
allocated to a temporary to have its lifetime extended to the end of the
current block instead of just the current call expression.

Reviewed By: silvas

Differential Revision: https://reviews.llvm.org/D86392

mlir/include/mlir/Pass/Pass.h

@@ -24,8 +24,11 @@ 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.
@@ -36,6 +39,10 @@ struct PassExecutionState {
 
   /// 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
 
@@ -156,6 +163,13 @@ protected:
   /// 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();

mlir/include/mlir/Pass/PassManager.h

@@ -36,6 +36,7 @@ class PassInstrumentor;
 
 namespace detail {
 struct OpPassManagerImpl;
+struct PassExecutionState;
 } // end namespace detail
 
 //===----------------------------------------------------------------------===//
@@ -119,6 +120,7 @@ private:
 
   /// Allow access to the constructor.
   friend class PassManager;
+  friend class Pass;
 
   /// Allow access.
   friend detail::OpPassManagerImpl;

mlir/lib/Pass/Pass.cpp

@@ -357,8 +357,22 @@ LogicalResult OpToOpPassAdaptor::run(Pass *pass, Operation *op,
     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.
+  auto dynamic_pipeline_callback =
+      [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 processing";
+          return failure();
+        }
+        AnalysisManager nestedAm = am.nest(root);
+        return OpToOpPassAdaptor::runPipeline(pipeline.getPasses(), root,
+                                              nestedAm);
+      };
+  pass->passState.emplace(op, am, dynamic_pipeline_callback);
   // Instrument before the pass has run.
   PassInstrumentor *pi = am.getPassInstrumentor();
   if (pi)
@@ -839,8 +853,6 @@ PassInstrumentor *AnalysisManager::getPassInstrumentor() const {
 
 /// 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

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

@@ -0,0 +1,11 @@
+// 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

@@ -0,0 +1,28 @@
+// 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

@@ -0,0 +1,44 @@
+// 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

@@ -11,6 +11,7 @@ add_mlir_library(MLIRTestTransforms
   TestConvertGPUKernelToCubin.cpp
   TestConvertGPUKernelToHsaco.cpp
   TestDominance.cpp
+  TestDynamicPipeline.cpp
   TestLoopFusion.cpp
   TestGpuMemoryPromotion.cpp
   TestGpuParallelLoopMapping.cpp

mlir/test/lib/Transforms/TestDynamicPipeline.cpp

@@ -0,0 +1,112 @@
+//===------ 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 {
+    OpPassManager pm(ModuleOp::getOperationName(), 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().getIdentifier(), 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() || !op->isKnownIsolatedFromAbove())
+          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

@@ -52,6 +52,7 @@ void registerTestConvertGPUKernelToCubinPass();
 void registerTestConvertGPUKernelToHsacoPass();
 void registerTestDominancePass();
 void registerTestDialect(DialectRegistry &);
+void registerTestDynamicPipelinePass();
 void registerTestExpandTanhPass();
 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
@@ -108,6 +109,7 @@ void registerTestPasses() {
   registerTestAffineLoopParametricTilingPass();
   registerTestBufferPlacementPreparationPass();
   registerTestDominancePass();
+  registerTestDynamicPipelinePass();
   registerTestFunc();
   registerTestExpandTanhPass();
   registerTestGpuMemoryPromotionPass();