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
This commit is contained in:
parent
b289dc5306
commit
fb1de7ed92
|
@ -24,8 +24,11 @@ class OpToOpPassAdaptor;
|
||||||
/// The state for a single execution of a pass. This provides a unified
|
/// The state for a single execution of a pass. This provides a unified
|
||||||
/// interface for accessing and initializing necessary state for pass execution.
|
/// interface for accessing and initializing necessary state for pass execution.
|
||||||
struct PassExecutionState {
|
struct PassExecutionState {
|
||||||
PassExecutionState(Operation *ir, AnalysisManager analysisManager)
|
PassExecutionState(Operation *ir, AnalysisManager analysisManager,
|
||||||
: irAndPassFailed(ir, false), 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
|
/// The current operation being transformed and a bool for if the pass
|
||||||
/// signaled a failure.
|
/// signaled a failure.
|
||||||
|
@ -36,6 +39,10 @@ struct PassExecutionState {
|
||||||
|
|
||||||
/// The set of preserved analyses for the current execution.
|
/// The set of preserved analyses for the current execution.
|
||||||
detail::PreservedAnalyses preservedAnalyses;
|
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
|
} // namespace detail
|
||||||
|
|
||||||
|
@ -156,6 +163,13 @@ protected:
|
||||||
/// The polymorphic API that runs the pass over the currently held operation.
|
/// The polymorphic API that runs the pass over the currently held operation.
|
||||||
virtual void runOnOperation() = 0;
|
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.
|
/// A clone method to create a copy of this pass.
|
||||||
std::unique_ptr<Pass> clone() const {
|
std::unique_ptr<Pass> clone() const {
|
||||||
auto newInst = clonePass();
|
auto newInst = clonePass();
|
||||||
|
|
|
@ -36,6 +36,7 @@ class PassInstrumentor;
|
||||||
|
|
||||||
namespace detail {
|
namespace detail {
|
||||||
struct OpPassManagerImpl;
|
struct OpPassManagerImpl;
|
||||||
|
struct PassExecutionState;
|
||||||
} // end namespace detail
|
} // end namespace detail
|
||||||
|
|
||||||
//===----------------------------------------------------------------------===//
|
//===----------------------------------------------------------------------===//
|
||||||
|
@ -119,6 +120,7 @@ private:
|
||||||
|
|
||||||
/// Allow access to the constructor.
|
/// Allow access to the constructor.
|
||||||
friend class PassManager;
|
friend class PassManager;
|
||||||
|
friend class Pass;
|
||||||
|
|
||||||
/// Allow access.
|
/// Allow access.
|
||||||
friend detail::OpPassManagerImpl;
|
friend detail::OpPassManagerImpl;
|
||||||
|
|
|
@ -357,8 +357,22 @@ LogicalResult OpToOpPassAdaptor::run(Pass *pass, Operation *op,
|
||||||
return op->emitOpError() << "trying to schedule a pass on an operation not "
|
return op->emitOpError() << "trying to schedule a pass on an operation not "
|
||||||
"marked as 'IsolatedFromAbove'";
|
"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.
|
// Instrument before the pass has run.
|
||||||
PassInstrumentor *pi = am.getPassInstrumentor();
|
PassInstrumentor *pi = am.getPassInstrumentor();
|
||||||
if (pi)
|
if (pi)
|
||||||
|
@ -839,8 +853,6 @@ PassInstrumentor *AnalysisManager::getPassInstrumentor() const {
|
||||||
|
|
||||||
/// Get an analysis manager for the given child operation.
|
/// Get an analysis manager for the given child operation.
|
||||||
AnalysisManager AnalysisManager::nest(Operation *op) {
|
AnalysisManager AnalysisManager::nest(Operation *op) {
|
||||||
assert(op->getParentOp() == impl->getOperation() &&
|
|
||||||
"'op' has a different parent operation");
|
|
||||||
auto it = impl->childAnalyses.find(op);
|
auto it = impl->childAnalyses.find(op);
|
||||||
if (it == impl->childAnalyses.end())
|
if (it == impl->childAnalyses.end())
|
||||||
it = impl->childAnalyses
|
it = impl->childAnalyses
|
||||||
|
|
11
mlir/test/Pass/dynamic-pipeline-fail-on-parent.mlir
Normal file
11
mlir/test/Pass/dynamic-pipeline-fail-on-parent.mlir
Normal file
|
@ -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()
|
||||||
|
}
|
||||||
|
}
|
28
mlir/test/Pass/dynamic-pipeline-nested.mlir
Normal file
28
mlir/test/Pass/dynamic-pipeline-nested.mlir
Normal file
|
@ -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()
|
||||||
|
}
|
44
mlir/test/Pass/dynamic-pipeline.mlir
Normal file
44
mlir/test/Pass/dynamic-pipeline.mlir
Normal file
|
@ -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
|
||||||
|
}
|
||||||
|
}
|
|
@ -11,6 +11,7 @@ add_mlir_library(MLIRTestTransforms
|
||||||
TestConvertGPUKernelToCubin.cpp
|
TestConvertGPUKernelToCubin.cpp
|
||||||
TestConvertGPUKernelToHsaco.cpp
|
TestConvertGPUKernelToHsaco.cpp
|
||||||
TestDominance.cpp
|
TestDominance.cpp
|
||||||
|
TestDynamicPipeline.cpp
|
||||||
TestLoopFusion.cpp
|
TestLoopFusion.cpp
|
||||||
TestGpuMemoryPromotion.cpp
|
TestGpuMemoryPromotion.cpp
|
||||||
TestGpuParallelLoopMapping.cpp
|
TestGpuParallelLoopMapping.cpp
|
||||||
|
|
112
mlir/test/lib/Transforms/TestDynamicPipeline.cpp
Normal file
112
mlir/test/lib/Transforms/TestDynamicPipeline.cpp
Normal file
|
@ -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 ®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
|
|
@ -52,6 +52,7 @@ void registerTestConvertGPUKernelToCubinPass();
|
||||||
void registerTestConvertGPUKernelToHsacoPass();
|
void registerTestConvertGPUKernelToHsacoPass();
|
||||||
void registerTestDominancePass();
|
void registerTestDominancePass();
|
||||||
void registerTestDialect(DialectRegistry &);
|
void registerTestDialect(DialectRegistry &);
|
||||||
|
void registerTestDynamicPipelinePass();
|
||||||
void registerTestExpandTanhPass();
|
void registerTestExpandTanhPass();
|
||||||
void registerTestFunc();
|
void registerTestFunc();
|
||||||
void registerTestGpuMemoryPromotionPass();
|
void registerTestGpuMemoryPromotionPass();
|
||||||
|
@ -108,6 +109,7 @@ void registerTestPasses() {
|
||||||
registerTestAffineLoopParametricTilingPass();
|
registerTestAffineLoopParametricTilingPass();
|
||||||
registerTestBufferPlacementPreparationPass();
|
registerTestBufferPlacementPreparationPass();
|
||||||
registerTestDominancePass();
|
registerTestDominancePass();
|
||||||
|
registerTestDynamicPipelinePass();
|
||||||
registerTestFunc();
|
registerTestFunc();
|
||||||
registerTestExpandTanhPass();
|
registerTestExpandTanhPass();
|
||||||
registerTestGpuMemoryPromotionPass();
|
registerTestGpuMemoryPromotionPass();
|
||||||
|
|
Loading…
Reference in a new issue