436 lines · cpp
1//===- PGOCtxProfLowering.cpp - Contextual PGO Instr. Lowering ------------===//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-exception6//7//===----------------------------------------------------------------------===//8//9 10#include "llvm/Transforms/Instrumentation/PGOCtxProfLowering.h"11#include "llvm/ADT/STLExtras.h"12#include "llvm/Analysis/CFG.h"13#include "llvm/Analysis/CtxProfAnalysis.h"14#include "llvm/Analysis/OptimizationRemarkEmitter.h"15#include "llvm/IR/Analysis.h"16#include "llvm/IR/Constants.h"17#include "llvm/IR/DiagnosticInfo.h"18#include "llvm/IR/GlobalValue.h"19#include "llvm/IR/IRBuilder.h"20#include "llvm/IR/InstrTypes.h"21#include "llvm/IR/Instructions.h"22#include "llvm/IR/IntrinsicInst.h"23#include "llvm/IR/Module.h"24#include "llvm/IR/PassManager.h"25#include "llvm/ProfileData/CtxInstrContextNode.h"26#include "llvm/ProfileData/InstrProf.h"27#include "llvm/Support/CommandLine.h"28#include <utility>29 30using namespace llvm;31 32#define DEBUG_TYPE "ctx-instr-lower"33 34static cl::list<std::string> ContextRoots(35 "profile-context-root", cl::Hidden,36 cl::desc(37 "A function name, assumed to be global, which will be treated as the "38 "root of an interesting graph, which will be profiled independently "39 "from other similar graphs."));40 41bool PGOCtxProfLoweringPass::isCtxIRPGOInstrEnabled() {42 return !ContextRoots.empty();43}44 45// the names of symbols we expect in compiler-rt. Using a namespace for46// readability.47namespace CompilerRtAPINames {48static auto StartCtx = "__llvm_ctx_profile_start_context";49static auto ReleaseCtx = "__llvm_ctx_profile_release_context";50static auto GetCtx = "__llvm_ctx_profile_get_context";51static auto ExpectedCalleeTLS = "__llvm_ctx_profile_expected_callee";52static auto CallsiteTLS = "__llvm_ctx_profile_callsite";53} // namespace CompilerRtAPINames54 55namespace {56// The lowering logic and state.57class CtxInstrumentationLowerer final {58 Module &M;59 ModuleAnalysisManager &MAM;60 Type *ContextNodeTy = nullptr;61 StructType *FunctionDataTy = nullptr;62 63 DenseSet<const Function *> ContextRootSet;64 Function *StartCtx = nullptr;65 Function *GetCtx = nullptr;66 Function *ReleaseCtx = nullptr;67 GlobalVariable *ExpectedCalleeTLS = nullptr;68 GlobalVariable *CallsiteInfoTLS = nullptr;69 Constant *CannotBeRootInitializer = nullptr;70 71public:72 CtxInstrumentationLowerer(Module &M, ModuleAnalysisManager &MAM);73 // return true if lowering happened (i.e. a change was made)74 bool lowerFunction(Function &F);75};76 77// llvm.instrprof.increment[.step] captures the total number of counters as one78// of its parameters, and llvm.instrprof.callsite captures the total number of79// callsites. Those values are the same for instances of those intrinsics in80// this function. Find the first instance of each and return them.81std::pair<uint32_t, uint32_t> getNumCountersAndCallsites(const Function &F) {82 uint32_t NumCounters = 0;83 uint32_t NumCallsites = 0;84 for (const auto &BB : F) {85 for (const auto &I : BB) {86 if (const auto *Incr = dyn_cast<InstrProfIncrementInst>(&I)) {87 uint32_t V =88 static_cast<uint32_t>(Incr->getNumCounters()->getZExtValue());89 assert((!NumCounters || V == NumCounters) &&90 "expected all llvm.instrprof.increment[.step] intrinsics to "91 "have the same total nr of counters parameter");92 NumCounters = V;93 } else if (const auto *CSIntr = dyn_cast<InstrProfCallsite>(&I)) {94 uint32_t V =95 static_cast<uint32_t>(CSIntr->getNumCounters()->getZExtValue());96 assert((!NumCallsites || V == NumCallsites) &&97 "expected all llvm.instrprof.callsite intrinsics to have the "98 "same total nr of callsites parameter");99 NumCallsites = V;100 }101#ifdef NDEBUG102 if (NumCounters && NumCallsites)103 return std::make_pair(NumCounters, NumCallsites);104#endif105 }106 }107 return {NumCounters, NumCallsites};108}109 110void emitUnsupportedRootError(const Function &F, StringRef Reason) {111 F.getContext().emitError("[ctxprof] The function " + F.getName() +112 " was indicated as context root but " + Reason +113 ", which is not supported.");114}115} // namespace116 117// set up tie-in with compiler-rt.118// NOTE!!!119// These have to match compiler-rt/lib/ctx_profile/CtxInstrProfiling.h120CtxInstrumentationLowerer::CtxInstrumentationLowerer(Module &M,121 ModuleAnalysisManager &MAM)122 : M(M), MAM(MAM) {123 auto *PointerTy = PointerType::get(M.getContext(), 0);124 auto *SanitizerMutexType = Type::getInt8Ty(M.getContext());125 auto *I32Ty = Type::getInt32Ty(M.getContext());126 auto *I64Ty = Type::getInt64Ty(M.getContext());127 128#define _PTRDECL(_, __) PointerTy,129#define _VOLATILE_PTRDECL(_, __) PointerTy,130#define _CONTEXT_ROOT PointerTy,131#define _MUTEXDECL(_) SanitizerMutexType,132 133 FunctionDataTy = StructType::get(134 M.getContext(), {CTXPROF_FUNCTION_DATA(_PTRDECL, _CONTEXT_ROOT,135 _VOLATILE_PTRDECL, _MUTEXDECL)});136#undef _PTRDECL137#undef _CONTEXT_ROOT138#undef _VOLATILE_PTRDECL139#undef _MUTEXDECL140 141#define _PTRDECL(_, __) Constant::getNullValue(PointerTy),142#define _VOLATILE_PTRDECL(_, __) _PTRDECL(_, __)143#define _MUTEXDECL(_) Constant::getNullValue(SanitizerMutexType),144#define _CONTEXT_ROOT \145 Constant::getIntegerValue( \146 PointerTy, \147 APInt(M.getDataLayout().getPointerTypeSizeInBits(PointerTy), 1U)),148 CannotBeRootInitializer = ConstantStruct::get(149 FunctionDataTy, {CTXPROF_FUNCTION_DATA(_PTRDECL, _CONTEXT_ROOT,150 _VOLATILE_PTRDECL, _MUTEXDECL)});151#undef _PTRDECL152#undef _CONTEXT_ROOT153#undef _VOLATILE_PTRDECL154#undef _MUTEXDECL155 156 // The Context header.157 ContextNodeTy = StructType::get(M.getContext(), {158 I64Ty, /*Guid*/159 PointerTy, /*Next*/160 I32Ty, /*NumCounters*/161 I32Ty, /*NumCallsites*/162 });163 164 // Define a global for each entrypoint. We'll reuse the entrypoint's name165 // as prefix. We assume the entrypoint names to be unique.166 for (const auto &Fname : ContextRoots) {167 if (const auto *F = M.getFunction(Fname)) {168 if (F->isDeclaration())169 continue;170 ContextRootSet.insert(F);171 for (const auto &BB : *F)172 for (const auto &I : BB)173 if (const auto *CB = dyn_cast<CallBase>(&I))174 if (CB->isMustTailCall())175 emitUnsupportedRootError(*F, "it features musttail calls");176 }177 }178 179 // Declare the functions we will call.180 StartCtx = cast<Function>(181 M.getOrInsertFunction(182 CompilerRtAPINames::StartCtx,183 FunctionType::get(PointerTy,184 {PointerTy, /*FunctionData*/185 I64Ty, /*Guid*/ I32Ty,186 /*NumCounters*/ I32Ty /*NumCallsites*/},187 false))188 .getCallee());189 GetCtx = cast<Function>(190 M.getOrInsertFunction(CompilerRtAPINames::GetCtx,191 FunctionType::get(PointerTy,192 {PointerTy, /*FunctionData*/193 PointerTy, /*Callee*/194 I64Ty, /*Guid*/195 I32Ty, /*NumCounters*/196 I32Ty}, /*NumCallsites*/197 false))198 .getCallee());199 ReleaseCtx = cast<Function>(200 M.getOrInsertFunction(CompilerRtAPINames::ReleaseCtx,201 FunctionType::get(Type::getVoidTy(M.getContext()),202 {203 PointerTy, /*FunctionData*/204 },205 false))206 .getCallee());207 208 // Declare the TLSes we will need to use.209 CallsiteInfoTLS =210 new GlobalVariable(M, PointerTy, false, GlobalValue::ExternalLinkage,211 nullptr, CompilerRtAPINames::CallsiteTLS);212 CallsiteInfoTLS->setThreadLocal(true);213 CallsiteInfoTLS->setVisibility(llvm::GlobalValue::HiddenVisibility);214 ExpectedCalleeTLS =215 new GlobalVariable(M, PointerTy, false, GlobalValue::ExternalLinkage,216 nullptr, CompilerRtAPINames::ExpectedCalleeTLS);217 ExpectedCalleeTLS->setThreadLocal(true);218 ExpectedCalleeTLS->setVisibility(llvm::GlobalValue::HiddenVisibility);219}220 221PreservedAnalyses PGOCtxProfLoweringPass::run(Module &M,222 ModuleAnalysisManager &MAM) {223 CtxInstrumentationLowerer Lowerer(M, MAM);224 bool Changed = false;225 for (auto &F : M)226 Changed |= Lowerer.lowerFunction(F);227 return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();228}229 230bool CtxInstrumentationLowerer::lowerFunction(Function &F) {231 if (F.isDeclaration())232 return false;233 234 // Probably pointless to try to do anything here, unlikely to be235 // performance-affecting.236 if (!llvm::canReturn(F)) {237 for (auto &BB : F)238 for (auto &I : make_early_inc_range(BB))239 if (isa<InstrProfCntrInstBase>(&I))240 I.eraseFromParent();241 if (ContextRootSet.contains(&F))242 emitUnsupportedRootError(F, "it does not return");243 return true;244 }245 246 auto &FAM = MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();247 auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);248 249 Value *Guid = nullptr;250 auto [NumCounters, NumCallsites] = getNumCountersAndCallsites(F);251 252 Value *Context = nullptr;253 Value *RealContext = nullptr;254 255 StructType *ThisContextType = nullptr;256 Value *TheRootFunctionData = nullptr;257 Value *ExpectedCalleeTLSAddr = nullptr;258 Value *CallsiteInfoTLSAddr = nullptr;259 const bool HasMusttail = [&F]() {260 for (auto &BB : F)261 for (auto &I : BB)262 if (auto *CB = dyn_cast<CallBase>(&I))263 if (CB->isMustTailCall())264 return true;265 return false;266 }();267 268 if (HasMusttail && ContextRootSet.contains(&F)) {269 F.getContext().emitError(270 "[ctx_prof] A function with musttail calls was explicitly requested as "271 "root. That is not supported because we cannot instrument a return "272 "instruction to release the context: " +273 F.getName());274 return false;275 }276 auto &Head = F.getEntryBlock();277 for (auto &I : Head) {278 // Find the increment intrinsic in the entry basic block.279 if (auto *Mark = dyn_cast<InstrProfIncrementInst>(&I)) {280 assert(Mark->getIndex()->isZero());281 282 IRBuilder<> Builder(Mark);283 Guid = Builder.getInt64(284 AssignGUIDPass::getGUID(cast<Function>(*Mark->getNameValue())));285 // The type of the context of this function is now knowable since we have286 // NumCallsites and NumCounters. We declare it here because it's more287 // convenient - we have the Builder.288 ThisContextType = StructType::get(289 F.getContext(),290 {ContextNodeTy, ArrayType::get(Builder.getInt64Ty(), NumCounters),291 ArrayType::get(Builder.getPtrTy(), NumCallsites)});292 // Figure out which way we obtain the context object for this function -293 // if it's an entrypoint, then we call StartCtx, otherwise GetCtx. In the294 // former case, we also set TheRootFunctionData since we need to release295 // it at the end (plus it can be used to know if we have an entrypoint or296 // a regular function). Don't set a name, they end up taking a lot of297 // space and we don't need them.298 299 // Zero-initialize the FunctionData, except for functions that have300 // musttail calls. There, we set the CtxRoot field to 1, which will be301 // treated as a "can't be set as root".302 TheRootFunctionData = new GlobalVariable(303 M, FunctionDataTy, false, GlobalVariable::InternalLinkage,304 HasMusttail ? CannotBeRootInitializer305 : Constant::getNullValue(FunctionDataTy));306 307 if (ContextRootSet.contains(&F)) {308 Context = Builder.CreateCall(309 StartCtx, {TheRootFunctionData, Guid, Builder.getInt32(NumCounters),310 Builder.getInt32(NumCallsites)});311 ORE.emit(312 [&] { return OptimizationRemark(DEBUG_TYPE, "Entrypoint", &F); });313 } else {314 Context = Builder.CreateCall(GetCtx, {TheRootFunctionData, &F, Guid,315 Builder.getInt32(NumCounters),316 Builder.getInt32(NumCallsites)});317 ORE.emit([&] {318 return OptimizationRemark(DEBUG_TYPE, "RegularFunction", &F);319 });320 }321 // The context could be scratch.322 auto *CtxAsInt = Builder.CreatePtrToInt(Context, Builder.getInt64Ty());323 if (NumCallsites > 0) {324 // Figure out which index of the TLS 2-element buffers to use.325 // Scratch context => we use index == 1. Real contexts => index == 0.326 auto *Index = Builder.CreateAnd(CtxAsInt, Builder.getInt64(1));327 // The GEPs corresponding to that index, in the respective TLS.328 ExpectedCalleeTLSAddr = Builder.CreateGEP(329 PointerType::getUnqual(F.getContext()),330 Builder.CreateThreadLocalAddress(ExpectedCalleeTLS), {Index});331 CallsiteInfoTLSAddr = Builder.CreateGEP(332 Builder.getInt32Ty(),333 Builder.CreateThreadLocalAddress(CallsiteInfoTLS), {Index});334 }335 // Because the context pointer may have LSB set (to indicate scratch),336 // clear it for the value we use as base address for the counter vector.337 // This way, if later we want to have "real" (not clobbered) buffers338 // acting as scratch, the lowering (at least this part of it that deals339 // with counters) stays the same.340 RealContext = Builder.CreateIntToPtr(341 Builder.CreateAnd(CtxAsInt, Builder.getInt64(-2)),342 PointerType::getUnqual(F.getContext()));343 I.eraseFromParent();344 break;345 }346 }347 if (!Context) {348 ORE.emit([&] {349 return OptimizationRemarkMissed(DEBUG_TYPE, "Skip", &F)350 << "Function doesn't have instrumentation, skipping";351 });352 return false;353 }354 355 bool ContextWasReleased = false;356 for (auto &BB : F) {357 for (auto &I : llvm::make_early_inc_range(BB)) {358 if (auto *Instr = dyn_cast<InstrProfCntrInstBase>(&I)) {359 IRBuilder<> Builder(Instr);360 switch (Instr->getIntrinsicID()) {361 case llvm::Intrinsic::instrprof_increment:362 case llvm::Intrinsic::instrprof_increment_step: {363 // Increments (or increment-steps) are just a typical load - increment364 // - store in the RealContext.365 auto *AsStep = cast<InstrProfIncrementInst>(Instr);366 auto *GEP = Builder.CreateGEP(367 ThisContextType, RealContext,368 {Builder.getInt32(0), Builder.getInt32(1), AsStep->getIndex()});369 Builder.CreateStore(370 Builder.CreateAdd(Builder.CreateLoad(Builder.getInt64Ty(), GEP),371 AsStep->getStep()),372 GEP);373 } break;374 case llvm::Intrinsic::instrprof_callsite:375 // callsite lowering: write the called value in the expected callee376 // TLS we treat the TLS as volatile because of signal handlers and to377 // avoid these being moved away from the callsite they decorate.378 auto *CSIntrinsic = dyn_cast<InstrProfCallsite>(Instr);379 Builder.CreateStore(CSIntrinsic->getCallee(), ExpectedCalleeTLSAddr,380 true);381 // write the GEP of the slot in the sub-contexts portion of the382 // context in TLS. Now, here, we use the actual Context value - as383 // returned from compiler-rt - which may have the LSB set if the384 // Context was scratch. Since the header of the context object and385 // then the values are all 8-aligned (or, really, insofar as we care,386 // they are even) - if the context is scratch (meaning, an odd value),387 // so will the GEP. This is important because this is then visible to388 // compiler-rt which will produce scratch contexts for callers that389 // have a scratch context.390 Builder.CreateStore(391 Builder.CreateGEP(ThisContextType, Context,392 {Builder.getInt32(0), Builder.getInt32(2),393 CSIntrinsic->getIndex()}),394 CallsiteInfoTLSAddr, true);395 break;396 }397 I.eraseFromParent();398 } else if (!HasMusttail && isa<ReturnInst>(I)) {399 // Remember to release the context if we are an entrypoint.400 IRBuilder<> Builder(&I);401 Builder.CreateCall(ReleaseCtx, {TheRootFunctionData});402 ContextWasReleased = true;403 }404 }405 }406 if (!HasMusttail && !ContextWasReleased)407 F.getContext().emitError(408 "[ctx_prof] A function that doesn't have musttail calls was "409 "instrumented but it has no `ret` "410 "instructions above which to release the context: " +411 F.getName());412 return true;413}414 415PreservedAnalyses NoinlineNonPrevailing::run(Module &M,416 ModuleAnalysisManager &MAM) {417 bool Changed = false;418 for (auto &F : M) {419 if (F.isDeclaration())420 continue;421 if (F.hasFnAttribute(Attribute::NoInline))422 continue;423 if (!F.isWeakForLinker())424 continue;425 426 if (F.hasFnAttribute(Attribute::AlwaysInline))427 F.removeFnAttr(Attribute::AlwaysInline);428 429 F.addFnAttr(Attribute::NoInline);430 Changed = true;431 }432 if (Changed)433 return PreservedAnalyses::none();434 return PreservedAnalyses::all();435}436