| //===- MLInlineAdvisor.h - ML - based InlineAdvisor factories ---*- C++ -*-===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #ifndef LLVM_ANALYSIS_MLINLINEADVISOR_H |
| #define LLVM_ANALYSIS_MLINLINEADVISOR_H |
| |
| #include "llvm/Analysis/CallGraph.h" |
| #include "llvm/Analysis/InlineAdvisor.h" |
| #include "llvm/Analysis/MLModelRunner.h" |
| #include "llvm/IR/PassManager.h" |
| |
| #include <memory> |
| #include <unordered_map> |
| |
| namespace llvm { |
| class Module; |
| class MLInlineAdvice; |
| |
| class MLInlineAdvisor : public InlineAdvisor { |
| public: |
| MLInlineAdvisor(Module &M, ModuleAnalysisManager &MAM, |
| std::unique_ptr<MLModelRunner> ModelRunner); |
| |
| CallGraph *callGraph() const { return CG.get(); } |
| virtual ~MLInlineAdvisor() = default; |
| |
| void onPassEntry() override; |
| |
| int64_t getIRSize(const Function &F) const { return F.getInstructionCount(); } |
| void onSuccessfulInlining(const MLInlineAdvice &Advice, |
| bool CalleeWasDeleted); |
| |
| bool isForcedToStop() const { return ForceStop; } |
| int64_t getLocalCalls(Function &F); |
| const MLModelRunner &getModelRunner() const { return *ModelRunner.get(); } |
| void onModuleInvalidated() override { Invalid = true; } |
| |
| protected: |
| std::unique_ptr<InlineAdvice> getAdviceImpl(CallBase &CB) override; |
| |
| std::unique_ptr<InlineAdvice> getMandatoryAdvice(CallBase &CB, |
| bool Advice) override; |
| |
| virtual std::unique_ptr<MLInlineAdvice> getMandatoryAdviceImpl(CallBase &CB); |
| |
| virtual std::unique_ptr<MLInlineAdvice> |
| getAdviceFromModel(CallBase &CB, OptimizationRemarkEmitter &ORE); |
| |
| std::unique_ptr<MLModelRunner> ModelRunner; |
| |
| private: |
| int64_t getModuleIRSize() const; |
| |
| bool Invalid = true; |
| std::unique_ptr<CallGraph> CG; |
| |
| int64_t NodeCount = 0; |
| int64_t EdgeCount = 0; |
| std::map<const Function *, unsigned> FunctionLevels; |
| const int32_t InitialIRSize = 0; |
| int32_t CurrentIRSize = 0; |
| |
| bool ForceStop = false; |
| }; |
| |
| /// InlineAdvice that tracks changes post inlining. For that reason, it only |
| /// overrides the "successful inlining" extension points. |
| class MLInlineAdvice : public InlineAdvice { |
| public: |
| MLInlineAdvice(MLInlineAdvisor *Advisor, CallBase &CB, |
| OptimizationRemarkEmitter &ORE, bool Recommendation) |
| : InlineAdvice(Advisor, CB, ORE, Recommendation), |
| CallerIRSize(Advisor->isForcedToStop() ? 0 |
| : Advisor->getIRSize(*Caller)), |
| CalleeIRSize(Advisor->isForcedToStop() ? 0 |
| : Advisor->getIRSize(*Callee)), |
| CallerAndCalleeEdges(Advisor->isForcedToStop() |
| ? 0 |
| : (Advisor->getLocalCalls(*Caller) + |
| Advisor->getLocalCalls(*Callee))) {} |
| virtual ~MLInlineAdvice() = default; |
| |
| void recordInliningImpl() override; |
| void recordInliningWithCalleeDeletedImpl() override; |
| void recordUnsuccessfulInliningImpl(const InlineResult &Result) override; |
| void recordUnattemptedInliningImpl() override; |
| |
| Function *getCaller() const { return Caller; } |
| Function *getCallee() const { return Callee; } |
| |
| const int64_t CallerIRSize; |
| const int64_t CalleeIRSize; |
| const int64_t CallerAndCalleeEdges; |
| |
| private: |
| void reportContextForRemark(DiagnosticInfoOptimizationBase &OR); |
| |
| MLInlineAdvisor *getAdvisor() const { |
| return static_cast<MLInlineAdvisor *>(Advisor); |
| }; |
| }; |
| |
| } // namespace llvm |
| |
| #endif // LLVM_ANALYSIS_MLINLINEADVISOR_H |