diff --git a/mlir/include/mlir/Pass/Pass.h b/mlir/include/mlir/Pass/Pass.h --- a/mlir/include/mlir/Pass/Pass.h +++ b/mlir/include/mlir/Pass/Pass.h @@ -24,8 +24,11 @@ /// 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 + 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 @@ /// 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 pipelineExecutor; }; } // namespace detail @@ -156,6 +163,13 @@ /// 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 clone() const { auto newInst = clonePass(); diff --git a/mlir/include/mlir/Pass/PassManager.h b/mlir/include/mlir/Pass/PassManager.h --- a/mlir/include/mlir/Pass/PassManager.h +++ b/mlir/include/mlir/Pass/PassManager.h @@ -36,6 +36,7 @@ namespace detail { struct OpPassManagerImpl; +struct PassExecutionState; } // end namespace detail //===----------------------------------------------------------------------===// @@ -119,6 +120,7 @@ /// 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 --- a/mlir/lib/Pass/Pass.cpp +++ b/mlir/lib/Pass/Pass.cpp @@ -357,8 +357,21 @@ 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) @@ -839,8 +852,6 @@ /// 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 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 --- /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 --- /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 --- /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 --- a/mlir/test/lib/Transforms/CMakeLists.txt +++ b/mlir/test/lib/Transforms/CMakeLists.txt @@ -11,6 +11,7 @@ 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 --- /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> { +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(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( + 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 pm; + + Option runOnNestedOp{ + *this, "run-on-nested-operations", + llvm::cl::desc("This will apply the pipeline on nested operations under " + "the visited operation.")}; + Option 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 pipeline{ + *this, "dynamic-pipeline", + llvm::cl::desc("The pipeline description that " + "will run on the filtered function.")}; + ListOption 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( + "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 --- a/mlir/tools/mlir-opt/mlir-opt.cpp +++ b/mlir/tools/mlir-opt/mlir-opt.cpp @@ -52,6 +52,7 @@ void registerTestConvertGPUKernelToHsacoPass(); void registerTestDominancePass(); void registerTestDialect(DialectRegistry &); +void registerTestDynamicPipelinePass(); void registerTestExpandTanhPass(); void registerTestFunc(); void registerTestGpuMemoryPromotionPass(); @@ -108,6 +109,7 @@ registerTestAffineLoopParametricTilingPass(); registerTestBufferPlacementPreparationPass(); registerTestDominancePass(); + registerTestDynamicPipelinePass(); registerTestFunc(); registerTestExpandTanhPass(); registerTestGpuMemoryPromotionPass();