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