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 if (!Invalid) 120 return; 121 NodeCount = 0; 122 EdgeCount = 0; 123 for (auto &F : M) 124 if (!F.isDeclaration()) { 125 ++NodeCount; 126 EdgeCount += getLocalCalls(F); 127 } 128 Invalid = false; 129 } 130 131 int64_t MLInlineAdvisor::getLocalCalls(Function &F) { 132 return FAM.getResult<FunctionPropertiesAnalysis>(F) 133 .DirectCallsToDefinedFunctions; 134 } 135 136 // Update the internal state of the advisor, and force invalidate feature 137 // analysis. Currently, we maintain minimal (and very simple) global state - the 138 // number of functions and the number of static calls. We also keep track of the 139 // total IR size in this module, to stop misbehaving policies at a certain bloat 140 // factor (SizeIncreaseThreshold) 141 void MLInlineAdvisor::onSuccessfulInlining(const MLInlineAdvice &Advice, 142 bool CalleeWasDeleted) { 143 assert(!ForceStop); 144 Function *Caller = Advice.getCaller(); 145 Function *Callee = Advice.getCallee(); 146 147 // The caller features aren't valid anymore. 148 { 149 PreservedAnalyses PA = PreservedAnalyses::all(); 150 PA.abandon<FunctionPropertiesAnalysis>(); 151 FAM.invalidate(*Caller, PA); 152 } 153 int64_t IRSizeAfter = 154 getIRSize(*Caller) + (CalleeWasDeleted ? 0 : Advice.CalleeIRSize); 155 CurrentIRSize += IRSizeAfter - (Advice.CallerIRSize + Advice.CalleeIRSize); 156 if (CurrentIRSize > SizeIncreaseThreshold * InitialIRSize) 157 ForceStop = true; 158 159 // We can delta-update module-wide features. We know the inlining only changed 160 // the caller, and maybe the callee (by deleting the latter). 161 // Nodes are simple to update. 162 // For edges, we 'forget' the edges that the caller and callee used to have 163 // before inlining, and add back what they currently have together. 164 int64_t NewCallerAndCalleeEdges = 165 FAM.getResult<FunctionPropertiesAnalysis>(*Caller) 166 .DirectCallsToDefinedFunctions; 167 168 if (CalleeWasDeleted) 169 --NodeCount; 170 else 171 NewCallerAndCalleeEdges += 172 FAM.getResult<FunctionPropertiesAnalysis>(*Callee) 173 .DirectCallsToDefinedFunctions; 174 EdgeCount += (NewCallerAndCalleeEdges - Advice.CallerAndCalleeEdges); 175 assert(CurrentIRSize >= 0 && EdgeCount >= 0 && NodeCount >= 0); 176 } 177 178 int64_t MLInlineAdvisor::getModuleIRSize() const { 179 int64_t Ret = 0; 180 for (auto &F : CG->getModule()) 181 if (!F.isDeclaration()) 182 Ret += getIRSize(F); 183 return Ret; 184 } 185 186 std::unique_ptr<InlineAdvice> MLInlineAdvisor::getAdviceImpl(CallBase &CB) { 187 auto &Caller = *CB.getCaller(); 188 auto &Callee = *CB.getCalledFunction(); 189 190 auto GetAssumptionCache = [&](Function &F) -> AssumptionCache & { 191 return FAM.getResult<AssumptionAnalysis>(F); 192 }; 193 auto &TIR = FAM.getResult<TargetIRAnalysis>(Callee); 194 auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(Caller); 195 196 auto MandatoryKind = InlineAdvisor::getMandatoryKind(CB, FAM, ORE); 197 // If this is a "never inline" case, there won't be any changes to internal 198 // state we need to track, so we can just return the base InlineAdvice, which 199 // will do nothing interesting. 200 // Same thing if this is a recursive case. 201 if (MandatoryKind == InlineAdvisor::MandatoryInliningKind::Never || 202 &Caller == &Callee) 203 return getMandatoryAdvice(CB, false); 204 205 bool Mandatory = 206 MandatoryKind == InlineAdvisor::MandatoryInliningKind::Always; 207 208 // If we need to stop, we won't want to track anymore any state changes, so 209 // we just return the base InlineAdvice, which acts as a noop. 210 if (ForceStop) { 211 ORE.emit([&] { 212 return OptimizationRemarkMissed(DEBUG_TYPE, "ForceStop", &CB) 213 << "Won't attempt inlining because module size grew too much."; 214 }); 215 return std::make_unique<InlineAdvice>(this, CB, ORE, Mandatory); 216 } 217 218 int CostEstimate = 0; 219 if (!Mandatory) { 220 auto IsCallSiteInlinable = 221 llvm::getInliningCostEstimate(CB, TIR, GetAssumptionCache); 222 if (!IsCallSiteInlinable) { 223 // We can't inline this for correctness reasons, so return the base 224 // InlineAdvice, as we don't care about tracking any state changes (which 225 // won't happen). 226 return std::make_unique<InlineAdvice>(this, CB, ORE, false); 227 } 228 CostEstimate = *IsCallSiteInlinable; 229 } 230 231 const auto CostFeatures = 232 llvm::getInliningCostFeatures(CB, TIR, GetAssumptionCache); 233 if (!CostFeatures) { 234 return std::make_unique<InlineAdvice>(this, CB, ORE, false); 235 } 236 237 if (Mandatory) 238 return getMandatoryAdvice(CB, true); 239 240 auto NrCtantParams = 0; 241 for (auto I = CB.arg_begin(), E = CB.arg_end(); I != E; ++I) { 242 NrCtantParams += (isa<Constant>(*I)); 243 } 244 245 auto &CallerBefore = FAM.getResult<FunctionPropertiesAnalysis>(Caller); 246 auto &CalleeBefore = FAM.getResult<FunctionPropertiesAnalysis>(Callee); 247 248 ModelRunner->setFeature(FeatureIndex::CalleeBasicBlockCount, 249 CalleeBefore.BasicBlockCount); 250 ModelRunner->setFeature(FeatureIndex::CallSiteHeight, 251 FunctionLevels[&Caller]); 252 ModelRunner->setFeature(FeatureIndex::NodeCount, NodeCount); 253 ModelRunner->setFeature(FeatureIndex::NrCtantParams, NrCtantParams); 254 ModelRunner->setFeature(FeatureIndex::EdgeCount, EdgeCount); 255 ModelRunner->setFeature(FeatureIndex::CallerUsers, CallerBefore.Uses); 256 ModelRunner->setFeature(FeatureIndex::CallerConditionallyExecutedBlocks, 257 CallerBefore.BlocksReachedFromConditionalInstruction); 258 ModelRunner->setFeature(FeatureIndex::CallerBasicBlockCount, 259 CallerBefore.BasicBlockCount); 260 ModelRunner->setFeature(FeatureIndex::CalleeConditionallyExecutedBlocks, 261 CalleeBefore.BlocksReachedFromConditionalInstruction); 262 ModelRunner->setFeature(FeatureIndex::CalleeUsers, CalleeBefore.Uses); 263 ModelRunner->setFeature(FeatureIndex::CostEstimate, CostEstimate); 264 265 // Add the cost features 266 for (size_t I = 0; 267 I < static_cast<size_t>(InlineCostFeatureIndex::NumberOfFeatures); ++I) { 268 ModelRunner->setFeature( 269 inlineCostFeatureToMlFeature(static_cast<InlineCostFeatureIndex>(I)), 270 CostFeatures->at(I)); 271 } 272 273 return getAdviceFromModel(CB, ORE); 274 } 275 276 std::unique_ptr<MLInlineAdvice> 277 MLInlineAdvisor::getAdviceFromModel(CallBase &CB, 278 OptimizationRemarkEmitter &ORE) { 279 return std::make_unique<MLInlineAdvice>(this, CB, ORE, ModelRunner->run()); 280 } 281 282 std::unique_ptr<InlineAdvice> MLInlineAdvisor::getMandatoryAdvice(CallBase &CB, 283 bool Advice) { 284 // Make sure we track inlinings in all cases - mandatory or not. 285 if (Advice && !ForceStop) 286 return getMandatoryAdviceImpl(CB); 287 288 // If this is a "never inline" case, there won't be any changes to internal 289 // state we need to track, so we can just return the base InlineAdvice, which 290 // will do nothing interesting. 291 // Same if we are forced to stop - we don't track anymore. 292 return std::make_unique<InlineAdvice>(this, CB, getCallerORE(CB), Advice); 293 } 294 295 std::unique_ptr<MLInlineAdvice> 296 MLInlineAdvisor::getMandatoryAdviceImpl(CallBase &CB) { 297 return std::make_unique<MLInlineAdvice>(this, CB, getCallerORE(CB), true); 298 } 299 300 void MLInlineAdvice::reportContextForRemark( 301 DiagnosticInfoOptimizationBase &OR) { 302 using namespace ore; 303 OR << NV("Callee", Callee->getName()); 304 for (size_t I = 0; I < NumberOfFeatures; ++I) 305 OR << NV(FeatureNameMap[I], getAdvisor()->getModelRunner().getFeature(I)); 306 OR << NV("ShouldInline", isInliningRecommended()); 307 } 308 309 void MLInlineAdvice::recordInliningImpl() { 310 ORE.emit([&]() { 311 OptimizationRemark R(DEBUG_TYPE, "InliningSuccess", DLoc, Block); 312 reportContextForRemark(R); 313 return R; 314 }); 315 getAdvisor()->onSuccessfulInlining(*this, /*CalleeWasDeleted*/ false); 316 } 317 318 void MLInlineAdvice::recordInliningWithCalleeDeletedImpl() { 319 ORE.emit([&]() { 320 OptimizationRemark R(DEBUG_TYPE, "InliningSuccessWithCalleeDeleted", DLoc, 321 Block); 322 reportContextForRemark(R); 323 return R; 324 }); 325 getAdvisor()->onSuccessfulInlining(*this, /*CalleeWasDeleted*/ true); 326 } 327 328 void MLInlineAdvice::recordUnsuccessfulInliningImpl( 329 const InlineResult &Result) { 330 ORE.emit([&]() { 331 OptimizationRemarkMissed R(DEBUG_TYPE, "InliningAttemptedAndUnsuccessful", 332 DLoc, Block); 333 reportContextForRemark(R); 334 return R; 335 }); 336 } 337 void MLInlineAdvice::recordUnattemptedInliningImpl() { 338 ORE.emit([&]() { 339 OptimizationRemarkMissed R(DEBUG_TYPE, "IniningNotAttempted", DLoc, Block); 340 reportContextForRemark(R); 341 return R; 342 }); 343 } 344 #endif // defined(LLVM_HAVE_TF_AOT) || defined(LLVM_HAVE_TF_API) 345