Skip to content

Commit 2542d34

Browse files
vzakharivzakhari
authored
[mlir][inliner] Refactor MLIR inliner pass and utils. (#84059)
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.
1 parent 064c2e7 commit 2542d34

File tree

6 files changed

+500
-315
lines changed

6 files changed

+500
-315
lines changed

‎mlir/include/mlir/Transforms/Inliner.h

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
//===- Inliner.h - Inliner pass utilities -----------------------*- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// This header file declares utility structures for the inliner pass.
10+
//
11+
//===----------------------------------------------------------------------===//
12+
13+
#ifndef MLIR_TRANSFORMS_INLINER_H
14+
#define MLIR_TRANSFORMS_INLINER_H
15+
16+
#include "mlir/Analysis/CallGraph.h"
17+
#include "mlir/Interfaces/CallInterfaces.h"
18+
#include "mlir/Pass/AnalysisManager.h"
19+
#include "mlir/Pass/PassManager.h"
20+
#include "mlir/Support/LogicalResult.h"
21+
#include "llvm/ADT/StringMap.h"
22+
23+
namespace mlir {
24+
class OpPassManager;
25+
class Operation;
26+
27+
class InlinerConfig {
28+
public:
29+
using DefaultPipelineTy = std::function<void(OpPassManager &)>;
30+
using OpPipelinesTy = llvm::StringMap<OpPassManager>;
31+
32+
InlinerConfig() = default;
33+
InlinerConfig(DefaultPipelineTy defaultPipeline,
34+
unsigned maxInliningIterations)
35+
: defaultPipeline(std::move(defaultPipeline)),
36+
maxInliningIterations(maxInliningIterations) {}
37+
38+
const DefaultPipelineTy &getDefaultPipeline() const {
39+
return defaultPipeline;
40+
}
41+
const OpPipelinesTy &getOpPipelines() const { return opPipelines; }
42+
unsigned getMaxInliningIterations() const { return maxInliningIterations; }
43+
void setDefaultPipeline(DefaultPipelineTy pipeline) {
44+
defaultPipeline = std::move(pipeline);
45+
}
46+
void setOpPipelines(OpPipelinesTy pipelines) {
47+
opPipelines = std::move(pipelines);
48+
}
49+
void setMaxInliningIterations(unsigned max) { maxInliningIterations = max; }
50+
51+
private:
52+
/// An optional function that constructs an optimization pipeline for
53+
/// a given operation. This optimization pipeline is applied
54+
/// only to those callable operations that do not have dedicated
55+
/// optimization pipeline in opPipelines (based on the operation name).
56+
DefaultPipelineTy defaultPipeline;
57+
/// A map of operation names to pass pipelines to use when optimizing
58+
/// callable operations of these types. This provides a specialized pipeline
59+
/// instead of the one produced by defaultPipeline.
60+
OpPipelinesTy opPipelines;
61+
/// For SCC-based inlining algorithms, specifies maximum number of iterations
62+
/// when inlining within an SCC.
63+
unsigned maxInliningIterations{0};
64+
};
65+
66+
/// This is an implementation of the inliner
67+
/// that operates bottom up over the Strongly Connected Components(SCCs)
68+
/// of the CallGraph. This enables a more incremental propagation
69+
/// of inlining decisions from the leafs to the roots of the callgraph.
70+
class Inliner {
71+
public:
72+
using RunPipelineHelperTy = std::function<LogicalResult(
73+
Pass &pass, OpPassManager &pipeline, Operation *op)>;
74+
75+
Inliner(Operation *op, CallGraph &cg, Pass &pass, AnalysisManager am,
76+
RunPipelineHelperTy runPipelineHelper, const InlinerConfig &config)
77+
: op(op), cg(cg), pass(pass), am(am),
78+
runPipelineHelper(std::move(runPipelineHelper)), config(config) {}
79+
Inliner(Inliner &) = delete;
80+
void operator=(const Inliner &) = delete;
81+
82+
/// Perform inlining on a OpTrait::SymbolTable operation.
83+
LogicalResult doInlining();
84+
85+
/// This struct represents a resolved call to a given callgraph node. Given
86+
/// that the call does not actually contain a direct reference to the
87+
/// Region(CallGraphNode) that it is dispatching to, we need to resolve them
88+
/// explicitly.
89+
struct ResolvedCall {
90+
ResolvedCall(CallOpInterface call, CallGraphNode *sourceNode,
91+
CallGraphNode *targetNode)
92+
: call(call), sourceNode(sourceNode), targetNode(targetNode) {}
93+
CallOpInterface call;
94+
CallGraphNode *sourceNode, *targetNode;
95+
};
96+
97+
protected:
98+
/// An OpTrait::SymbolTable operation to run the inlining on.
99+
Operation *op;
100+
/// A CallGraph analysis for the given operation.
101+
CallGraph &cg;
102+
/// A reference to the pass using this inliner.
103+
Pass &pass;
104+
/// Analysis manager for the given operation instance.
105+
AnalysisManager am;
106+
/// A callback for running a nested pass pipeline on the operation
107+
/// contained within the main operation.
108+
const RunPipelineHelperTy runPipelineHelper;
109+
/// The inliner configuration parameters.
110+
const InlinerConfig &config;
111+
112+
private:
113+
/// Forward declaration of the class providing the actual implementation.
114+
class Impl;
115+
116+
public:
117+
};
118+
} // namespace mlir
119+
120+
#endif // MLIR_TRANSFORMS_INLINER_H

‎mlir/include/mlir/Transforms/Passes.td

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -269,7 +269,9 @@ def Inliner : Pass<"inline"> {
269269
let constructor = "mlir::createInlinerPass()";
270270
let options = [
271271
Option<"defaultPipelineStr", "default-pipeline", "std::string",
272-
/*default=*/"\"canonicalize\"", "The default optimizer pipeline used for callables">,
272+
/*default=*/"\"canonicalize\"",
273+
"The optimizer pipeline used for callables that do not have "
274+
"a dedicated optimizer pipeline in opPipelineList">,
273275
ListOption<"opPipelineList", "op-pipelines", "OpPassManager",
274276
"Callable operation specific optimizer pipelines (in the form "
275277
"of `dialect.op(pipeline)`)">,

‎mlir/lib/Transforms/CMakeLists.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ add_mlir_library(MLIRTransforms
55
ControlFlowSink.cpp
66
CSE.cpp
77
GenerateRuntimeVerification.cpp
8-
Inliner.cpp
8+
InlinerPass.cpp
99
LocationSnapshot.cpp
1010
LoopInvariantCodeMotion.cpp
1111
Mem2Reg.cpp

‎mlir/lib/Transforms/InlinerPass.cpp

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
//===- InlinerPass.cpp - Pass to inline function calls --------------------===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// This file implements a basic inlining algorithm that operates bottom up over
10+
// the Strongly Connect Components(SCCs) of the CallGraph. This enables a more
11+
// incremental propagation of inlining decisions from the leafs to the roots of
12+
// the callgraph.
13+
//
14+
//===----------------------------------------------------------------------===//
15+
16+
#include "mlir/Transforms/Passes.h"
17+
18+
#include "mlir/Analysis/CallGraph.h"
19+
#include "mlir/Pass/PassManager.h"
20+
#include "mlir/Transforms/Inliner.h"
21+
22+
namespace mlir {
23+
#define GEN_PASS_DEF_INLINER
24+
#include "mlir/Transforms/Passes.h.inc"
25+
} // namespace mlir
26+
27+
using namespace mlir;
28+
29+
/// This function implements the inliner optimization pipeline.
30+
static void defaultInlinerOptPipeline(OpPassManager &pm) {
31+
pm.addPass(createCanonicalizerPass());
32+
}
33+
34+
//===----------------------------------------------------------------------===//
35+
// InlinerPass
36+
//===----------------------------------------------------------------------===//
37+
38+
namespace {
39+
class InlinerPass : public impl::InlinerBase<InlinerPass> {
40+
public:
41+
InlinerPass();
42+
InlinerPass(const InlinerPass &) = default;
43+
InlinerPass(std::function<void(OpPassManager &)> defaultPipeline);
44+
InlinerPass(std::function<void(OpPassManager &)> defaultPipeline,
45+
llvm::StringMap<OpPassManager> opPipelines);
46+
void runOnOperation() override;
47+
48+
/// A callback provided to the inliner driver to execute
49+
/// the specified pass pipeline on the given operation
50+
/// within the context of the current inliner pass,
51+
/// which is passed as the first argument.
52+
/// runPipeline API is protected within the Pass class,
53+
/// so this helper is required to call it from the foreign
54+
/// inliner driver.
55+
static LogicalResult runPipelineHelper(Pass &pass, OpPassManager &pipeline,
56+
Operation *op) {
57+
return mlir::cast<InlinerPass>(pass).runPipeline(pipeline, op);
58+
}
59+
60+
private:
61+
/// Attempt to initialize the options of this pass from the given string.
62+
/// Derived classes may override this method to hook into the point at which
63+
/// options are initialized, but should generally always invoke this base
64+
/// class variant.
65+
LogicalResult initializeOptions(StringRef options) override;
66+
67+
/// Inliner configuration parameters created from the pass options.
68+
InlinerConfig config;
69+
};
70+
} // namespace
71+
72+
InlinerPass::InlinerPass() : InlinerPass(defaultInlinerOptPipeline) {}
73+
74+
InlinerPass::InlinerPass(
75+
std::function<void(OpPassManager &)> defaultPipelineArg)
76+
: InlinerPass(std::move(defaultPipelineArg),
77+
llvm::StringMap<OpPassManager>{}) {}
78+
79+
InlinerPass::InlinerPass(std::function<void(OpPassManager &)> defaultPipeline,
80+
llvm::StringMap<OpPassManager> opPipelines)
81+
: config(std::move(defaultPipeline), maxInliningIterations) {
82+
if (opPipelines.empty())
83+
return;
84+
85+
// Update the option for the op specific optimization pipelines.
86+
for (auto &it : opPipelines)
87+
opPipelineList.addValue(it.second);
88+
config.setOpPipelines(std::move(opPipelines));
89+
}
90+
91+
void InlinerPass::runOnOperation() {
92+
CallGraph &cg = getAnalysis<CallGraph>();
93+
94+
// The inliner should only be run on operations that define a symbol table,
95+
// as the callgraph will need to resolve references.
96+
Operation *op = getOperation();
97+
if (!op->hasTrait<OpTrait::SymbolTable>()) {
98+
op->emitOpError() << " was scheduled to run under the inliner, but does "
99+
"not define a symbol table";
100+
return signalPassFailure();
101+
}
102+
103+
// Get an instance of the inliner.
104+
Inliner inliner(op, cg, *this, getAnalysisManager(), runPipelineHelper,
105+
config);
106+
107+
// Run the inlining.
108+
if (failed(inliner.doInlining()))
109+
signalPassFailure();
110+
return;
111+
}
112+
113+
LogicalResult InlinerPass::initializeOptions(StringRef options) {
114+
if (failed(Pass::initializeOptions(options)))
115+
return failure();
116+
117+
// Initialize the pipeline builder for operations without the dedicated
118+
// optimization pipeline in opPipelineList to use the option string.
119+
// TODO: Use a generic pass manager for the pre-inline pipeline, and remove
120+
// this.
121+
if (!defaultPipelineStr.empty()) {
122+
std::string defaultPipelineCopy = defaultPipelineStr;
123+
config.setDefaultPipeline([=](OpPassManager &pm) {
124+
(void)parsePassPipeline(defaultPipelineCopy, pm);
125+
});
126+
} else if (defaultPipelineStr.getNumOccurrences()) {
127+
config.setDefaultPipeline(nullptr);
128+
}
129+
130+
// Initialize the op specific pass pipelines.
131+
llvm::StringMap<OpPassManager> pipelines;
132+
for (OpPassManager pipeline : opPipelineList)
133+
if (!pipeline.empty())
134+
pipelines.try_emplace(pipeline.getOpAnchorName(), pipeline);
135+
config.setOpPipelines(std::move(pipelines));
136+
137+
config.setMaxInliningIterations(maxInliningIterations);
138+
139+
return success();
140+
}
141+
142+
std::unique_ptr<Pass> mlir::createInlinerPass() {
143+
return std::make_unique<InlinerPass>();
144+
}
145+
std::unique_ptr<Pass>
146+
mlir::createInlinerPass(llvm::StringMap<OpPassManager> opPipelines) {
147+
return std::make_unique<InlinerPass>(defaultInlinerOptPipeline,
148+
std::move(opPipelines));
149+
}
150+
std::unique_ptr<Pass> mlir::createInlinerPass(
151+
llvm::StringMap<OpPassManager> opPipelines,
152+
std::function<void(OpPassManager &)> defaultPipelineBuilder) {
153+
return std::make_unique<InlinerPass>(std::move(defaultPipelineBuilder),
154+
std::move(opPipelines));
155+
}

‎mlir/lib/Transforms/Utils/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ add_mlir_library(MLIRTransformUtils
55
DialectConversion.cpp
66
FoldUtils.cpp
77
GreedyPatternRewriteDriver.cpp
8+
Inliner.cpp
89
InliningUtils.cpp
910
LoopInvariantCodeMotionUtils.cpp
1011
OneToNTypeConversion.cpp

0 commit comments

Comments
 (0)