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/Analysis/MLInlineAdvisor.h"
15#include "llvm/ADT/SCCIterator.h"
16#include "llvm/Analysis/AssumptionCache.h"
17#include "llvm/Analysis/BlockFrequencyInfo.h"
18#include "llvm/Analysis/CallGraph.h"
19#include "llvm/Analysis/FunctionPropertiesAnalysis.h"
20#include "llvm/Analysis/InlineCost.h"
21#include "llvm/Analysis/InlineModelFeatureMaps.h"
22#include "llvm/Analysis/LazyCallGraph.h"
23#include "llvm/Analysis/LoopInfo.h"
24#include "llvm/Analysis/MLModelRunner.h"
25#include "llvm/Analysis/OptimizationRemarkEmitter.h"
26#include "llvm/Analysis/ProfileSummaryInfo.h"
27#include "llvm/Analysis/ReleaseModeModelRunner.h"
28#include "llvm/Analysis/TargetTransformInfo.h"
29#include "llvm/Analysis/TensorSpec.h"
30#include "llvm/Analysis/Utils/MLGOUtils.h"
31#include "llvm/IR/Dominators.h"
32#include "llvm/IR/InstIterator.h"
33#include "llvm/IR/Module.h"
34#include "llvm/IR/PassManager.h"
35#include "llvm/Support/CommandLine.h"
36
37using namespace llvm;
38
39static cl::opt<std::string> InteractiveChannelBaseName(
40 "inliner-interactive-channel-base", cl::Hidden,
41 cl::desc(
42 "Base file path for the interactive mode. The incoming filename should "
43 "have the name <inliner-interactive-channel-base>.in, while the "
44 "outgoing name should be <inliner-interactive-channel-base>.out"));
45static const std::string InclDefaultMsg =
46 (Twine("In interactive mode, also send the default policy decision: ") +
47 DefaultDecisionName + ".")
48 .str();
49static cl::opt<bool>
50 InteractiveIncludeDefault("inliner-interactive-include-default", cl::Hidden,
51 cl::desc(InclDefaultMsg));
52
53enum class SkipMLPolicyCriteria { Never, IfCallerIsNotCold };
54
55static cl::opt<SkipMLPolicyCriteria> SkipPolicy(
56 "ml-inliner-skip-policy", cl::Hidden, cl::init(Val: SkipMLPolicyCriteria::Never),
57 cl::values(clEnumValN(SkipMLPolicyCriteria::Never, "never", "never"),
58 clEnumValN(SkipMLPolicyCriteria::IfCallerIsNotCold,
59 "if-caller-not-cold", "if the caller is not cold")));
60
61static cl::opt<std::string> ModelSelector("ml-inliner-model-selector",
62 cl::Hidden, cl::init(Val: ""));
63
64static cl::opt<bool> StopImmediatelyForTest("ml-inliner-stop-immediately",
65 cl::Hidden);
66
67#if defined(LLVM_HAVE_TF_AOT_INLINERSIZEMODEL)
68// codegen-ed file
69#include "InlinerSizeModel.h" // NOLINT
70using CompiledModelType = llvm::InlinerSizeModel;
71#else
72using CompiledModelType = NoopSavedModelImpl;
73#endif
74
75#if defined(LLVM_HAVE_MLIR_LOWERING_INLINER)
76constexpr bool HaveMLIRLoweringInliner = true;
77#include "llvm/Analysis/EmitCModelRunner.h"
78#include "llvm/Analysis/InlinerModels.h"
79
80enum class EmitCModelChoice {
81 Default,
82#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) CLASS_NAME,
83#include "llvm/Analysis/InlinerModels.def"
84};
85
86static llvm::cl::opt<EmitCModelChoice> SelectedMLGOModel(
87 "mlgo-model", llvm::cl::desc("Select the MLGO model to execute:"),
88 llvm::cl::init(EmitCModelChoice::Default),
89 llvm::cl::values(clEnumValN(EmitCModelChoice::Default, "default",
90 "Use standard heuristic")
91#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) \
92 , clEnumValN(EmitCModelChoice::CLASS_NAME, CLI_FLAG, \
93 "Use the " CLI_FLAG " MLGO model")
94#include "llvm/Analysis/InlinerModels.def"
95 ));
96
97static std::unique_ptr<MLModelRunner>
98createEmitCModelRunner(LLVMContext &Ctx,
99 const std::vector<TensorSpec> &InputFeatures) {
100 switch (SelectedMLGOModel) {
101 case EmitCModelChoice::Default:
102 return nullptr;
103#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) \
104 case EmitCModelChoice::CLASS_NAME: \
105 return std::make_unique<EmitCModelRunner<CLASS_NAME>>(Ctx, InputFeatures);
106#include "llvm/Analysis/InlinerModels.def"
107 }
108 llvm_unreachable("Unknown MLGO model type!");
109}
110#else
111constexpr bool HaveMLIRLoweringInliner = false;
112enum class EmitCModelChoice { Default };
113static const EmitCModelChoice SelectedMLGOModel = EmitCModelChoice::Default;
114static inline std::unique_ptr<MLModelRunner>
115createEmitCModelRunner(LLVMContext &, const std::vector<TensorSpec> &) {
116 return nullptr;
117}
118#endif
119
120std::unique_ptr<InlineAdvisor>
121llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM,
122 std::function<bool(CallBase &)> GetDefaultAdvice) {
123 if (!isReleaseModelValid<CompiledModelType>(InteractiveChannelBaseName,
124 SelectedModel: SelectedMLGOModel))
125 return nullptr;
126 auto RunnerFactory = [&](const std::vector<TensorSpec> &InputFeatures)
127 -> std::unique_ptr<MLModelRunner> {
128 return createReleaseModeModelRunner<CompiledModelType,
129 HaveMLIRLoweringInliner>(
130 Ctx&: M.getContext(), InputFeatures, DecisionName, InteractiveChannelBaseName,
131 InteractiveDecisionSpec: InlineDecisionSpec, CreateEmitCModelRunner&: createEmitCModelRunner,
132 Options: EmbeddedModelRunnerOptions().setModelSelector(ModelSelector));
133 };
134 return std::make_unique<MLInlineAdvisor>(args&: M, args&: MAM, args&: RunnerFactory,
135 args&: GetDefaultAdvice);
136}
137
138#define DEBUG_TYPE "inline-ml"
139
140static cl::opt<float> SizeIncreaseThreshold(
141 "ml-advisor-size-increase-threshold", cl::Hidden,
142 cl::desc("Maximum factor by which expected native size may increase before "
143 "blocking any further inlining."),
144 cl::init(Val: 2.0));
145
146static cl::opt<bool> KeepFPICache(
147 "ml-advisor-keep-fpi-cache", cl::Hidden,
148 cl::desc(
149 "For test - keep the ML Inline advisor's FunctionPropertiesInfo cache"),
150 cl::init(Val: false));
151
152const std::vector<TensorSpec> &MLInlineAdvisor::getInitialFeatureMap() {
153 // clang-format off
154static std::vector<TensorSpec> FeatureMap{
155#define POPULATE_NAMES(DTYPE, SHAPE, NAME, __) TensorSpec::createSpec<DTYPE>(#NAME, SHAPE),
156// InlineCost features - these must come first
157 INLINE_COST_FEATURE_ITERATOR(POPULATE_NAMES)
158
159// Non-cost features
160 INLINE_FEATURE_ITERATOR(POPULATE_NAMES)
161#undef POPULATE_NAMES
162};
163 // clang-format on
164 return FeatureMap;
165}
166
167const char *const llvm::DecisionName = "inlining_decision";
168const TensorSpec llvm::InlineDecisionSpec =
169 TensorSpec::createSpec<int64_t>(Name: DecisionName, Shape: {1});
170const char *const llvm::DefaultDecisionName = "inlining_default";
171const TensorSpec llvm::DefaultDecisionSpec =
172 TensorSpec::createSpec<int64_t>(Name: DefaultDecisionName, Shape: {1});
173const char *const llvm::RewardName = "delta_size";
174
175CallBase *getInlinableCS(Instruction &I) {
176 if (auto *CS = dyn_cast<CallBase>(Val: &I))
177 if (Function *Callee = CS->getCalledFunction()) {
178 if (!Callee->isDeclaration()) {
179 return CS;
180 }
181 }
182 return nullptr;
183}
184
185MLInlineAdvisor::MLInlineAdvisor(
186 Module &M, ModuleAnalysisManager &MAM,
187 std::function<
188 std::unique_ptr<MLModelRunner>(const std::vector<TensorSpec> &)>
189 GetModelRunner,
190 std::function<bool(CallBase &)> GetDefaultAdvice)
191 : InlineAdvisor(
192 M, MAM.getResult<FunctionAnalysisManagerModuleProxy>(IR&: M).getManager()),
193 GetDefaultAdvice(GetDefaultAdvice), FeatureMap(getInitialFeatureMap()),
194 CG(MAM.getResult<LazyCallGraphAnalysis>(IR&: M)),
195 UseIR2Vec(MAM.getCachedResult<IR2VecVocabAnalysis>(IR&: M) != nullptr),
196 InitialIRSize(getModuleIRSize()), CurrentIRSize(InitialIRSize),
197 PSI(MAM.getResult<ProfileSummaryAnalysis>(IR&: M)) {
198 // Extract the 'call site height' feature - the position of a call site
199 // relative to the farthest statically reachable SCC node. We don't mutate
200 // this value while inlining happens. Empirically, this feature proved
201 // critical in behavioral cloning - i.e. training a model to mimic the manual
202 // heuristic's decisions - and, thus, equally important for training for
203 // improvement.
204 CallGraph CGraph(M);
205 for (auto I = scc_begin(G: &CGraph); !I.isAtEnd(); ++I) {
206 const std::vector<CallGraphNode *> &CGNodes = *I;
207 unsigned Level = 0;
208 for (auto *CGNode : CGNodes) {
209 Function *F = CGNode->getFunction();
210 if (!F || F->isDeclaration())
211 continue;
212 for (auto &I : instructions(F)) {
213 if (auto *CS = getInlinableCS(I)) {
214 auto *Called = CS->getCalledFunction();
215 auto Pos = FunctionLevels.find(x: &CG.get(F&: *Called));
216 // In bottom up traversal, an inlinable callee is either in the
217 // same SCC, or to a function in a visited SCC. So not finding its
218 // level means we haven't visited it yet, meaning it's in this SCC.
219 if (Pos == FunctionLevels.end())
220 continue;
221 Level = std::max(a: Level, b: Pos->second + 1);
222 }
223 }
224 }
225 for (auto *CGNode : CGNodes) {
226 Function *F = CGNode->getFunction();
227 if (F && !F->isDeclaration())
228 FunctionLevels[&CG.get(F&: *F)] = Level;
229 }
230 }
231 for (auto KVP : FunctionLevels) {
232 AllNodes.insert(V: KVP.first);
233 EdgeCount += getLocalCalls(F&: KVP.first->getFunction());
234 }
235 NodeCount = AllNodes.size();
236
237 if (auto *IR2VecVocabResult = MAM.getCachedResult<IR2VecVocabAnalysis>(IR&: M)) {
238 if (!IR2VecVocabResult->isValid()) {
239 M.getContext().emitError(ErrorStr: "IR2VecVocabAnalysis is not valid");
240 return;
241 }
242 // Add the IR2Vec features to the feature map
243 auto IR2VecDim = IR2VecVocabResult->getDimension();
244 FeatureMap.push_back(
245 x: TensorSpec::createSpec<float>(Name: "callee_embedding", Shape: {IR2VecDim}));
246 FeatureMap.push_back(
247 x: TensorSpec::createSpec<float>(Name: "caller_embedding", Shape: {IR2VecDim}));
248 }
249 if (InteractiveIncludeDefault)
250 FeatureMap.push_back(x: DefaultDecisionSpec);
251
252 ModelRunner = GetModelRunner(getFeatureMap());
253 if (!ModelRunner) {
254 M.getContext().emitError(ErrorStr: "Could not create model runner");
255 return;
256 }
257 ModelRunner->switchContext(Name: "");
258 ForceStop = StopImmediatelyForTest;
259}
260
261unsigned MLInlineAdvisor::getInitialFunctionLevel(const Function &F) const {
262 return CG.lookup(F) ? FunctionLevels.at(k: CG.lookup(F)) : 0;
263}
264
265void MLInlineAdvisor::onPassEntry(LazyCallGraph::SCC *CurSCC) {
266 if (!CurSCC || ForceStop)
267 return;
268 FPICache.clear();
269 // Function passes executed between InlinerPass runs may have changed the
270 // module-wide features.
271 // The cgscc pass manager rules are such that:
272 // - if a pass leads to merging SCCs, then the pipeline is restarted on the
273 // merged SCC
274 // - if a pass leads to splitting the SCC, then we continue with one of the
275 // splits
276 // This means that the NodesInLastSCC is a superset (not strict) of the nodes
277 // that subsequent passes would have processed
278 // - in addition, if new Nodes were created by a pass (e.g. CoroSplit),
279 // they'd be adjacent to Nodes in the last SCC. So we just need to check the
280 // boundary of Nodes in NodesInLastSCC for Nodes we haven't seen. We don't
281 // care about the nature of the Edge (call or ref). `FunctionLevels`-wise, we
282 // record them at the same level as the original node (this is a choice, may
283 // need revisiting).
284 // - nodes are only deleted at the end of a call graph walk where they are
285 // batch deleted, so we shouldn't see any dead nodes here.
286 while (!NodesInLastSCC.empty()) {
287 const auto *N = *NodesInLastSCC.begin();
288 assert(!N->isDead());
289 NodesInLastSCC.erase(Ptr: N);
290 EdgeCount += getLocalCalls(F&: N->getFunction());
291 const auto NLevel = FunctionLevels.at(k: N);
292 for (const auto &E : *(*N)) {
293 const auto *AdjNode = &E.getNode();
294 assert(!AdjNode->isDead() && !AdjNode->getFunction().isDeclaration());
295 auto I = AllNodes.insert(V: AdjNode);
296 // We've discovered a new function.
297 if (I.second) {
298 ++NodeCount;
299 NodesInLastSCC.insert(Ptr: AdjNode);
300 FunctionLevels[AdjNode] = NLevel;
301 }
302 }
303 }
304
305 EdgeCount -= EdgesOfLastSeenNodes;
306 EdgesOfLastSeenNodes = 0;
307
308 // (Re)use NodesInLastSCC to remember the nodes in the SCC right now,
309 // in case the SCC is split before onPassExit and some nodes are split out
310 assert(NodesInLastSCC.empty());
311 for (const auto &N : *CurSCC)
312 NodesInLastSCC.insert(Ptr: &N);
313}
314
315void MLInlineAdvisor::onPassExit(LazyCallGraph::SCC *CurSCC) {
316 // No need to keep this around - function passes will invalidate it.
317 if (!KeepFPICache)
318 FPICache.clear();
319 if (!CurSCC || ForceStop)
320 return;
321 // Keep track of the nodes and edges we last saw. Then, in onPassEntry,
322 // we update the node count and edge count from the subset of these nodes that
323 // survived.
324 EdgesOfLastSeenNodes = 0;
325
326 // Check on nodes that were in SCC onPassEntry
327 for (const LazyCallGraph::Node *N : NodesInLastSCC) {
328 assert(!N->isDead());
329 EdgesOfLastSeenNodes += getLocalCalls(F&: N->getFunction());
330 }
331
332 // Check on nodes that may have got added to SCC
333 for (const auto &N : *CurSCC) {
334 assert(!N.isDead());
335 auto I = NodesInLastSCC.insert(Ptr: &N);
336 if (I.second)
337 EdgesOfLastSeenNodes += getLocalCalls(F&: N.getFunction());
338 }
339 assert(NodeCount >= NodesInLastSCC.size());
340 assert(EdgeCount >= EdgesOfLastSeenNodes);
341}
342
343int64_t MLInlineAdvisor::getLocalCalls(Function &F) {
344 return getCachedFPI(F).DirectCallsToDefinedFunctions;
345}
346
347// Update the internal state of the advisor, and force invalidate feature
348// analysis. Currently, we maintain minimal (and very simple) global state - the
349// number of functions and the number of static calls. We also keep track of the
350// total IR size in this module, to stop misbehaving policies at a certain bloat
351// factor (SizeIncreaseThreshold)
352void MLInlineAdvisor::onSuccessfulInlining(const MLInlineAdvice &Advice,
353 bool CalleeWasDeleted) {
354 assert(!ForceStop);
355 Function *Caller = Advice.getCaller();
356 Function *Callee = Advice.getCallee();
357 // The caller features aren't valid anymore.
358 {
359 PreservedAnalyses PA = PreservedAnalyses::all();
360 PA.abandon<FunctionPropertiesAnalysis>();
361 PA.abandon<LoopAnalysis>();
362 FAM.invalidate(IR&: *Caller, PA);
363 }
364 Advice.updateCachedCallerFPI(FAM);
365 if (Caller == Callee) {
366 assert(!CalleeWasDeleted);
367 // We double-counted CallerAndCalleeEdges - since the caller and callee
368 // would be the same
369 assert(Advice.CallerAndCalleeEdges % 2 == 0);
370 CurrentIRSize += getIRSize(F&: *Caller) - Advice.CallerIRSize;
371 EdgeCount += getCachedFPI(*Caller).DirectCallsToDefinedFunctions -
372 Advice.CallerAndCalleeEdges / 2;
373 // The NodeCount would stay the same.
374 } else {
375 int64_t IRSizeAfter =
376 getIRSize(F&: *Caller) + (CalleeWasDeleted ? 0 : Advice.CalleeIRSize);
377 CurrentIRSize += IRSizeAfter - (Advice.CallerIRSize + Advice.CalleeIRSize);
378
379 // We can delta-update module-wide features. We know the inlining only
380 // changed the caller, and maybe the callee (by deleting the latter). Nodes
381 // are simple to update. For edges, we 'forget' the edges that the caller
382 // and callee used to have before inlining, and add back what they currently
383 // have together.
384 int64_t NewCallerAndCalleeEdges =
385 getCachedFPI(*Caller).DirectCallsToDefinedFunctions;
386
387 // A dead function's node is not actually removed from the call graph until
388 // the end of the call graph walk, but the node no longer belongs to any
389 // valid SCC.
390 if (CalleeWasDeleted) {
391 --NodeCount;
392 NodesInLastSCC.erase(Ptr: CG.lookup(F: *Callee));
393 DeadFunctions.insert(V: Callee);
394 } else {
395 NewCallerAndCalleeEdges +=
396 getCachedFPI(*Callee).DirectCallsToDefinedFunctions;
397 }
398 EdgeCount += (NewCallerAndCalleeEdges - Advice.CallerAndCalleeEdges);
399 }
400 if (CurrentIRSize > SizeIncreaseThreshold * InitialIRSize)
401 ForceStop = true;
402
403 assert(CurrentIRSize >= 0 && EdgeCount >= 0 && NodeCount >= 0);
404}
405
406int64_t MLInlineAdvisor::getModuleIRSize() const {
407 int64_t Ret = 0;
408 for (auto &F : M)
409 if (!F.isDeclaration())
410 Ret += getIRSize(F);
411 return Ret;
412}
413
414FunctionPropertiesInfo &MLInlineAdvisor::getCachedFPI(Function &F) const {
415 auto InsertPair = FPICache.try_emplace(k: &F);
416 if (!InsertPair.second)
417 return InsertPair.first->second;
418 InsertPair.first->second = FAM.getResult<FunctionPropertiesAnalysis>(IR&: F);
419 return InsertPair.first->second;
420}
421
422std::unique_ptr<InlineAdvice> MLInlineAdvisor::getAdviceImpl(CallBase &CB) {
423 if (auto Skip = getSkipAdviceIfUnreachableCallsite(CB))
424 return Skip;
425
426 auto &Caller = *CB.getCaller();
427 auto &Callee = *CB.getCalledFunction();
428
429 auto GetAssumptionCache = [&](Function &F) -> AssumptionCache & {
430 return FAM.getResult<AssumptionAnalysis>(IR&: F);
431 };
432 auto &TIR = FAM.getResult<TargetIRAnalysis>(IR&: Callee);
433 auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(IR&: Caller);
434
435 if (SkipPolicy == SkipMLPolicyCriteria::IfCallerIsNotCold) {
436 if (!PSI.isFunctionEntryCold(F: &Caller)) {
437 // Return a MLInlineAdvice, despite delegating to the default advice,
438 // because we need to keep track of the internal state. This is different
439 // from the other instances where we return a "default" InlineAdvice,
440 // which happen at points we won't come back to the MLAdvisor for
441 // decisions requiring that state.
442 return ForceStop ? std::make_unique<InlineAdvice>(args: this, args&: CB, args&: ORE,
443 args: GetDefaultAdvice(CB))
444 : std::make_unique<MLInlineAdvice>(args: this, args&: CB, args&: ORE,
445 args: GetDefaultAdvice(CB));
446 }
447 }
448 auto MandatoryKind = InlineAdvisor::getMandatoryKind(CB, FAM, ORE);
449 // If this is a "never inline" case, there won't be any changes to internal
450 // state we need to track, so we can just return the base InlineAdvice, which
451 // will do nothing interesting.
452 // Same thing if this is a recursive case.
453 if (MandatoryKind == InlineAdvisor::MandatoryInliningKind::Never ||
454 &Caller == &Callee)
455 return getMandatoryAdvice(CB, Advice: false);
456
457 bool Mandatory =
458 MandatoryKind == InlineAdvisor::MandatoryInliningKind::Always;
459
460 // If we need to stop, we won't want to track anymore any state changes, so
461 // we just return the base InlineAdvice, which acts as a noop.
462 if (ForceStop) {
463 ORE.emit(RemarkBuilder: [&] {
464 return OptimizationRemarkMissed(DEBUG_TYPE, "ForceStop", &CB)
465 << "Won't attempt inlining because module size grew too much.";
466 });
467 return std::make_unique<InlineAdvice>(args: this, args&: CB, args&: ORE, args&: Mandatory);
468 }
469
470 int CostEstimate = 0;
471 if (!Mandatory) {
472 auto IsCallSiteInlinable =
473 llvm::getInliningCostEstimate(Call&: CB, CalleeTTI&: TIR, GetAssumptionCache);
474 if (!IsCallSiteInlinable) {
475 // We can't inline this for correctness reasons, so return the base
476 // InlineAdvice, as we don't care about tracking any state changes (which
477 // won't happen).
478 return std::make_unique<InlineAdvice>(args: this, args&: CB, args&: ORE, args: false);
479 }
480 CostEstimate = *IsCallSiteInlinable;
481 }
482
483 const auto CostFeatures =
484 llvm::getInliningCostFeatures(Call&: CB, CalleeTTI&: TIR, GetAssumptionCache);
485 if (!CostFeatures) {
486 return std::make_unique<InlineAdvice>(args: this, args&: CB, args&: ORE, args: false);
487 }
488
489 if (Mandatory)
490 return getMandatoryAdvice(CB, Advice: true);
491
492 auto NumCtantParams = 0;
493 for (auto I = CB.arg_begin(), E = CB.arg_end(); I != E; ++I) {
494 NumCtantParams += (isa<Constant>(Val: *I));
495 }
496
497 auto &CallerBefore = getCachedFPI(F&: Caller);
498 auto &CalleeBefore = getCachedFPI(F&: Callee);
499
500 *ModelRunner->getTensor<int64_t>(FeatureID: FeatureIndex::callee_basic_block_count) =
501 CalleeBefore.BasicBlockCount;
502 *ModelRunner->getTensor<int64_t>(FeatureID: FeatureIndex::callsite_height) =
503 getInitialFunctionLevel(F: Caller);
504 *ModelRunner->getTensor<int64_t>(FeatureID: FeatureIndex::node_count) = NodeCount;
505 *ModelRunner->getTensor<int64_t>(FeatureID: FeatureIndex::nr_ctant_params) =
506 NumCtantParams;
507 *ModelRunner->getTensor<int64_t>(FeatureID: FeatureIndex::edge_count) = EdgeCount;
508 *ModelRunner->getTensor<int64_t>(FeatureID: FeatureIndex::caller_users) =
509 CallerBefore.Uses;
510 *ModelRunner->getTensor<int64_t>(
511 FeatureID: FeatureIndex::caller_conditionally_executed_blocks) =
512 CallerBefore.BlocksReachedFromConditionalInstruction;
513 *ModelRunner->getTensor<int64_t>(FeatureID: FeatureIndex::caller_basic_block_count) =
514 CallerBefore.BasicBlockCount;
515 *ModelRunner->getTensor<int64_t>(
516 FeatureID: FeatureIndex::callee_conditionally_executed_blocks) =
517 CalleeBefore.BlocksReachedFromConditionalInstruction;
518 *ModelRunner->getTensor<int64_t>(FeatureID: FeatureIndex::callee_users) =
519 CalleeBefore.Uses;
520 *ModelRunner->getTensor<int64_t>(FeatureID: FeatureIndex::cost_estimate) = CostEstimate;
521 *ModelRunner->getTensor<int64_t>(FeatureID: FeatureIndex::is_callee_avail_external) =
522 Callee.hasAvailableExternallyLinkage();
523 *ModelRunner->getTensor<int64_t>(FeatureID: FeatureIndex::is_caller_avail_external) =
524 Caller.hasAvailableExternallyLinkage();
525
526 if (UseIR2Vec) {
527 // Python side expects float embeddings. The IR2Vec embeddings are doubles
528 // as of now due to the restriction of fromJSON method used by the
529 // readVocabulary method in ir2vec::Embeddings.
530 auto setEmbedding = [&](const ir2vec::Embedding &Embedding,
531 FeatureIndex Index) {
532 llvm::transform(Range: Embedding, d_first: ModelRunner->getTensor<float>(FeatureID: Index),
533 F: [](double Val) { return static_cast<float>(Val); });
534 };
535
536 setEmbedding(CalleeBefore.getFunctionEmbedding(),
537 FeatureIndex::callee_embedding);
538 setEmbedding(CallerBefore.getFunctionEmbedding(),
539 FeatureIndex::caller_embedding);
540 }
541
542 // Add the cost features
543 for (size_t I = 0;
544 I < static_cast<size_t>(InlineCostFeatureIndex::NumberOfFeatures); ++I) {
545 *ModelRunner->getTensor<int64_t>(FeatureID: inlineCostFeatureToMlFeature(
546 Feature: static_cast<InlineCostFeatureIndex>(I))) = CostFeatures->at(n: I);
547 }
548 // This one would have been set up to be right at the end.
549 if (!InteractiveChannelBaseName.empty() && InteractiveIncludeDefault)
550 *ModelRunner->getTensor<int64_t>(FeatureID: getFeatureMap().size() - 1) =
551 GetDefaultAdvice(CB);
552 return getAdviceFromModel(CB, ORE);
553}
554
555std::unique_ptr<MLInlineAdvice>
556MLInlineAdvisor::getAdviceFromModel(CallBase &CB,
557 OptimizationRemarkEmitter &ORE) {
558 return std::make_unique<MLInlineAdvice>(
559 args: this, args&: CB, args&: ORE, args: static_cast<bool>(ModelRunner->evaluate<int64_t>()));
560}
561
562std::unique_ptr<InlineAdvice>
563MLInlineAdvisor::getSkipAdviceIfUnreachableCallsite(CallBase &CB) {
564 if (!FAM.getResult<DominatorTreeAnalysis>(IR&: *CB.getCaller())
565 .isReachableFromEntry(A: CB.getParent()))
566 return std::make_unique<InlineAdvice>(args: this, args&: CB, args&: getCallerORE(CB), args: false);
567 return nullptr;
568}
569
570std::unique_ptr<InlineAdvice> MLInlineAdvisor::getMandatoryAdvice(CallBase &CB,
571 bool Advice) {
572 // Make sure we track inlinings in all cases - mandatory or not.
573 if (auto Skip = getSkipAdviceIfUnreachableCallsite(CB))
574 return Skip;
575 if (Advice && !ForceStop)
576 return getMandatoryAdviceImpl(CB);
577
578 // If this is a "never inline" case, there won't be any changes to internal
579 // state we need to track, so we can just return the base InlineAdvice, which
580 // will do nothing interesting.
581 // Same if we are forced to stop - we don't track anymore.
582 return std::make_unique<InlineAdvice>(args: this, args&: CB, args&: getCallerORE(CB), args&: Advice);
583}
584
585std::unique_ptr<MLInlineAdvice>
586MLInlineAdvisor::getMandatoryAdviceImpl(CallBase &CB) {
587 return std::make_unique<MLInlineAdvice>(args: this, args&: CB, args&: getCallerORE(CB), args: true);
588}
589
590void MLInlineAdvisor::print(raw_ostream &OS) const {
591 OS << "[MLInlineAdvisor] Nodes: " << NodeCount << " Edges: " << EdgeCount
592 << " EdgesOfLastSeenNodes: " << EdgesOfLastSeenNodes << "\n";
593 OS << "[MLInlineAdvisor] FPI:\n";
594 for (auto I : FPICache) {
595 OS << I.first->getName() << ":\n";
596 I.second.print(OS);
597 OS << "\n";
598 }
599 OS << "\n";
600 OS << "[MLInlineAdvisor] FuncLevels:\n";
601 for (auto I : FunctionLevels)
602 OS << (DeadFunctions.contains(V: &I.first->getFunction())
603 ? "<deleted>"
604 : I.first->getFunction().getName())
605 << " : " << I.second << "\n";
606
607 OS << "\n";
608}
609
610MLInlineAdvice::MLInlineAdvice(MLInlineAdvisor *Advisor, CallBase &CB,
611 OptimizationRemarkEmitter &ORE,
612 bool Recommendation)
613 : InlineAdvice(Advisor, CB, ORE, Recommendation),
614 CallerIRSize(Advisor->isForcedToStop() ? 0 : Advisor->getIRSize(F&: *Caller)),
615 CalleeIRSize(Advisor->isForcedToStop() ? 0 : Advisor->getIRSize(F&: *Callee)),
616 CallerAndCalleeEdges(Advisor->isForcedToStop()
617 ? 0
618 : (Advisor->getLocalCalls(F&: *Caller) +
619 Advisor->getLocalCalls(F&: *Callee))),
620 PreInlineCallerFPI(Advisor->getCachedFPI(F&: *Caller)) {
621 if (Recommendation)
622 FPU.emplace(args&: Advisor->getCachedFPI(F&: *getCaller()), args&: CB);
623}
624
625void MLInlineAdvice::reportContextForRemark(
626 DiagnosticInfoOptimizationBase &OR) {
627 using namespace ore;
628 OR << NV("Callee", Callee->getName());
629 for (size_t I = 0; I < getAdvisor()->getFeatureMap().size(); ++I)
630 OR << NV(getAdvisor()->getFeatureMap()[I].name(),
631 *getAdvisor()->getModelRunner().getTensor<int64_t>(FeatureID: I));
632 OR << NV("ShouldInline", isInliningRecommended());
633}
634
635void MLInlineAdvice::updateCachedCallerFPI(FunctionAnalysisManager &FAM) const {
636 FPU->finish(FAM);
637}
638
639void MLInlineAdvice::recordInliningImpl() {
640 ORE.emit(RemarkBuilder: [&]() {
641 OptimizationRemark R(DEBUG_TYPE, "InliningSuccess", DLoc, Block);
642 reportContextForRemark(OR&: R);
643 return R;
644 });
645 getAdvisor()->onSuccessfulInlining(Advice: *this, /*CalleeWasDeleted*/ false);
646}
647
648void MLInlineAdvice::recordInliningWithCalleeDeletedImpl() {
649 ORE.emit(RemarkBuilder: [&]() {
650 OptimizationRemark R(DEBUG_TYPE, "InliningSuccessWithCalleeDeleted", DLoc,
651 Block);
652 reportContextForRemark(OR&: R);
653 return R;
654 });
655 getAdvisor()->onSuccessfulInlining(Advice: *this, /*CalleeWasDeleted*/ true);
656}
657
658void MLInlineAdvice::recordUnsuccessfulInliningImpl(
659 const InlineResult &Result) {
660 getAdvisor()->getCachedFPI(F&: *Caller) = PreInlineCallerFPI;
661 ORE.emit(RemarkBuilder: [&]() {
662 OptimizationRemarkMissed R(DEBUG_TYPE, "InliningAttemptedAndUnsuccessful",
663 DLoc, Block);
664 reportContextForRemark(OR&: R);
665 return R;
666 });
667}
668void MLInlineAdvice::recordUnattemptedInliningImpl() {
669 assert(!FPU);
670 ORE.emit(RemarkBuilder: [&]() {
671 OptimizationRemarkMissed R(DEBUG_TYPE, "IniningNotAttempted", DLoc, Block);
672 reportContextForRemark(OR&: R);
673 return R;
674 });
675}
676