[Mlir-commits] [mlir] [mlir][inliner] Refactor MLIR inliner pass and utils. (PR #84059)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Mar 5 11:33:37 PST 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-affine
Author: Slava Zakharin (vzakhari)
<details>
<summary>Changes</summary>
This is just code refactoring done as a preparation for adding
MLIR inliner cost model hook(s).
Related discussion: https://discourse.llvm.org/t/inliner-cost-model/2992
The logic of SCC-based MLIR inliner is separated into the Inliner
implementation. The MLIR inliner pass becomes, well, just a pass
that invokes the SCC-based MLIR inliner.
---
Patch is 43.35 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/84059.diff
11 Files Affected:
- (added) mlir/include/mlir/Transforms/Inliner.h (+124)
- (modified) mlir/include/mlir/Transforms/Passes.td (+4-2)
- (modified) mlir/lib/Transforms/CMakeLists.txt (+1-1)
- (added) mlir/lib/Transforms/InlinerPass.cpp (+158)
- (modified) mlir/lib/Transforms/Utils/CMakeLists.txt (+1)
- (renamed) mlir/lib/Transforms/Utils/Inliner.cpp (+221-313)
- (modified) mlir/test/Dialect/Affine/inlining.mlir (+1-1)
- (modified) mlir/test/Dialect/SPIRV/Transforms/inlining.mlir (+1-1)
- (modified) mlir/test/Transforms/inlining-dump-default-pipeline.mlir (+1-1)
- (modified) mlir/test/Transforms/inlining-recursive.mlir (+2-2)
- (modified) mlir/test/Transforms/inlining.mlir (+3-3)
``````````diff
diff --git a/mlir/include/mlir/Transforms/Inliner.h b/mlir/include/mlir/Transforms/Inliner.h
new file mode 100644
index 00000000000000..a74a32fcfec430
--- /dev/null
+++ b/mlir/include/mlir/Transforms/Inliner.h
@@ -0,0 +1,124 @@
+//===- Inliner.h - Inliner pass utilities -----------------------*- C++ -*-===//
+//
+// 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 header file declares utility structures for the inliner pass.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TRANSFORMS_INLINER_H
+#define MLIR_TRANSFORMS_INLINER_H
+
+#include "mlir/Analysis/CallGraph.h"
+#include "mlir/Interfaces/CallInterfaces.h"
+#include "mlir/Pass/AnalysisManager.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/StringMap.h"
+
+namespace mlir {
+class OpPassManager;
+class Operation;
+
+class InlinerConfig {
+public:
+ using PreInlineCallableOptPipelineTy = std::function<void(OpPassManager &)>;
+ using OpPipelinesTy = llvm::StringMap<OpPassManager>;
+
+ InlinerConfig() = default;
+ InlinerConfig(PreInlineCallableOptPipelineTy preInlineCallableOptPipeline,
+ unsigned maxInliningIterations)
+ : preInlineCallableOptPipeline(std::move(preInlineCallableOptPipeline)),
+ maxInliningIterations(maxInliningIterations) {}
+
+ const PreInlineCallableOptPipelineTy &
+ getPreInlineCallableOptPipeline() const {
+ return preInlineCallableOptPipeline;
+ }
+ const OpPipelinesTy &getOpPipelines() const { return opPipelines; }
+ unsigned getMaxInliningIterations() const { return maxInliningIterations; }
+ void
+ setPreInlineCallableOptPipeline(PreInlineCallableOptPipelineTy pipeline) {
+ preInlineCallableOptPipeline = std::move(pipeline);
+ }
+ void setOpPipelines(OpPipelinesTy pipelines) {
+ opPipelines = std::move(pipelines);
+ }
+ void setMaxInliningIterations(unsigned max) { maxInliningIterations = max; }
+
+private:
+ /// An optional function that constructs an optimization pipeline for
+ /// a given operation.
+ PreInlineCallableOptPipelineTy preInlineCallableOptPipeline;
+ /// A map of operation names to pass pipelines to use when optimizing
+ /// callable operations of these types. This provides a specialized pipeline
+ /// instead of the one produced by preInlineCallableOptPipeline.
+ OpPipelinesTy opPipelines;
+ /// For SCC-based inlining algorithms, specifies maximum number of iterations
+ /// when inlining within an SCC.
+ unsigned maxInliningIterations{0};
+};
+
+/// This is an implementation of the inliner
+/// that operates bottom up over the Strongly Connected Components(SCCs)
+/// of the CallGraph. This enables a more incremental propagation
+/// of inlining decisions from the leafs to the roots of the callgraph.
+///
+/// Derived implementations may rely on the same algorithm, but override
+/// the provided hooks to tune various algorithm aspects.
+class Inliner {
+public:
+ using RunPipelineHelperTy = std::function<LogicalResult(
+ Pass &pass, OpPassManager &pipeline, Operation *op)>;
+
+ virtual ~Inliner() {}
+ Inliner(Operation *op, CallGraph &cg, Pass &pass, AnalysisManager am,
+ RunPipelineHelperTy runPipelineHelper, const InlinerConfig &config)
+ : op(op), cg(cg), pass(pass), am(am),
+ runPipelineHelper(std::move(runPipelineHelper)), config(config) {}
+ Inliner(Inliner &) = delete;
+ void operator=(const Inliner &) = delete;
+
+ /// Perform inlining on a OpTrait::SymbolTable operation.
+ LogicalResult doInlining();
+
+ /// This struct represents a resolved call to a given callgraph node. Given
+ /// that the call does not actually contain a direct reference to the
+ /// Region(CallGraphNode) that it is dispatching to, we need to resolve them
+ /// explicitly.
+ struct ResolvedCall {
+ ResolvedCall(CallOpInterface call, CallGraphNode *sourceNode,
+ CallGraphNode *targetNode)
+ : call(call), sourceNode(sourceNode), targetNode(targetNode) {}
+ CallOpInterface call;
+ CallGraphNode *sourceNode, *targetNode;
+ };
+
+protected:
+ /// An OpTrait::SymbolTable operation to run the inlining on.
+ Operation *op;
+ /// A CallGraph analysis for the given operation.
+ CallGraph &cg;
+ /// A reference to the pass using this inliner.
+ Pass &pass;
+ /// Analysis manager for the given operation instance.
+ AnalysisManager am;
+ /// A callback for running a nested pass pipeline on the operation
+ /// contained within the main operation.
+ const RunPipelineHelperTy runPipelineHelper;
+ /// The inliner configuration parameters.
+ const InlinerConfig &config;
+
+private:
+ /// Forward declaration of the class providing the actual implementation.
+ class Impl;
+
+public:
+};
+} // namespace mlir
+
+#endif // MLIR_TRANSFORMS_INLINER_H
diff --git a/mlir/include/mlir/Transforms/Passes.td b/mlir/include/mlir/Transforms/Passes.td
index 2d2d54fb8fb5ea..1e61026d7fc3c6 100644
--- a/mlir/include/mlir/Transforms/Passes.td
+++ b/mlir/include/mlir/Transforms/Passes.td
@@ -268,8 +268,10 @@ def Inliner : Pass<"inline"> {
let summary = "Inline function calls";
let constructor = "mlir::createInlinerPass()";
let options = [
- Option<"defaultPipelineStr", "default-pipeline", "std::string",
- /*default=*/"\"canonicalize\"", "The default optimizer pipeline used for callables">,
+ Option<"preInlineCallableOptPipelineStr", "pre-inline-pipeline",
+ "std::string", /*default=*/"\"canonicalize\"",
+ "The optimizer pipeline used for callables that do not have "
+ "a dedicated optimizer pipeline in opPipelineList">,
ListOption<"opPipelineList", "op-pipelines", "OpPassManager",
"Callable operation specific optimizer pipelines (in the form "
"of `dialect.op(pipeline)`)">,
diff --git a/mlir/lib/Transforms/CMakeLists.txt b/mlir/lib/Transforms/CMakeLists.txt
index af51a4ab1157f1..6c32ecf8a2a2f1 100644
--- a/mlir/lib/Transforms/CMakeLists.txt
+++ b/mlir/lib/Transforms/CMakeLists.txt
@@ -5,7 +5,7 @@ add_mlir_library(MLIRTransforms
ControlFlowSink.cpp
CSE.cpp
GenerateRuntimeVerification.cpp
- Inliner.cpp
+ InlinerPass.cpp
LocationSnapshot.cpp
LoopInvariantCodeMotion.cpp
Mem2Reg.cpp
diff --git a/mlir/lib/Transforms/InlinerPass.cpp b/mlir/lib/Transforms/InlinerPass.cpp
new file mode 100644
index 00000000000000..5df598f1526b80
--- /dev/null
+++ b/mlir/lib/Transforms/InlinerPass.cpp
@@ -0,0 +1,158 @@
+//===- InlinerPass.cpp - Pass to inline function calls --------------------===//
+//
+// 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 basic inlining algorithm that operates bottom up over
+// the Strongly Connect Components(SCCs) of the CallGraph. This enables a more
+// incremental propagation of inlining decisions from the leafs to the roots of
+// the callgraph.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Transforms/Passes.h"
+
+#include "mlir/Analysis/CallGraph.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/Inliner.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_INLINER
+#include "mlir/Transforms/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+
+/// This function implements the inliner optimization pipeline.
+static void defaultInlinerOptPipeline(OpPassManager &pm) {
+ pm.addPass(createCanonicalizerPass());
+}
+
+//===----------------------------------------------------------------------===//
+// InlinerPass
+//===----------------------------------------------------------------------===//
+
+namespace {
+class InlinerPass : public impl::InlinerBase<InlinerPass> {
+public:
+ InlinerPass();
+ InlinerPass(const InlinerPass &) = default;
+ InlinerPass(
+ std::function<void(OpPassManager &)> preInlineCallableOptPipeline);
+ InlinerPass(std::function<void(OpPassManager &)> preInlineCallableOptPipeline,
+ llvm::StringMap<OpPassManager> opPipelines);
+ void runOnOperation() override;
+
+ /// A callback provided to the inliner driver to execute
+ /// the specified pass pipeline on the given operation
+ /// within the context of the current inliner pass,
+ /// which is passed as the first argument.
+ /// runPipeline API is protected within the Pass class,
+ /// so this helper is required to call it from the foreign
+ /// inliner driver.
+ static LogicalResult runPipelineHelper(Pass &pass, OpPassManager &pipeline,
+ Operation *op) {
+ return mlir::cast<InlinerPass>(pass).runPipeline(pipeline, op);
+ }
+
+private:
+ /// Attempt to initialize the options of this pass from the given string.
+ /// Derived classes may override this method to hook into the point at which
+ /// options are initialized, but should generally always invoke this base
+ /// class variant.
+ LogicalResult initializeOptions(StringRef options) override;
+
+ /// Inliner configuration parameters created from the pass options.
+ InlinerConfig config;
+};
+} // namespace
+
+InlinerPass::InlinerPass() : InlinerPass(defaultInlinerOptPipeline) {}
+
+InlinerPass::InlinerPass(
+ std::function<void(OpPassManager &)> preInlineCallableOptPipelineArg)
+ : InlinerPass(std::move(preInlineCallableOptPipelineArg),
+ llvm::StringMap<OpPassManager>{}) {}
+
+InlinerPass::InlinerPass(
+ std::function<void(OpPassManager &)> preInlineCallableOptPipeline,
+ llvm::StringMap<OpPassManager> opPipelines)
+ : config(std::move(preInlineCallableOptPipeline), maxInliningIterations) {
+ if (opPipelines.empty())
+ return;
+
+ // Update the option for the op specific optimization pipelines.
+ for (auto &it : opPipelines)
+ opPipelineList.addValue(it.second);
+ config.setOpPipelines(std::move(opPipelines));
+}
+
+void InlinerPass::runOnOperation() {
+ CallGraph &cg = getAnalysis<CallGraph>();
+
+ // The inliner should only be run on operations that define a symbol table,
+ // as the callgraph will need to resolve references.
+ Operation *op = getOperation();
+ if (!op->hasTrait<OpTrait::SymbolTable>()) {
+ op->emitOpError() << " was scheduled to run under the inliner, but does "
+ "not define a symbol table";
+ return signalPassFailure();
+ }
+
+ // Get an instance of the inliner.
+ Inliner inliner(op, cg, *this, getAnalysisManager(), runPipelineHelper,
+ config);
+
+ // Run the inlining.
+ if (failed(inliner.doInlining()))
+ signalPassFailure();
+ return;
+}
+
+LogicalResult InlinerPass::initializeOptions(StringRef options) {
+ if (failed(Pass::initializeOptions(options)))
+ return failure();
+
+ // Initialize the pipeline builder for operations without the dedicated
+ // optimization pipeline in opPipelineList to use the option string.
+ // TODO: Use a generic pass manager for the pre-inline pipeline, and remove
+ // this.
+ if (!preInlineCallableOptPipelineStr.empty()) {
+ std::string preInlineCallableOptPipelineCopy =
+ preInlineCallableOptPipelineStr;
+ config.setPreInlineCallableOptPipeline([=](OpPassManager &pm) {
+ (void)parsePassPipeline(preInlineCallableOptPipelineCopy, pm);
+ });
+ } else if (preInlineCallableOptPipelineStr.getNumOccurrences()) {
+ config.setPreInlineCallableOptPipeline(nullptr);
+ }
+
+ // Initialize the op specific pass pipelines.
+ llvm::StringMap<OpPassManager> pipelines;
+ for (OpPassManager pipeline : opPipelineList)
+ if (!pipeline.empty())
+ pipelines.try_emplace(pipeline.getOpAnchorName(), pipeline);
+ config.setOpPipelines(std::move(pipelines));
+
+ config.setMaxInliningIterations(maxInliningIterations);
+
+ return success();
+}
+
+std::unique_ptr<Pass> mlir::createInlinerPass() {
+ return std::make_unique<InlinerPass>();
+}
+std::unique_ptr<Pass>
+mlir::createInlinerPass(llvm::StringMap<OpPassManager> opPipelines) {
+ return std::make_unique<InlinerPass>(defaultInlinerOptPipeline,
+ std::move(opPipelines));
+}
+std::unique_ptr<Pass> mlir::createInlinerPass(
+ llvm::StringMap<OpPassManager> opPipelines,
+ std::function<void(OpPassManager &)> preInlineCallableOptPipelineBuilder) {
+ return std::make_unique<InlinerPass>(
+ std::move(preInlineCallableOptPipelineBuilder), std::move(opPipelines));
+}
diff --git a/mlir/lib/Transforms/Utils/CMakeLists.txt b/mlir/lib/Transforms/Utils/CMakeLists.txt
index 1c608e0634a67e..d6aac0e2da4f5a 100644
--- a/mlir/lib/Transforms/Utils/CMakeLists.txt
+++ b/mlir/lib/Transforms/Utils/CMakeLists.txt
@@ -5,6 +5,7 @@ add_mlir_library(MLIRTransformUtils
DialectConversion.cpp
FoldUtils.cpp
GreedyPatternRewriteDriver.cpp
+ Inliner.cpp
InliningUtils.cpp
LoopInvariantCodeMotionUtils.cpp
OneToNTypeConversion.cpp
diff --git a/mlir/lib/Transforms/Inliner.cpp b/mlir/lib/Transforms/Utils/Inliner.cpp
similarity index 76%
rename from mlir/lib/Transforms/Inliner.cpp
rename to mlir/lib/Transforms/Utils/Inliner.cpp
index b32b0fc28c78b0..be30159c9ac976 100644
--- a/mlir/lib/Transforms/Inliner.cpp
+++ b/mlir/lib/Transforms/Utils/Inliner.cpp
@@ -1,4 +1,4 @@
-//===- Inliner.cpp - Pass to inline function calls ------------------------===//
+//===- Inliner.cpp ---- SCC-based inliner ---------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
@@ -6,38 +6,29 @@
//
//===----------------------------------------------------------------------===//
//
-// This file implements a basic inlining algorithm that operates bottom up over
-// the Strongly Connect Components(SCCs) of the CallGraph. This enables a more
-// incremental propagation of inlining decisions from the leafs to the roots of
-// the callgraph.
+// This file implements Inliner that uses a basic inlining
+// algorithm that operates bottom up over the Strongly Connect Components(SCCs)
+// of the CallGraph. This enables a more incremental propagation of inlining
+// decisions from the leafs to the roots of the callgraph.
//
//===----------------------------------------------------------------------===//
-#include "mlir/Transforms/Passes.h"
-
-#include "mlir/Analysis/CallGraph.h"
+#include "mlir/Transforms/Inliner.h"
#include "mlir/IR/Threading.h"
#include "mlir/Interfaces/CallInterfaces.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
-#include "mlir/Pass/PassManager.h"
+#include "mlir/Pass/Pass.h"
#include "mlir/Support/DebugStringHelper.h"
#include "mlir/Transforms/InliningUtils.h"
#include "llvm/ADT/SCCIterator.h"
+#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Support/Debug.h"
-namespace mlir {
-#define GEN_PASS_DEF_INLINER
-#include "mlir/Transforms/Passes.h.inc"
-} // namespace mlir
-
#define DEBUG_TYPE "inlining"
using namespace mlir;
-/// This function implements the default inliner optimization pipeline.
-static void defaultInlinerOptPipeline(OpPassManager &pm) {
- pm.addPass(createCanonicalizerPass());
-}
+using ResolvedCall = Inliner::ResolvedCall;
//===----------------------------------------------------------------------===//
// Symbol Use Tracking
@@ -306,20 +297,6 @@ static LogicalResult runTransformOnCGSCCs(
return success();
}
-namespace {
-/// This struct represents a resolved call to a given callgraph node. Given that
-/// the call does not actually contain a direct reference to the
-/// Region(CallGraphNode) that it is dispatching to, we need to resolve them
-/// explicitly.
-struct ResolvedCall {
- ResolvedCall(CallOpInterface call, CallGraphNode *sourceNode,
- CallGraphNode *targetNode)
- : call(call), sourceNode(sourceNode), targetNode(targetNode) {}
- CallOpInterface call;
- CallGraphNode *sourceNode, *targetNode;
-};
-} // namespace
-
/// Collect all of the callable operations within the given range of blocks. If
/// `traverseNestedCGNodes` is true, this will also collect call operations
/// inside of nested callgraph nodes.
@@ -368,7 +345,7 @@ static void collectCallOps(iterator_range<Region::iterator> blocks,
}
//===----------------------------------------------------------------------===//
-// Inliner
+// InlinerInterfaceImpl
//===----------------------------------------------------------------------===//
#ifndef NDEBUG
@@ -397,9 +374,9 @@ static bool inlineHistoryIncludes(
namespace {
/// This class provides a specialization of the main inlining interface.
-struct Inliner : public InlinerInterface {
- Inliner(MLIRContext *context, CallGraph &cg,
- SymbolTableCollection &symbolTable)
+struct InlinerInterfaceImpl : public InlinerInterface {
+ InlinerInterfaceImpl(MLIRContext *context, CallGraph &cg,
+ SymbolTableCollection &symbolTable)
: InlinerInterface(context), cg(cg), symbolTable(symbolTable) {}
/// Process a set of blocks that have been inlined. This callback is invoked
@@ -442,45 +419,173 @@ struct Inliner : public InlinerInterface {
};
} // namespace
-/// Returns true if the given call should be inlined.
-static bool shouldInline(ResolvedCall &resolvedCall) {
- // Don't allow inlining terminator calls. We currently don't support this
- // case.
- if (resolvedCall.call->hasTrait<OpTrait::IsTerminator>())
- return false;
+namespace mlir {
- // Don't allow inlining if the target is an ancestor of the call. This
- // prevents inlining recursively.
- Region *callableRegion = resolvedCall.targetNode->getCallableRegion();
- if (callableRegion->isAncestor(resolvedCall.call->getParentRegion()))
- return false;
+class Inliner::Impl {
+public:
+ Impl(Inliner &inliner) : inliner(inliner) {}
- // Don't allow inlining if the callee has multiple blocks (unstructured
- // control flow) but we cannot be sure that the caller region supports that.
- bool calleeHasMultipleBlocks =
- llvm::hasNItemsOrMore(*callableRegion, /*N=*/2);
- // If both parent ops have the same type, it is safe to inline. Otherwise,
- // decide based on whether the op has the SingleBlock trait or not.
- // Note: This check does currently not account for SizedRegion/MaxSizedRegion.
- auto callerRegionSupportsMultipleBlocks = [&]() {
- return callableRegion->getParentOp()->getName() ==
- resolvedCall.call->getParentOp()->getName() ||
- !resolvedCall.call->getParentOp()
- ->mightHaveTrait<OpTrait::SingleBlock>();
- };
- if (calleeHasMultipleBlocks && !callerRegionSupportsMultipleBlocks())
- return false;
+ /// Attempt to inline calls within the given scc, and run simplifications,
+ /// until a fixed point is reached. This allows for the inlining of newly
+ /// devirtualized calls. Returns failure if there was a fatal error during
+ /// inlining.
+ LogicalResult inlineSCC(InlinerInterfaceImpl &inlinerIface,
+ CGUseList &useList, CallGraphSCC ¤tSCC,
+ MLIRContext *context);
- // Otherwise, inline.
- return true;
+private:
+ /// Optimize the nodes within the given SCC with one of the held optimization
+ /// pass pipelines. Returns failure if an error occurred during the
+ /// optimization of the SCC, success otherwise.
+ LogicalResult optimizeSCC(CallGraph &cg, CGUseList &useList,
+ CallGraphSCC ¤tSCC, MLIRContext *context);
+
+ /// Optimize the nodes within the given SCC in parallel. Returns failure if an
+ /// error occurred during the optimization of the SCC, success otherwise.
+ LogicalResult optimizeSCCAsync(MutableArrayRef<CallGraphNode *> nodesToVisit,
+ MLIRContext *context);
+
+ /// Optimize the given callable node with one of the pass managers provided
+ /// with `pipelines`, or the generic pre-inline pipeline. Returns failure if
+ /// an error occurred during the optimization of the callable, success
+ /// otherwise.
+ LogicalRes...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/84059
More information about the Mlir-commits
mailing list