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/InlineCost.h"
24 #include "llvm/Analysis/InlineFeaturesAnalysis.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 
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 
66 MLInlineAdvisor::MLInlineAdvisor(Module &M, ModuleAnalysisManager &MAM,
67                                  std::unique_ptr<MLModelRunner> Runner)
68     : InlineAdvisor(
69           MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager()),
70       M(M), 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 
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 
120 int64_t MLInlineAdvisor::getLocalCalls(Function &F) {
121   return FAM.getResult<InlineFeaturesAnalysis>(F).DirectCallsToDefinedFunctions;
122 }
123 
124 // Update the internal state of the advisor, and force invalidate feature
125 // analysis. Currently, we maintain minimal (and very simple) global state - the
126 // number of functions and the number of static calls. We also keep track of the
127 // total IR size in this module, to stop misbehaving policies at a certain bloat
128 // factor (SizeIncreaseThreshold)
129 void MLInlineAdvisor::onSuccessfulInlining(const MLInlineAdvice &Advice,
130                                            bool CalleeWasDeleted) {
131   assert(!ForceStop);
132   Function *Caller = Advice.getCaller();
133   Function *Callee = Advice.getCallee();
134 
135   // The caller features aren't valid anymore.
136   FAM.invalidate<InlineFeaturesAnalysis>(*Caller);
137   int64_t IRSizeAfter =
138       getIRSize(*Caller) + (CalleeWasDeleted ? 0 : Advice.CalleeIRSize);
139   CurrentIRSize += IRSizeAfter - (Advice.CallerIRSize + Advice.CalleeIRSize);
140   if (CurrentIRSize > SizeIncreaseThreshold * InitialIRSize)
141     ForceStop = true;
142 
143   // We can delta-update module-wide features. We know the inlining only changed
144   // the caller, and maybe the callee (by deleting the latter).
145   // Nodes are simple to update.
146   // For edges, we 'forget' the edges that the caller and callee used to have
147   // before inlining, and add back what they currently have together.
148   int64_t NewCallerAndCalleeEdges =
149       FAM.getResult<InlineFeaturesAnalysis>(*Caller)
150           .DirectCallsToDefinedFunctions;
151 
152   if (CalleeWasDeleted)
153     --NodeCount;
154   else
155     NewCallerAndCalleeEdges += FAM.getResult<InlineFeaturesAnalysis>(*Callee)
156                                    .DirectCallsToDefinedFunctions;
157   EdgeCount += (NewCallerAndCalleeEdges - Advice.CallerAndCalleeEdges);
158   assert(CurrentIRSize >= 0 && EdgeCount >= 0 && NodeCount >= 0);
159 }
160 
161 int64_t MLInlineAdvisor::getModuleIRSize() const {
162   int64_t Ret = 0;
163   for (auto &F : CG->getModule())
164     if (!F.isDeclaration())
165       Ret += getIRSize(F);
166   return Ret;
167 }
168 
169 std::unique_ptr<InlineAdvice> MLInlineAdvisor::getAdvice(CallBase &CB) {
170   auto &Caller = *CB.getCaller();
171   auto &Callee = *CB.getCalledFunction();
172 
173   auto GetAssumptionCache = [&](Function &F) -> AssumptionCache & {
174     return FAM.getResult<AssumptionAnalysis>(F);
175   };
176   auto GetTLI = [&](Function &F) -> const TargetLibraryInfo & {
177     return FAM.getResult<TargetLibraryAnalysis>(F);
178   };
179 
180   auto &TIR = FAM.getResult<TargetIRAnalysis>(Callee);
181   auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(Caller);
182 
183   auto TrivialDecision =
184       llvm::getAttributeBasedInliningDecision(CB, &Callee, TIR, GetTLI);
185 
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 ((TrivialDecision.hasValue() && !TrivialDecision->isSuccess()) ||
191       &Caller == &Callee)
192     return std::make_unique<InlineAdvice>(this, CB, ORE, false);
193 
194   bool Mandatory = TrivialDecision.hasValue() && TrivialDecision->isSuccess();
195 
196   // If we need to stop, we won't want to track anymore any state changes, so
197   // we just return the base InlineAdvice, which acts as a noop.
198   if (ForceStop) {
199     ORE.emit([&] {
200       return OptimizationRemarkMissed(DEBUG_TYPE, "ForceStop", &CB)
201              << "Won't attempt inlining because module size grew too much.";
202     });
203     return std::make_unique<InlineAdvice>(this, CB, ORE, Mandatory);
204   }
205 
206   int CostEstimate = 0;
207   if (!Mandatory) {
208     auto IsCallSiteInlinable =
209         llvm::getInliningCostEstimate(CB, TIR, GetAssumptionCache);
210     if (!IsCallSiteInlinable) {
211       // We can't inline this for correctness reasons, so return the base
212       // InlineAdvice, as we don't care about tracking any state changes (which
213       // won't happen).
214       return std::make_unique<InlineAdvice>(this, CB, ORE, false);
215     }
216     CostEstimate = *IsCallSiteInlinable;
217   }
218 
219   if (Mandatory)
220     return getMandatoryAdvice(CB, ORE);
221 
222   auto NrCtantParams = 0;
223   for (auto I = CB.arg_begin(), E = CB.arg_end(); I != E; ++I) {
224     NrCtantParams += (isa<Constant>(*I));
225   }
226 
227   auto &CallerBefore = FAM.getResult<InlineFeaturesAnalysis>(Caller);
228   auto &CalleeBefore = FAM.getResult<InlineFeaturesAnalysis>(Callee);
229 
230   ModelRunner->setFeature(FeatureIndex::CalleeBasicBlockCount,
231                           CalleeBefore.BasicBlockCount);
232   ModelRunner->setFeature(FeatureIndex::CallSiteHeight,
233                           FunctionLevels[&Caller]);
234   ModelRunner->setFeature(FeatureIndex::NodeCount, NodeCount);
235   ModelRunner->setFeature(FeatureIndex::NrCtantParams, NrCtantParams);
236   ModelRunner->setFeature(FeatureIndex::CostEstimate, CostEstimate);
237   ModelRunner->setFeature(FeatureIndex::EdgeCount, EdgeCount);
238   ModelRunner->setFeature(FeatureIndex::CallerUsers, CallerBefore.Uses);
239   ModelRunner->setFeature(FeatureIndex::CallerConditionallyExecutedBlocks,
240                           CallerBefore.BlocksReachedFromConditionalInstruction);
241   ModelRunner->setFeature(FeatureIndex::CallerBasicBlockCount,
242                           CallerBefore.BasicBlockCount);
243   ModelRunner->setFeature(FeatureIndex::CalleeConditionallyExecutedBlocks,
244                           CalleeBefore.BlocksReachedFromConditionalInstruction);
245   ModelRunner->setFeature(FeatureIndex::CalleeUsers, CalleeBefore.Uses);
246   return getAdviceFromModel(CB, ORE);
247 }
248 
249 std::unique_ptr<MLInlineAdvice>
250 MLInlineAdvisor::getAdviceFromModel(CallBase &CB,
251                                     OptimizationRemarkEmitter &ORE) {
252   return std::make_unique<MLInlineAdvice>(this, CB, ORE, ModelRunner->run());
253 }
254 
255 std::unique_ptr<MLInlineAdvice>
256 MLInlineAdvisor::getMandatoryAdvice(CallBase &CB,
257                                     OptimizationRemarkEmitter &ORE) {
258   return std::make_unique<MLInlineAdvice>(this, CB, ORE, true);
259 }
260 
261 void MLInlineAdvice::reportContextForRemark(
262     DiagnosticInfoOptimizationBase &OR) {
263   using namespace ore;
264   OR << NV("Callee", Callee->getName());
265   for (size_t I = 0; I < NumberOfFeatures; ++I)
266     OR << NV(FeatureNameMap[I], getAdvisor()->getModelRunner().getFeature(I));
267   OR << NV("ShouldInline", isInliningRecommended());
268 }
269 
270 void MLInlineAdvice::recordInliningImpl() {
271   ORE.emit([&]() {
272     OptimizationRemark R(DEBUG_TYPE, "InliningSuccess", DLoc, Block);
273     reportContextForRemark(R);
274     return R;
275   });
276   getAdvisor()->onSuccessfulInlining(*this, /*CalleeWasDeleted*/ false);
277 }
278 
279 void MLInlineAdvice::recordInliningWithCalleeDeletedImpl() {
280   ORE.emit([&]() {
281     OptimizationRemark R(DEBUG_TYPE, "InliningSuccessWithCalleeDeleted", DLoc,
282                          Block);
283     reportContextForRemark(R);
284     return R;
285   });
286   getAdvisor()->onSuccessfulInlining(*this, /*CalleeWasDeleted*/ true);
287 }
288 
289 void MLInlineAdvice::recordUnsuccessfulInliningImpl(
290     const InlineResult &Result) {
291   ORE.emit([&]() {
292     OptimizationRemarkMissed R(DEBUG_TYPE, "InliningAttemptedAndUnsuccessful",
293                                DLoc, Block);
294     reportContextForRemark(R);
295     return R;
296   });
297 }
298 void MLInlineAdvice::recordUnattemptedInliningImpl() {
299   ORE.emit([&]() {
300     OptimizationRemarkMissed R(DEBUG_TYPE, "IniningNotAttempted", DLoc, Block);
301     reportContextForRemark(R);
302     return R;
303   });
304 }
305 #endif // defined(LLVM_HAVE_TF_AOT) || defined(LLVM_HAVE_TF_API)
306