1//===- CoroAnnotationElide.cpp - Elide attributed safe coroutine calls ----===//
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// \file
10// This pass transforms all Call or Invoke instructions that are annotated
11// "coro_elide_safe" to call the `.noalloc` variant of coroutine instead.
12// The frame of the callee coroutine is allocated inside the caller. A pointer
13// to the allocated frame will be passed into the `.noalloc` ramp function.
14//
15//===----------------------------------------------------------------------===//
16
17#include "llvm/Transforms/Coroutines/CoroAnnotationElide.h"
18
19#include "llvm/Analysis/CGSCCPassManager.h"
20#include "llvm/Analysis/LazyCallGraph.h"
21#include "llvm/Analysis/OptimizationRemarkEmitter.h"
22#include "llvm/IR/Analysis.h"
23#include "llvm/IR/IRBuilder.h"
24#include "llvm/IR/Instruction.h"
25#include "llvm/IR/Module.h"
26#include "llvm/IR/PassManager.h"
27#include "llvm/Support/CommandLine.h"
28#include "llvm/Transforms/Utils/CallGraphUpdater.h"
29#include "llvm/Transforms/Utils/Cloning.h"
30
31#include <cassert>
32
33using namespace llvm;
34
35#define DEBUG_TYPE "coro-annotation-elide"
36
37static cl::opt<float> CoroElideBranchRatio(
38 "coro-elide-branch-ratio", cl::init(Val: 0.55), cl::Hidden,
39 cl::desc("Minimum BranchProbability to consider a elide a coroutine."));
40extern cl::opt<unsigned> MinBlockCounterExecution;
41
42static Instruction *getFirstNonAllocaInTheEntryBlock(Function *F) {
43 for (Instruction &I : F->getEntryBlock())
44 if (!isa<AllocaInst>(Val: &I))
45 return &I;
46 llvm_unreachable("no terminator in the entry block");
47}
48
49// Create an alloca in the caller, using FrameSize and FrameAlign as the callee
50// coroutine's activation frame.
51static Value *allocateFrameInCaller(Function *Caller, uint64_t FrameSize,
52 Align FrameAlign) {
53 LLVMContext &C = Caller->getContext();
54 BasicBlock::iterator InsertPt =
55 getFirstNonAllocaInTheEntryBlock(F: Caller)->getIterator();
56 const DataLayout &DL = Caller->getDataLayout();
57 auto FrameTy = ArrayType::get(ElementType: Type::getInt8Ty(C), NumElements: FrameSize);
58 auto *Frame = new AllocaInst(FrameTy, DL.getAllocaAddrSpace(), "", InsertPt);
59 Frame->setAlignment(FrameAlign);
60 return Frame;
61}
62
63// Given a call or invoke instruction to the elide safe coroutine, this function
64// does the following:
65// - Allocate a frame for the callee coroutine in the caller using alloca.
66// - Replace the old CB with a new Call or Invoke to `NewCallee`, with the
67// pointer to the frame as an additional argument to NewCallee.
68static void processCall(CallBase *CB, Function *Caller, Function *NewCallee,
69 uint64_t FrameSize, Align FrameAlign) {
70 // TODO: generate the lifetime intrinsics for the new frame. This will require
71 // introduction of two pesudo lifetime intrinsics in the frontend around the
72 // `co_await` expression and convert them to real lifetime intrinsics here.
73 auto *FramePtr = allocateFrameInCaller(Caller, FrameSize, FrameAlign);
74 auto NewCBInsertPt = CB->getIterator();
75 llvm::CallBase *NewCB = nullptr;
76 SmallVector<Value *, 4> NewArgs;
77 NewArgs.append(in_start: CB->arg_begin(), in_end: CB->arg_end());
78 NewArgs.push_back(Elt: FramePtr);
79
80 if (auto *CI = dyn_cast<CallInst>(Val: CB)) {
81 auto *NewCI = CallInst::Create(Ty: NewCallee->getFunctionType(), Func: NewCallee,
82 Args: NewArgs, NameStr: "", InsertBefore: NewCBInsertPt);
83 NewCI->setTailCallKind(CI->getTailCallKind());
84 NewCB = NewCI;
85 } else if (auto *II = dyn_cast<InvokeInst>(Val: CB)) {
86 NewCB = InvokeInst::Create(Ty: NewCallee->getFunctionType(), Func: NewCallee,
87 IfNormal: II->getNormalDest(), IfException: II->getUnwindDest(),
88 Args: NewArgs, Bundles: {}, NameStr: "", InsertBefore: NewCBInsertPt);
89 } else {
90 llvm_unreachable("CallBase should either be Call or Invoke!");
91 }
92
93 NewCB->setCalledFunction(FTy: NewCallee->getFunctionType(), Fn: NewCallee);
94 NewCB->setCallingConv(CB->getCallingConv());
95 NewCB->setAttributes(CB->getAttributes());
96 NewCB->setDebugLoc(CB->getDebugLoc());
97 std::copy(first: CB->bundle_op_info_begin(), last: CB->bundle_op_info_end(),
98 result: NewCB->bundle_op_info_begin());
99
100 NewCB->removeFnAttr(Kind: llvm::Attribute::CoroElideSafe);
101 CB->replaceAllUsesWith(V: NewCB);
102
103 InlineFunctionInfo IFI;
104 InlineResult IR = InlineFunction(CB&: *NewCB, IFI);
105 if (IR.isSuccess()) {
106 CB->eraseFromParent();
107 } else {
108 NewCB->replaceAllUsesWith(V: CB);
109 NewCB->eraseFromParent();
110 }
111}
112
113PreservedAnalyses CoroAnnotationElidePass::run(LazyCallGraph::SCC &C,
114 CGSCCAnalysisManager &AM,
115 LazyCallGraph &CG,
116 CGSCCUpdateResult &UR) {
117 bool Changed = false;
118 CallGraphUpdater CGUpdater;
119 CGUpdater.initialize(LCG&: CG, SCC&: C, AM, UR);
120
121 auto &FAM =
122 AM.getResult<FunctionAnalysisManagerCGSCCProxy>(IR&: C, ExtraArgs&: CG).getManager();
123
124 for (LazyCallGraph::Node &N : C) {
125 Function *Callee = &N.getFunction();
126 Function *NewCallee = Callee->getParent()->getFunction(
127 Name: (Callee->getName() + ".noalloc").str());
128 if (!NewCallee)
129 continue;
130
131 SmallVector<CallBase *, 4> Users;
132 for (auto *U : Callee->users()) {
133 if (auto *CB = dyn_cast<CallBase>(Val: U)) {
134 if (CB->getCalledFunction() == Callee)
135 Users.push_back(Elt: CB);
136 }
137 }
138 auto FramePtrArgPosition = NewCallee->arg_size() - 1;
139 auto FrameSize =
140 NewCallee->getParamDereferenceableBytes(ArgNo: FramePtrArgPosition);
141 auto FrameAlign =
142 NewCallee->getParamAlign(ArgNo: FramePtrArgPosition).valueOrOne();
143
144 auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(IR&: *Callee);
145
146 for (auto *CB : Users) {
147 auto *Caller = CB->getFunction();
148 if (!Caller)
149 continue;
150
151 bool IsCallerPresplitCoroutine = Caller->isPresplitCoroutine();
152 bool HasAttr = CB->hasFnAttr(Kind: llvm::Attribute::CoroElideSafe);
153 if (IsCallerPresplitCoroutine && HasAttr) {
154 auto &BFI = FAM.getResult<BlockFrequencyAnalysis>(IR&: *Caller);
155
156 auto BlockFreq = BFI.getBlockFreq(BB: CB->getParent()).getFrequency();
157 auto EntryFreq = BFI.getEntryFreq().getFrequency();
158 uint64_t MinFreq =
159 static_cast<uint64_t>(EntryFreq * CoroElideBranchRatio);
160
161 if (BlockFreq < MinFreq) {
162 ORE.emit(RemarkBuilder: [&]() {
163 return OptimizationRemarkMissed(
164 DEBUG_TYPE, "CoroAnnotationElideUnlikely", Caller)
165 << "'" << ore::NV("callee", Callee->getName())
166 << "' not elided in '"
167 << ore::NV("caller", Caller->getName())
168 << "' because of low frequency: "
169 << ore::NV("block_freq", BlockFreq)
170 << " (threshold: " << ore::NV("min_freq", MinFreq) << ")";
171 });
172 continue;
173 }
174
175 auto *CallerN = CG.lookup(F: *Caller);
176 auto *CallerC = CallerN ? CG.lookupSCC(N&: *CallerN) : nullptr;
177 // If CallerC is nullptr, it means LazyCallGraph hasn't visited Caller
178 // yet. Skip the call graph update.
179 auto ShouldUpdateCallGraph = !!CallerC;
180 processCall(CB, Caller, NewCallee, FrameSize, FrameAlign);
181
182 ORE.emit(RemarkBuilder: [&]() {
183 return OptimizationRemark(DEBUG_TYPE, "CoroAnnotationElide", Caller)
184 << "'" << ore::NV("callee", Callee->getName())
185 << "' elided in '" << ore::NV("caller", Caller->getName())
186 << "' (block_freq: " << ore::NV("block_freq", BlockFreq)
187 << ")";
188 });
189
190 FAM.invalidate(IR&: *Caller, PA: PreservedAnalyses::none());
191 Changed = true;
192 if (ShouldUpdateCallGraph)
193 updateCGAndAnalysisManagerForCGSCCPass(G&: CG, C&: *CallerC, N&: *CallerN, AM, UR,
194 FAM);
195
196 } else {
197 ORE.emit(RemarkBuilder: [&]() {
198 return OptimizationRemarkMissed(DEBUG_TYPE, "CoroAnnotationElide",
199 Caller)
200 << "'" << ore::NV("callee", Callee->getName())
201 << "' not elided in '" << ore::NV("caller", Caller->getName())
202 << "' (caller_presplit="
203 << ore::NV("caller_presplit", IsCallerPresplitCoroutine)
204 << ", elide_safe_attr=" << ore::NV("elide_safe_attr", HasAttr)
205 << ")";
206 });
207 }
208 }
209 }
210
211 return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
212}
213