1 //===- MLInlineAdvisor.h - ML - based InlineAdvisor factories ---*- C++ -*-===//
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 #ifndef LLVM_ANALYSIS_MLINLINEADVISOR_H
10 #define LLVM_ANALYSIS_MLINLINEADVISOR_H
11 
12 #include "llvm/Analysis/FunctionPropertiesAnalysis.h"
13 #include "llvm/Analysis/InlineAdvisor.h"
14 #include "llvm/Analysis/LazyCallGraph.h"
15 #include "llvm/Analysis/MLModelRunner.h"
16 #include "llvm/IR/PassManager.h"
17 
18 #include <deque>
19 #include <map>
20 #include <memory>
21 
22 namespace llvm {
23 class DiagnosticInfoOptimizationBase;
24 class Module;
25 class MLInlineAdvice;
26 
27 class MLInlineAdvisor : public InlineAdvisor {
28 public:
29   MLInlineAdvisor(Module &M, ModuleAnalysisManager &MAM,
30                   std::unique_ptr<MLModelRunner> ModelRunner);
31 
32   virtual ~MLInlineAdvisor() = default;
33 
34   void onPassEntry(LazyCallGraph::SCC *SCC) override;
35   void onPassExit(LazyCallGraph::SCC *SCC) override;
36 
37   int64_t getIRSize(Function &F) const {
38     return getCachedFPI(F).TotalInstructionCount;
39   }
40   void onSuccessfulInlining(const MLInlineAdvice &Advice,
41                             bool CalleeWasDeleted);
42 
43   bool isForcedToStop() const { return ForceStop; }
44   int64_t getLocalCalls(Function &F);
45   const MLModelRunner &getModelRunner() const { return *ModelRunner.get(); }
46   FunctionPropertiesInfo &getCachedFPI(Function &) const;
47 
48 protected:
49   std::unique_ptr<InlineAdvice> getAdviceImpl(CallBase &CB) override;
50 
51   std::unique_ptr<InlineAdvice> getMandatoryAdvice(CallBase &CB,
52                                                    bool Advice) override;
53 
54   virtual std::unique_ptr<MLInlineAdvice> getMandatoryAdviceImpl(CallBase &CB);
55 
56   virtual std::unique_ptr<MLInlineAdvice>
57   getAdviceFromModel(CallBase &CB, OptimizationRemarkEmitter &ORE);
58 
59   // Get the initial 'level' of the function, or 0 if the function has been
60   // introduced afterwards.
61   // TODO: should we keep this updated?
62   unsigned getInitialFunctionLevel(const Function &F) const;
63 
64   std::unique_ptr<MLModelRunner> ModelRunner;
65 
66 private:
67   int64_t getModuleIRSize() const;
68   std::unique_ptr<InlineAdvice>
69   getSkipAdviceIfUnreachableCallsite(CallBase &CB);
70   void print(raw_ostream &OS) const override;
71 
72   mutable DenseMap<const Function *, FunctionPropertiesInfo> FPICache;
73 
74   LazyCallGraph &CG;
75 
76   int64_t NodeCount = 0;
77   int64_t EdgeCount = 0;
78   int64_t EdgesOfLastSeenNodes = 0;
79 
80   std::map<const LazyCallGraph::Node *, unsigned> FunctionLevels;
81   const int32_t InitialIRSize = 0;
82   int32_t CurrentIRSize = 0;
83   llvm::SmallPtrSet<const LazyCallGraph::Node *, 1> NodesInLastSCC;
84   DenseSet<const LazyCallGraph::Node *> AllNodes;
85   bool ForceStop = false;
86 };
87 
88 /// InlineAdvice that tracks changes post inlining. For that reason, it only
89 /// overrides the "successful inlining" extension points.
90 class MLInlineAdvice : public InlineAdvice {
91 public:
92   MLInlineAdvice(MLInlineAdvisor *Advisor, CallBase &CB,
93                  OptimizationRemarkEmitter &ORE, bool Recommendation);
94   virtual ~MLInlineAdvice() = default;
95 
96   void recordInliningImpl() override;
97   void recordInliningWithCalleeDeletedImpl() override;
98   void recordUnsuccessfulInliningImpl(const InlineResult &Result) override;
99   void recordUnattemptedInliningImpl() override;
100 
101   Function *getCaller() const { return Caller; }
102   Function *getCallee() const { return Callee; }
103 
104   const int64_t CallerIRSize;
105   const int64_t CalleeIRSize;
106   const int64_t CallerAndCalleeEdges;
107   void updateCachedCallerFPI(FunctionAnalysisManager &FAM) const;
108 
109 private:
110   void reportContextForRemark(DiagnosticInfoOptimizationBase &OR);
111   MLInlineAdvisor *getAdvisor() const {
112     return static_cast<MLInlineAdvisor *>(Advisor);
113   };
114   // Make a copy of the FPI of the caller right before inlining. If inlining
115   // fails, we can just update the cache with that value.
116   const FunctionPropertiesInfo PreInlineCallerFPI;
117   Optional<FunctionPropertiesUpdater> FPU;
118 };
119 
120 } // namespace llvm
121 
122 #endif // LLVM_ANALYSIS_MLINLINEADVISOR_H
123