/// 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.
/// 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 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();
namespace detail {
struct OpPassManagerImpl;
+struct PassExecutionState;
} // end namespace detail
//===----------------------------------------------------------------------===//
/// Allow access to the constructor.
friend class PassManager;
+ friend class Pass;
/// Allow access.
friend detail::OpPassManagerImpl;
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)
/// 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
--- /dev/null
+// 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()
+}
+}
--- /dev/null
+// 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()
+}
--- /dev/null
+// 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
+ }
+}
TestConvertGPUKernelToCubin.cpp
TestConvertGPUKernelToHsaco.cpp
TestDominance.cpp
+ TestDynamicPipeline.cpp
TestLoopFusion.cpp
TestGpuMemoryPromotion.cpp
TestGpuParallelLoopMapping.cpp
--- /dev/null
+//===------ 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 ®istry) 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
void registerTestConvertGPUKernelToHsacoPass();
void registerTestDominancePass();
void registerTestDialect(DialectRegistry &);
+void registerTestDynamicPipelinePass();
void registerTestExpandTanhPass();
void registerTestFunc();
void registerTestGpuMemoryPromotionPass();
registerTestAffineLoopParametricTilingPass();
registerTestBufferPlacementPreparationPass();
registerTestDominancePass();
+ registerTestDynamicPipelinePass();
registerTestFunc();
registerTestExpandTanhPass();
registerTestGpuMemoryPromotionPass();