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