llvm/lib/Analysis/MLInlineAdvisor.cpp

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