[Mlir-commits] [mlir] fb1de7e - Implement a new kind of Pass: dynamic pass pipeline
Mehdi Amini
llvmlistbot at llvm.org
Tue Sep 22 11:52:07 PDT 2020
Author: Mehdi Amini
Date: 2020-09-22T18:51:54Z
New Revision: fb1de7ed921c3c1531c0cb28b9603af22ecd0529
URL: https://github.com/llvm/llvm-project/commit/fb1de7ed921c3c1531c0cb28b9603af22ecd0529
DIFF: https://github.com/llvm/llvm-project/commit/fb1de7ed921c3c1531c0cb28b9603af22ecd0529.diff
LOG: 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
Added:
mlir/test/Pass/dynamic-pipeline-fail-on-parent.mlir
mlir/test/Pass/dynamic-pipeline-nested.mlir
mlir/test/Pass/dynamic-pipeline.mlir
mlir/test/lib/Transforms/TestDynamicPipeline.cpp
Modified:
mlir/include/mlir/Pass/Pass.h
mlir/include/mlir/Pass/PassManager.h
mlir/lib/Pass/Pass.cpp
mlir/test/lib/Transforms/CMakeLists.txt
mlir/tools/mlir-opt/mlir-opt.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Pass/Pass.h b/mlir/include/mlir/Pass/Pass.h
index 526669fae363..e21e3e750270 100644
--- a/mlir/include/mlir/Pass/Pass.h
+++ b/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 @@ class Pass {
/// 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();
diff --git a/mlir/include/mlir/Pass/PassManager.h b/mlir/include/mlir/Pass/PassManager.h
index 9aace79f2053..9b9214fd16c7 100644
--- a/mlir/include/mlir/Pass/PassManager.h
+++ b/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 @@ class OpPassManager {
/// Allow access to the constructor.
friend class PassManager;
+ friend class Pass;
/// Allow access.
friend detail::OpPassManagerImpl;
diff --git a/mlir/lib/Pass/Pass.cpp b/mlir/lib/Pass/Pass.cpp
index 2b49d9309322..7098b4bd4388 100644
--- a/mlir/lib/Pass/Pass.cpp
+++ b/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
diff erent parent operation");
auto it = impl->childAnalyses.find(op);
if (it == impl->childAnalyses.end())
it = impl->childAnalyses
diff --git a/mlir/test/Pass/dynamic-pipeline-fail-on-parent.mlir b/mlir/test/Pass/dynamic-pipeline-fail-on-parent.mlir
new file mode 100644
index 000000000000..8885dab80538
--- /dev/null
+++ b/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()
+}
+}
diff --git a/mlir/test/Pass/dynamic-pipeline-nested.mlir b/mlir/test/Pass/dynamic-pipeline-nested.mlir
new file mode 100644
index 000000000000..7651ff403348
--- /dev/null
+++ b/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()
+}
diff --git a/mlir/test/Pass/dynamic-pipeline.mlir b/mlir/test/Pass/dynamic-pipeline.mlir
new file mode 100644
index 000000000000..6a84ccc5688c
--- /dev/null
+++ b/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
+ }
+}
diff --git a/mlir/test/lib/Transforms/CMakeLists.txt b/mlir/test/lib/Transforms/CMakeLists.txt
index 99424f1c9c06..3c82554fa13a 100644
--- a/mlir/test/lib/Transforms/CMakeLists.txt
+++ b/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
diff --git a/mlir/test/lib/Transforms/TestDynamicPipeline.cpp b/mlir/test/lib/Transforms/TestDynamicPipeline.cpp
new file mode 100644
index 000000000000..92e8861b5ad0
--- /dev/null
+++ b/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 ®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
diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
index 93934d40fe59..aed8b0ae818b 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/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();
More information about the Mlir-commits
mailing list