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