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