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 @@ -19,6 +19,8 @@ namespace mlir { namespace detail { +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 { @@ -154,6 +156,17 @@ /// The polymorphic API that runs the pass over the currently held operation. virtual void runOnOperation() = 0; + /// This is the equivalent to `runOnOperation` for a DynamicPass: it is + /// expected to return an OpPassManager containing the dynamic pipeline to + /// execute. + virtual OpPassManager *buildDynamicPipeline() { + llvm_unreachable("calling buildDynamicPipeline on a non-dynamic pass"); + } + + /// Return true if this is a dynamic pass: intended to be overriden only in + /// the DynamicPassWrapper class. + virtual bool isDynamic() const { return false; } + /// A clone method to create a copy of this pass. std::unique_ptr clone() const { auto newInst = clonePass(); @@ -249,9 +262,6 @@ void copyOptionValuesFrom(const Pass *other); private: - /// Forwarding function to execute this pass on the given operation. - LLVM_NODISCARD - LogicalResult run(Operation *op, AnalysisManager am); /// Out of line virtual method to ensure vtables and metadata are emitted to a /// single .o file. @@ -273,11 +283,11 @@ /// The pass options registered to this pass instance. detail::PassOptions passOptions; - /// Allow access to 'clone' and 'run'. + /// Allow access to 'clone'. friend class OpPassManager; - /// Allow access to 'run'. - friend class PassManager; + /// Allow access to 'passState'. + friend detail::OpToOpPassAdaptor; /// Allow access to 'passOptions'. friend class PassInfo; @@ -380,6 +390,49 @@ std::unique_ptr clonePass() const override { return std::make_unique(*static_cast(this)); } + +private: + /// Hide from derived class, this is only available to dynamic passes. + OpPassManager *buildDynamicPipeline() final { + llvm_unreachable("calling buildDynamicPipeline on a non-dynamic pass"); + } + /// Always return false: use DynamicPassWrapper for dynamic passes. + bool isDynamic() const final { return false; } +}; + +/// This class provides a CRTP wrapper around a base dynamic pass class to +/// define several necessary utility methods. +template +class DynamicPassWrapper : public BaseT { +public: + /// Support isa/dyn_cast functionality for the derived pass class. + static bool classof(const Pass *pass) { + return pass->getTypeID() == TypeID::get(); + } + +protected: + DynamicPassWrapper() : BaseT(TypeID::get()) {} + + /// Returns the derived pass name. + StringRef getName() const override { return llvm::getTypeName(); } + + /// A clone method to create a copy of this pass. + std::unique_ptr clonePass() const override { + return std::make_unique(*static_cast(this)); + } + + /// Build and return a dynamic pipeline. + virtual OpPassManager *buildDynamicPipeline() = 0; + +private: + /// Hide this method from dynamic passes. + void runOnOperation() final { + llvm::report_fatal_error( + "Unexpected use of runOnOperation in a dynamic pass"); + } + + /// Always returns true to indicate that it is a dynamic pass. + bool isDynamic() const final { return true; } }; } // end namespace mlir 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 //===----------------------------------------------------------------------===// @@ -47,6 +48,7 @@ /// other OpPassManagers or the top-level PassManager. class OpPassManager { public: + OpPassManager(OperationName name, bool verifyPasses); OpPassManager(OpPassManager &&rhs); OpPassManager(const OpPassManager &rhs); ~OpPassManager(); @@ -54,22 +56,19 @@ /// Iterator over the passes in this pass manager. using pass_iterator = - llvm::pointee_iterator>::iterator>; + llvm::pointee_iterator>::iterator>; pass_iterator begin(); pass_iterator end(); iterator_range getPasses() { return {begin(), end()}; } - using const_pass_iterator = llvm::pointee_iterator< - std::vector>::const_iterator>; + using const_pass_iterator = + llvm::pointee_iterator>::const_iterator>; const_pass_iterator begin() const; const_pass_iterator end() const; iterator_range getPasses() const { return {begin(), end()}; } - /// Run the held passes over the given operation. - LogicalResult run(Operation *op, AnalysisManager am); - /// Nest a new operation pass manager for the given operation kind under this /// pass manager. OpPassManager &nest(const OperationName &nestedName); @@ -115,13 +114,12 @@ void getDependentDialects(DialectRegistry &dialects) const; private: - OpPassManager(OperationName name, bool verifyPasses); - /// A pointer to an internal implementation instance. std::unique_ptr impl; /// 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 @@ -18,6 +18,7 @@ #include "mlir/IR/Verifier.h" #include "mlir/Support/FileUtilities.h" #include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/ScopeExit.h" #include "llvm/ADT/SetVector.h" #include "llvm/Support/CommandLine.h" #include "llvm/Support/CrashRecoveryContext.h" @@ -73,33 +74,6 @@ passOptions.print(os); } -/// Forwarding function to execute this pass. -LogicalResult Pass::run(Operation *op, AnalysisManager am) { - passState.emplace(op, am); - - // Instrument before the pass has run. - auto pi = am.getPassInstrumentor(); - if (pi) - pi->runBeforePass(this, op); - - // Invoke the virtual runOnOperation method. - runOnOperation(); - - // Invalidate any non preserved analyses. - am.invalidate(passState->preservedAnalyses); - - // Instrument after the pass has run. - bool passFailed = passState->irAndPassFailed.getInt(); - if (pi) { - if (passFailed) - pi->runAfterPassFailed(this, op); - else - pi->runAfterPass(this, op); - } - - // Return if the pass signaled a failure. - return failure(passFailed); -} //===----------------------------------------------------------------------===// // Verifier Passes @@ -286,24 +260,17 @@ OpPassManager::~OpPassManager() {} OpPassManager::pass_iterator OpPassManager::begin() { - return impl->passes.begin(); + return MutableArrayRef>{impl->passes}.begin(); +} +OpPassManager::pass_iterator OpPassManager::end() { + return MutableArrayRef>{impl->passes}.end(); } -OpPassManager::pass_iterator OpPassManager::end() { return impl->passes.end(); } OpPassManager::const_pass_iterator OpPassManager::begin() const { - return impl->passes.begin(); + return ArrayRef>{impl->passes}.begin(); } OpPassManager::const_pass_iterator OpPassManager::end() const { - return impl->passes.end(); -} - -/// Run all of the passes in this manager over the current operation. -LogicalResult OpPassManager::run(Operation *op, AnalysisManager am) { - // Run each of the held passes. - for (auto &pass : impl->passes) - if (failed(pass->run(op, am))) - return failure(); - return success(); + return ArrayRef>{impl->passes}.end(); } /// Nest a new operation pass manager for the given operation kind under this @@ -367,19 +334,59 @@ // OpToOpPassAdaptor //===----------------------------------------------------------------------===// -/// Utility to run the given operation and analysis manager on a provided op -/// pass manager. -static LogicalResult runPipeline(OpPassManager &pm, Operation *op, - AnalysisManager am) { +LogicalResult OpToOpPassAdaptor::run(Pass *pass, Operation *op, + AnalysisManager am) { + pass->passState.emplace(op, am); + + // Instrument before the pass has run. + PassInstrumentor *pi = am.getPassInstrumentor(); + if (pi) + pi->runBeforePass(pass, op); + + // Invoke the virtual either runOnOperation method or if a dynamic pass, + // retried the OpPassManager and execute it. + if (pass->isDynamic()) { + OpPassManager *dynamicPipeline = pass->buildDynamicPipeline(); + if (dynamicPipeline) + OpToOpPassAdaptor::runPipeline(dynamicPipeline->getPasses(), op, am); + } else { + pass->runOnOperation(); + } + + // Invalidate any non preserved analyses. + am.invalidate(pass->passState->preservedAnalyses); + + // Instrument after the pass has run. + bool passFailed = pass->passState->irAndPassFailed.getInt(); + if (pi) { + if (passFailed) + pi->runAfterPassFailed(pass, op); + else + pi->runAfterPass(pass, op); + } + + // Return if the pass signaled a failure. + return failure(passFailed); +} + +/// Run the given operation and analysis manager on a provided op pass manager. +LogicalResult OpToOpPassAdaptor::runPipeline( + iterator_range passes, Operation *op, + AnalysisManager am) { + auto scope_exit = llvm::make_scope_exit([&] { + // Clear out any computed operation analyses. These analyses won't be used + // any more in this pipeline, and this helps reduce the current working set + // of memory. If preserving these analyses becomes important in the future + // we can re-evaluate this. + am.clear (); + }); + // Run the pipeline over the provided operation. - auto result = pm.run(op, am); + for (Pass &pass : passes) + if (failed(run(&pass, op, am))) + return failure(); - // Clear out any computed operation analyses. These analyses won't be used - // any more in this pipeline, and this helps reduce the current working set - // of memory. If preserving these analyses becomes important in the future - // we can re-evaluate this. - am.clear(); - return result; + return success(); } /// Find an operation pass manager that can operate on an operation of the given @@ -457,7 +464,7 @@ // Run the held pipeline over the current operation. if (instrumentor) instrumentor->runBeforePipeline(mgr->getOpName(), parentInfo); - auto result = runPipeline(*mgr, &op, am.slice(&op)); + auto result = runPipeline(mgr->getPasses(), &op, am.slice(&op)); if (instrumentor) instrumentor->runAfterPipeline(mgr->getOpName(), parentInfo); @@ -536,7 +543,8 @@ if (instrumentor) instrumentor->runBeforePipeline(pm->getOpName(), parentInfo); - auto pipelineResult = runPipeline(*pm, it.first, it.second); + auto pipelineResult = + runPipeline(pm->getPasses(), it.first, it.second); if (instrumentor) instrumentor->runAfterPipeline(pm->getOpName(), parentInfo); @@ -709,7 +717,7 @@ llvm::CrashRecoveryContext recoveryContext; recoveryContext.RunSafelyOnThread([&] { for (std::unique_ptr &pass : passes) - if (failed(pass->run(module, am))) + if (failed(OpToOpPassAdaptor::run(pass.get(), module, am))) return; passManagerResult = success(); }); @@ -753,9 +761,10 @@ // If reproducer generation is enabled, run the pass manager with crash // handling enabled. - LogicalResult result = crashReproducerFileName - ? runWithCrashRecovery(module, am) - : OpPassManager::run(module, am); + LogicalResult result = + crashReproducerFileName + ? runWithCrashRecovery(module, am) + : OpToOpPassAdaptor::runPipeline(getPasses(), module, am); // Dump all of the pass statistics if necessary. if (passStatisticsMode) diff --git a/mlir/lib/Pass/PassDetail.h b/mlir/lib/Pass/PassDetail.h --- a/mlir/lib/Pass/PassDetail.h +++ b/mlir/lib/Pass/PassDetail.h @@ -62,12 +62,24 @@ /// Run this pass adaptor asynchronously. void runOnOperationAsyncImpl(); + /// Run the given operation and analysis manager on a single pass. + static LogicalResult run(Pass *pass, Operation *op, AnalysisManager am); + + /// Run the given operation and analysis manager on a provided op pass + /// manager. + static LogicalResult + runPipeline(iterator_range passes, + Operation *op, AnalysisManager am); + /// A set of adaptors to run. SmallVector mgrs; /// A set of executors, cloned from the main executor, that run asynchronously /// on different threads. This is used when threading is enabled. SmallVector, 8> asyncExecutors; + + // For accessing "runPipeline". + friend class mlir::PassManager; }; } // end namespace detail 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/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,72 @@ +//===------ 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 DynamicPassWrapper> { +public: + TestDynamicPipelinePass(){}; + TestDynamicPipelinePass(const TestDynamicPipelinePass &) {} + + OpPassManager *buildDynamicPipeline() override { + llvm::errs() << "Dynamic execute '" << pipeline << "' on " + << getOperation()->getName() << "\n"; + if (pipeline.empty()) + return nullptr; + auto symbolOp = dyn_cast(getOperation()); + if (!symbolOp) { + getOperation()->emitWarning() + << "Ignoring because not implementing SymbolOpInterface\n"; + return nullptr; + } + + auto opName = symbolOp.getName(); + if (!opNames.empty() && !llvm::is_contained(opNames, opName)) { + llvm::errs() << "dynamic-pipeline skip op name: " << opName << "\n"; + return nullptr; + } + if (!pm) { + pm = std::make_unique(getOperation()->getName(), false); + parsePassPipeline(pipeline, *pm, llvm::errs()); + } + return pm.get(); + } + + std::unique_ptr pm; + + 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 @@ -49,6 +49,7 @@ void registerTestConvertGPUKernelToHsacoPass(); void registerTestDominancePass(); void registerTestDialect(DialectRegistry &); +void registerTestDynamicPipelinePass(); void registerTestExpandTanhPass(); void registerTestFunc(); void registerTestGpuMemoryPromotionPass(); @@ -100,6 +101,7 @@ #endif registerTestBufferPlacementPreparationPass(); registerTestDominancePass(); + registerTestDynamicPipelinePass(); registerTestFunc(); registerTestExpandTanhPass(); registerTestGpuMemoryPromotionPass();