1 //===- MLInlineAdvisor.cpp - machine learned InlineAdvisor ----------------===//
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 the interface between the inliner and a learned model.
10 // It delegates model evaluation to either the AOT compiled model (the
11 // 'release' mode) or a runtime-loaded model (the 'development' case).
12 //
13 //===----------------------------------------------------------------------===//
14 #include "llvm/Config/config.h"
15 #if defined(LLVM_HAVE_TF_AOT) || defined(LLVM_HAVE_TF_API)
16
17 #include <limits>
18 #include <unordered_map>
19 #include <unordered_set>
20
21 #include "llvm/ADT/SCCIterator.h"
22 #include "llvm/Analysis/CallGraph.h"
23 #include "llvm/Analysis/FunctionPropertiesAnalysis.h"
24 #include "llvm/Analysis/InlineCost.h"
25 #include "llvm/Analysis/MLInlineAdvisor.h"
26 #include "llvm/Analysis/MLModelRunner.h"
27 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
28 #include "llvm/Analysis/TargetLibraryInfo.h"
29 #include "llvm/Analysis/TargetTransformInfo.h"
30 #include "llvm/IR/InstIterator.h"
31 #include "llvm/IR/Instructions.h"
32 #include "llvm/IR/PassManager.h"
33 #include "llvm/Support/CommandLine.h"
34 #include "llvm/Support/Path.h"
35
36 using namespace llvm;
37
38 #define DEBUG_TYPE "inline-ml"
39
40 static cl::opt<float> SizeIncreaseThreshold(
41 "ml-advisor-size-increase-threshold", cl::Hidden,
42 cl::desc("Maximum factor by which expected native size may increase before "
43 "blocking any further inlining."),
44 cl::init(2.0));
45
46 const std::array<std::string, NumberOfFeatures> llvm::FeatureNameMap{
47 #define POPULATE_NAMES(INDEX_NAME, NAME, COMMENT) NAME,
48 INLINE_FEATURE_ITERATOR(POPULATE_NAMES)
49 #undef POPULATE_NAMES
50 };
51
52 const char *const llvm::DecisionName = "inlining_decision";
53 const char *const llvm::DefaultDecisionName = "inlining_default";
54 const char *const llvm::RewardName = "delta_size";
55
getInlinableCS(Instruction & I)56 CallBase *getInlinableCS(Instruction &I) {
57 if (auto *CS = dyn_cast<CallBase>(&I))
58 if (Function *Callee = CS->getCalledFunction()) {
59 if (!Callee->isDeclaration()) {
60 return CS;
61 }
62 }
63 return nullptr;
64 }
65
MLInlineAdvisor(Module & M,ModuleAnalysisManager & MAM,std::unique_ptr<MLModelRunner> Runner)66 MLInlineAdvisor::MLInlineAdvisor(Module &M, ModuleAnalysisManager &MAM,
67 std::unique_ptr<MLModelRunner> Runner)
68 : InlineAdvisor(
69 M, MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager()),
70 ModelRunner(std::move(Runner)), CG(new CallGraph(M)),
71 InitialIRSize(getModuleIRSize()), CurrentIRSize(InitialIRSize) {
72 assert(ModelRunner);
73
74 // Extract the 'call site height' feature - the position of a call site
75 // relative to the farthest statically reachable SCC node. We don't mutate
76 // this value while inlining happens. Empirically, this feature proved
77 // critical in behavioral cloning - i.e. training a model to mimic the manual
78 // heuristic's decisions - and, thus, equally important for training for
79 // improvement.
80 for (auto I = scc_begin(CG.get()); !I.isAtEnd(); ++I) {
81 const std::vector<CallGraphNode *> &CGNodes = *I;
82 unsigned Level = 0;
83 for (auto *CGNode : CGNodes) {
84 Function *F = CGNode->getFunction();
85 if (!F || F->isDeclaration())
86 continue;
87 for (auto &I : instructions(F)) {
88 if (auto *CS = getInlinableCS(I)) {
89 auto *Called = CS->getCalledFunction();
90 auto Pos = FunctionLevels.find(Called);
91 // In bottom up traversal, an inlinable callee is either in the
92 // same SCC, or to a function in a visited SCC. So not finding its
93 // level means we haven't visited it yet, meaning it's in this SCC.
94 if (Pos == FunctionLevels.end())
95 continue;
96 Level = std::max(Level, Pos->second + 1);
97 }
98 }
99 }
100 for (auto *CGNode : CGNodes) {
101 Function *F = CGNode->getFunction();
102 if (F && !F->isDeclaration())
103 FunctionLevels[F] = Level;
104 }
105 }
106 }
107
onPassEntry()108 void MLInlineAdvisor::onPassEntry() {
109 // Function passes executed between InlinerPass runs may have changed the
110 // module-wide features.
111 NodeCount = 0;
112 EdgeCount = 0;
113 for (auto &F : M)
114 if (!F.isDeclaration()) {
115 ++NodeCount;
116 EdgeCount += getLocalCalls(F);
117 }
118 }
119
getLocalCalls(Function & F)120 int64_t MLInlineAdvisor::getLocalCalls(Function &F) {
121 return FAM.getResult<FunctionPropertiesAnalysis>(F)
122 .DirectCallsToDefinedFunctions;
123 }
124
125 // Update the internal state of the advisor, and force invalidate feature
126 // analysis. Currently, we maintain minimal (and very simple) global state - the
127 // number of functions and the number of static calls. We also keep track of the
128 // total IR size in this module, to stop misbehaving policies at a certain bloat
129 // factor (SizeIncreaseThreshold)
onSuccessfulInlining(const MLInlineAdvice & Advice,bool CalleeWasDeleted)130 void MLInlineAdvisor::onSuccessfulInlining(const MLInlineAdvice &Advice,
131 bool CalleeWasDeleted) {
132 assert(!ForceStop);
133 Function *Caller = Advice.getCaller();
134 Function *Callee = Advice.getCallee();
135
136 // The caller features aren't valid anymore.
137 {
138 PreservedAnalyses PA = PreservedAnalyses::all();
139 PA.abandon<FunctionPropertiesAnalysis>();
140 FAM.invalidate(*Caller, PA);
141 }
142 int64_t IRSizeAfter =
143 getIRSize(*Caller) + (CalleeWasDeleted ? 0 : Advice.CalleeIRSize);
144 CurrentIRSize += IRSizeAfter - (Advice.CallerIRSize + Advice.CalleeIRSize);
145 if (CurrentIRSize > SizeIncreaseThreshold * InitialIRSize)
146 ForceStop = true;
147
148 // We can delta-update module-wide features. We know the inlining only changed
149 // the caller, and maybe the callee (by deleting the latter).
150 // Nodes are simple to update.
151 // For edges, we 'forget' the edges that the caller and callee used to have
152 // before inlining, and add back what they currently have together.
153 int64_t NewCallerAndCalleeEdges =
154 FAM.getResult<FunctionPropertiesAnalysis>(*Caller)
155 .DirectCallsToDefinedFunctions;
156
157 if (CalleeWasDeleted)
158 --NodeCount;
159 else
160 NewCallerAndCalleeEdges +=
161 FAM.getResult<FunctionPropertiesAnalysis>(*Callee)
162 .DirectCallsToDefinedFunctions;
163 EdgeCount += (NewCallerAndCalleeEdges - Advice.CallerAndCalleeEdges);
164 assert(CurrentIRSize >= 0 && EdgeCount >= 0 && NodeCount >= 0);
165 }
166
getModuleIRSize() const167 int64_t MLInlineAdvisor::getModuleIRSize() const {
168 int64_t Ret = 0;
169 for (auto &F : CG->getModule())
170 if (!F.isDeclaration())
171 Ret += getIRSize(F);
172 return Ret;
173 }
174
getAdviceImpl(CallBase & CB)175 std::unique_ptr<InlineAdvice> MLInlineAdvisor::getAdviceImpl(CallBase &CB) {
176 auto &Caller = *CB.getCaller();
177 auto &Callee = *CB.getCalledFunction();
178
179 auto GetAssumptionCache = [&](Function &F) -> AssumptionCache & {
180 return FAM.getResult<AssumptionAnalysis>(F);
181 };
182 auto &TIR = FAM.getResult<TargetIRAnalysis>(Callee);
183 auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(Caller);
184
185 auto MandatoryKind = InlineAdvisor::getMandatoryKind(CB, FAM, ORE);
186 // If this is a "never inline" case, there won't be any changes to internal
187 // state we need to track, so we can just return the base InlineAdvice, which
188 // will do nothing interesting.
189 // Same thing if this is a recursive case.
190 if (MandatoryKind == InlineAdvisor::MandatoryInliningKind::Never ||
191 &Caller == &Callee)
192 return getMandatoryAdvice(CB, false);
193
194 bool Mandatory =
195 MandatoryKind == InlineAdvisor::MandatoryInliningKind::Always;
196
197 // If we need to stop, we won't want to track anymore any state changes, so
198 // we just return the base InlineAdvice, which acts as a noop.
199 if (ForceStop) {
200 ORE.emit([&] {
201 return OptimizationRemarkMissed(DEBUG_TYPE, "ForceStop", &CB)
202 << "Won't attempt inlining because module size grew too much.";
203 });
204 return std::make_unique<InlineAdvice>(this, CB, ORE, Mandatory);
205 }
206
207 int CostEstimate = 0;
208 if (!Mandatory) {
209 auto IsCallSiteInlinable =
210 llvm::getInliningCostEstimate(CB, TIR, GetAssumptionCache);
211 if (!IsCallSiteInlinable) {
212 // We can't inline this for correctness reasons, so return the base
213 // InlineAdvice, as we don't care about tracking any state changes (which
214 // won't happen).
215 return std::make_unique<InlineAdvice>(this, CB, ORE, false);
216 }
217 CostEstimate = *IsCallSiteInlinable;
218 }
219
220 if (Mandatory)
221 return getMandatoryAdvice(CB, true);
222
223 auto NrCtantParams = 0;
224 for (auto I = CB.arg_begin(), E = CB.arg_end(); I != E; ++I) {
225 NrCtantParams += (isa<Constant>(*I));
226 }
227
228 auto &CallerBefore = FAM.getResult<FunctionPropertiesAnalysis>(Caller);
229 auto &CalleeBefore = FAM.getResult<FunctionPropertiesAnalysis>(Callee);
230
231 ModelRunner->setFeature(FeatureIndex::CalleeBasicBlockCount,
232 CalleeBefore.BasicBlockCount);
233 ModelRunner->setFeature(FeatureIndex::CallSiteHeight,
234 FunctionLevels[&Caller]);
235 ModelRunner->setFeature(FeatureIndex::NodeCount, NodeCount);
236 ModelRunner->setFeature(FeatureIndex::NrCtantParams, NrCtantParams);
237 ModelRunner->setFeature(FeatureIndex::CostEstimate, CostEstimate);
238 ModelRunner->setFeature(FeatureIndex::EdgeCount, EdgeCount);
239 ModelRunner->setFeature(FeatureIndex::CallerUsers, CallerBefore.Uses);
240 ModelRunner->setFeature(FeatureIndex::CallerConditionallyExecutedBlocks,
241 CallerBefore.BlocksReachedFromConditionalInstruction);
242 ModelRunner->setFeature(FeatureIndex::CallerBasicBlockCount,
243 CallerBefore.BasicBlockCount);
244 ModelRunner->setFeature(FeatureIndex::CalleeConditionallyExecutedBlocks,
245 CalleeBefore.BlocksReachedFromConditionalInstruction);
246 ModelRunner->setFeature(FeatureIndex::CalleeUsers, CalleeBefore.Uses);
247 return getAdviceFromModel(CB, ORE);
248 }
249
250 std::unique_ptr<MLInlineAdvice>
getAdviceFromModel(CallBase & CB,OptimizationRemarkEmitter & ORE)251 MLInlineAdvisor::getAdviceFromModel(CallBase &CB,
252 OptimizationRemarkEmitter &ORE) {
253 return std::make_unique<MLInlineAdvice>(this, CB, ORE, ModelRunner->run());
254 }
255
getMandatoryAdvice(CallBase & CB,bool Advice)256 std::unique_ptr<InlineAdvice> MLInlineAdvisor::getMandatoryAdvice(CallBase &CB,
257 bool Advice) {
258 // Make sure we track inlinings in all cases - mandatory or not.
259 if (Advice && !ForceStop)
260 return getMandatoryAdviceImpl(CB);
261
262 // If this is a "never inline" case, there won't be any changes to internal
263 // state we need to track, so we can just return the base InlineAdvice, which
264 // will do nothing interesting.
265 // Same if we are forced to stop - we don't track anymore.
266 return std::make_unique<InlineAdvice>(this, CB, getCallerORE(CB), Advice);
267 }
268
269 std::unique_ptr<MLInlineAdvice>
getMandatoryAdviceImpl(CallBase & CB)270 MLInlineAdvisor::getMandatoryAdviceImpl(CallBase &CB) {
271 return std::make_unique<MLInlineAdvice>(this, CB, getCallerORE(CB), true);
272 }
273
reportContextForRemark(DiagnosticInfoOptimizationBase & OR)274 void MLInlineAdvice::reportContextForRemark(
275 DiagnosticInfoOptimizationBase &OR) {
276 using namespace ore;
277 OR << NV("Callee", Callee->getName());
278 for (size_t I = 0; I < NumberOfFeatures; ++I)
279 OR << NV(FeatureNameMap[I], getAdvisor()->getModelRunner().getFeature(I));
280 OR << NV("ShouldInline", isInliningRecommended());
281 }
282
recordInliningImpl()283 void MLInlineAdvice::recordInliningImpl() {
284 ORE.emit([&]() {
285 OptimizationRemark R(DEBUG_TYPE, "InliningSuccess", DLoc, Block);
286 reportContextForRemark(R);
287 return R;
288 });
289 getAdvisor()->onSuccessfulInlining(*this, /*CalleeWasDeleted*/ false);
290 }
291
recordInliningWithCalleeDeletedImpl()292 void MLInlineAdvice::recordInliningWithCalleeDeletedImpl() {
293 ORE.emit([&]() {
294 OptimizationRemark R(DEBUG_TYPE, "InliningSuccessWithCalleeDeleted", DLoc,
295 Block);
296 reportContextForRemark(R);
297 return R;
298 });
299 getAdvisor()->onSuccessfulInlining(*this, /*CalleeWasDeleted*/ true);
300 }
301
recordUnsuccessfulInliningImpl(const InlineResult & Result)302 void MLInlineAdvice::recordUnsuccessfulInliningImpl(
303 const InlineResult &Result) {
304 ORE.emit([&]() {
305 OptimizationRemarkMissed R(DEBUG_TYPE, "InliningAttemptedAndUnsuccessful",
306 DLoc, Block);
307 reportContextForRemark(R);
308 return R;
309 });
310 }
recordUnattemptedInliningImpl()311 void MLInlineAdvice::recordUnattemptedInliningImpl() {
312 ORE.emit([&]() {
313 OptimizationRemarkMissed R(DEBUG_TYPE, "IniningNotAttempted", DLoc, Block);
314 reportContextForRemark(R);
315 return R;
316 });
317 }
318 #endif // defined(LLVM_HAVE_TF_AOT) || defined(LLVM_HAVE_TF_API)
319