llvm/Analysis/MLInlineAdvisor.h

0001 //===- MLInlineAdvisor.h - ML - based InlineAdvisor factories ---*- C++ -*-===//
0002 //
0003 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0004 // See https://llvm.org/LICENSE.txt for license information.
0005 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0006 //
0007 //===----------------------------------------------------------------------===//
0008
0009 #ifndef LLVM_ANALYSIS_MLINLINEADVISOR_H
0010 #define LLVM_ANALYSIS_MLINLINEADVISOR_H
0011
0012 #include "llvm/Analysis/FunctionPropertiesAnalysis.h"
0013 #include "llvm/Analysis/InlineAdvisor.h"
0014 #include "llvm/Analysis/LazyCallGraph.h"
0015 #include "llvm/Analysis/MLModelRunner.h"
0016 #include "llvm/IR/PassManager.h"
0017
0018 #include <map>
0019 #include <memory>
0020 #include <optional>
0021
0022 namespace llvm {
0023 class DiagnosticInfoOptimizationBase;
0024 class Module;
0025 class MLInlineAdvice;
0026 class ProfileSummaryInfo;
0027
0028 class MLInlineAdvisor : public InlineAdvisor {
0029 public:
0030   MLInlineAdvisor(Module &M, ModuleAnalysisManager &MAM,
0031                   std::unique_ptr<MLModelRunner> ModelRunner,
0032                   std::function<bool(CallBase &)> GetDefaultAdvice);
0033
0034   virtual ~MLInlineAdvisor() = default;
0035
0036   void onPassEntry(LazyCallGraph::SCC *SCC) override;
0037   void onPassExit(LazyCallGraph::SCC *SCC) override;
0038
0039   int64_t getIRSize(Function &F) const {
0040     return getCachedFPI(F).TotalInstructionCount;
0041   }
0042   void onSuccessfulInlining(const MLInlineAdvice &Advice,
0043                             bool CalleeWasDeleted);
0044
0045   bool isForcedToStop() const { return ForceStop; }
0046   int64_t getLocalCalls(Function &F);
0047   const MLModelRunner &getModelRunner() const { return *ModelRunner; }
0048   FunctionPropertiesInfo &getCachedFPI(Function &) const;
0049
0050 protected:
0051   std::unique_ptr<InlineAdvice> getAdviceImpl(CallBase &CB) override;
0052
0053   std::unique_ptr<InlineAdvice> getMandatoryAdvice(CallBase &CB,
0054                                                    bool Advice) override;
0055
0056   virtual std::unique_ptr<MLInlineAdvice> getMandatoryAdviceImpl(CallBase &CB);
0057
0058   virtual std::unique_ptr<MLInlineAdvice>
0059   getAdviceFromModel(CallBase &CB, OptimizationRemarkEmitter &ORE);
0060
0061   // Get the initial 'level' of the function, or 0 if the function has been
0062   // introduced afterwards.
0063   // TODO: should we keep this updated?
0064   unsigned getInitialFunctionLevel(const Function &F) const;
0065
0066   std::unique_ptr<MLModelRunner> ModelRunner;
0067   std::function<bool(CallBase &)> GetDefaultAdvice;
0068
0069 private:
0070   int64_t getModuleIRSize() const;
0071   std::unique_ptr<InlineAdvice>
0072   getSkipAdviceIfUnreachableCallsite(CallBase &CB);
0073   void print(raw_ostream &OS) const override;
0074
0075   // Using std::map to benefit from its iterator / reference non-invalidating
0076   // semantics, which make it easy to use `getCachedFPI` results from multiple
0077   // calls without needing to copy to avoid invalidation effects.
0078   mutable std::map<const Function *, FunctionPropertiesInfo> FPICache;
0079
0080   LazyCallGraph &CG;
0081
0082   int64_t NodeCount = 0;
0083   int64_t EdgeCount = 0;
0084   int64_t EdgesOfLastSeenNodes = 0;
0085
0086   std::map<const LazyCallGraph::Node *, unsigned> FunctionLevels;
0087   const int32_t InitialIRSize = 0;
0088   int32_t CurrentIRSize = 0;
0089   llvm::SmallPtrSet<const LazyCallGraph::Node *, 1> NodesInLastSCC;
0090   DenseSet<const LazyCallGraph::Node *> AllNodes;
0091   DenseSet<Function *> DeadFunctions;
0092   bool ForceStop = false;
0093   ProfileSummaryInfo &PSI;
0094 };
0095
0096 /// InlineAdvice that tracks changes post inlining. For that reason, it only
0097 /// overrides the "successful inlining" extension points.
0098 class MLInlineAdvice : public InlineAdvice {
0099 public:
0100   MLInlineAdvice(MLInlineAdvisor *Advisor, CallBase &CB,
0101                  OptimizationRemarkEmitter &ORE, bool Recommendation);
0102   virtual ~MLInlineAdvice() = default;
0103
0104   void recordInliningImpl() override;
0105   void recordInliningWithCalleeDeletedImpl() override;
0106   void recordUnsuccessfulInliningImpl(const InlineResult &Result) override;
0107   void recordUnattemptedInliningImpl() override;
0108
0109   Function *getCaller() const { return Caller; }
0110   Function *getCallee() const { return Callee; }
0111
0112   const int64_t CallerIRSize;
0113   const int64_t CalleeIRSize;
0114   const int64_t CallerAndCalleeEdges;
0115   void updateCachedCallerFPI(FunctionAnalysisManager &FAM) const;
0116
0117 private:
0118   void reportContextForRemark(DiagnosticInfoOptimizationBase &OR);
0119   MLInlineAdvisor *getAdvisor() const {
0120     return static_cast<MLInlineAdvisor *>(Advisor);
0121   };
0122   // Make a copy of the FPI of the caller right before inlining. If inlining
0123   // fails, we can just update the cache with that value.
0124   const FunctionPropertiesInfo PreInlineCallerFPI;
0125   std::optional<FunctionPropertiesUpdater> FPU;
0126 };
0127
0128 } // namespace llvm
0129
0130 #endif // LLVM_ANALYSIS_MLINLINEADVISOR_H