1//===- CoroEarly.cpp - Coroutine Early Function Pass ----------------------===//
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#include "llvm/Transforms/Coroutines/CoroEarly.h"
10#include "CoroInternal.h"
11#include "llvm/IR/DIBuilder.h"
12#include "llvm/IR/Function.h"
13#include "llvm/IR/IRBuilder.h"
14#include "llvm/IR/InstIterator.h"
15#include "llvm/IR/Module.h"
16
17using namespace llvm;
18
19#define DEBUG_TYPE "coro-early"
20
21namespace {
22// Created on demand if the coro-early pass has work to do.
23class Lowerer : public coro::LowererBase {
24  IRBuilder<> Builder;
25  PointerType *const AnyResumeFnPtrTy;
26  Constant *NoopCoro = nullptr;
27
28  void lowerResumeOrDestroy(CallBase &CB, CoroSubFnInst::ResumeKind);
29  void lowerCoroPromise(CoroPromiseInst *Intrin);
30  void lowerCoroDone(IntrinsicInst *II);
31  void lowerCoroNoop(IntrinsicInst *II);
32
33public:
34  Lowerer(Module &M)
35      : LowererBase(M), Builder(Context),
36        AnyResumeFnPtrTy(FunctionType::get(Type::getVoidTy(Context), Int8Ptr,
37                                           /*isVarArg=*/false)
38                             ->getPointerTo()) {}
39  void lowerEarlyIntrinsics(Function &F);
40};
41}
42
43// Replace a direct call to coro.resume or coro.destroy with an indirect call to
44// an address returned by coro.subfn.addr intrinsic. This is done so that
45// CGPassManager recognizes devirtualization when CoroElide pass replaces a call
46// to coro.subfn.addr with an appropriate function address.
47void Lowerer::lowerResumeOrDestroy(CallBase &CB,
48                                   CoroSubFnInst::ResumeKind Index) {
49  Value *ResumeAddr = makeSubFnCall(CB.getArgOperand(0), Index, &CB);
50  CB.setCalledOperand(ResumeAddr);
51  CB.setCallingConv(CallingConv::Fast);
52}
53
54// Coroutine promise field is always at the fixed offset from the beginning of
55// the coroutine frame. i8* coro.promise(i8*, i1 from) intrinsic adds an offset
56// to a passed pointer to move from coroutine frame to coroutine promise and
57// vice versa. Since we don't know exactly which coroutine frame it is, we build
58// a coroutine frame mock up starting with two function pointers, followed by a
59// properly aligned coroutine promise field.
60// TODO: Handle the case when coroutine promise alloca has align override.
61void Lowerer::lowerCoroPromise(CoroPromiseInst *Intrin) {
62  Value *Operand = Intrin->getArgOperand(0);
63  Align Alignment = Intrin->getAlignment();
64  Type *Int8Ty = Builder.getInt8Ty();
65
66  auto *SampleStruct =
67      StructType::get(Context, {AnyResumeFnPtrTy, AnyResumeFnPtrTy, Int8Ty});
68  const DataLayout &DL = TheModule.getDataLayout();
69  int64_t Offset = alignTo(
70      DL.getStructLayout(SampleStruct)->getElementOffset(2), Alignment);
71  if (Intrin->isFromPromise())
72    Offset = -Offset;
73
74  Builder.SetInsertPoint(Intrin);
75  Value *Replacement =
76      Builder.CreateConstInBoundsGEP1_32(Int8Ty, Operand, Offset);
77
78  Intrin->replaceAllUsesWith(Replacement);
79  Intrin->eraseFromParent();
80}
81
82// When a coroutine reaches final suspend point, it zeros out ResumeFnAddr in
83// the coroutine frame (it is UB to resume from a final suspend point).
84// The llvm.coro.done intrinsic is used to check whether a coroutine is
85// suspended at the final suspend point or not.
86void Lowerer::lowerCoroDone(IntrinsicInst *II) {
87  Value *Operand = II->getArgOperand(0);
88
89  // ResumeFnAddr is the first pointer sized element of the coroutine frame.
90  static_assert(coro::Shape::SwitchFieldIndex::Resume == 0,
91                "resume function not at offset zero");
92  auto *FrameTy = Int8Ptr;
93  PointerType *FramePtrTy = FrameTy->getPointerTo();
94
95  Builder.SetInsertPoint(II);
96  auto *BCI = Builder.CreateBitCast(Operand, FramePtrTy);
97  auto *Load = Builder.CreateLoad(FrameTy, BCI);
98  auto *Cond = Builder.CreateICmpEQ(Load, NullPtr);
99
100  II->replaceAllUsesWith(Cond);
101  II->eraseFromParent();
102}
103
104static void buildDebugInfoForNoopResumeDestroyFunc(Function *NoopFn) {
105  Module &M = *NoopFn->getParent();
106  if (M.debug_compile_units().empty())
107     return;
108
109  DICompileUnit *CU = *M.debug_compile_units_begin();
110  DIBuilder DB(M, /*AllowUnresolved*/ false, CU);
111  std::array<Metadata *, 2> Params{nullptr, nullptr};
112  auto *SubroutineType =
113      DB.createSubroutineType(DB.getOrCreateTypeArray(Params));
114  StringRef Name = NoopFn->getName();
115  auto *SP = DB.createFunction(
116      CU, /*Name=*/Name, /*LinkageName=*/Name, /*File=*/ CU->getFile(),
117      /*LineNo=*/0, SubroutineType, /*ScopeLine=*/0, DINode::FlagArtificial,
118      DISubprogram::SPFlagDefinition);
119  NoopFn->setSubprogram(SP);
120  DB.finalize();
121}
122
123void Lowerer::lowerCoroNoop(IntrinsicInst *II) {
124  if (!NoopCoro) {
125    LLVMContext &C = Builder.getContext();
126    Module &M = *II->getModule();
127
128    // Create a noop.frame struct type.
129    StructType *FrameTy = StructType::create(C, "NoopCoro.Frame");
130    auto *FramePtrTy = FrameTy->getPointerTo();
131    auto *FnTy = FunctionType::get(Type::getVoidTy(C), FramePtrTy,
132                                   /*isVarArg=*/false);
133    auto *FnPtrTy = FnTy->getPointerTo();
134    FrameTy->setBody({FnPtrTy, FnPtrTy});
135
136    // Create a Noop function that does nothing.
137    Function *NoopFn =
138        Function::Create(FnTy, GlobalValue::LinkageTypes::PrivateLinkage,
139                         "__NoopCoro_ResumeDestroy", &M);
140    NoopFn->setCallingConv(CallingConv::Fast);
141    buildDebugInfoForNoopResumeDestroyFunc(NoopFn);
142    auto *Entry = BasicBlock::Create(C, "entry", NoopFn);
143    ReturnInst::Create(C, Entry);
144
145    // Create a constant struct for the frame.
146    Constant* Values[] = {NoopFn, NoopFn};
147    Constant* NoopCoroConst = ConstantStruct::get(FrameTy, Values);
148    NoopCoro = new GlobalVariable(M, NoopCoroConst->getType(), /*isConstant=*/true,
149                                GlobalVariable::PrivateLinkage, NoopCoroConst,
150                                "NoopCoro.Frame.Const");
151  }
152
153  Builder.SetInsertPoint(II);
154  auto *NoopCoroVoidPtr = Builder.CreateBitCast(NoopCoro, Int8Ptr);
155  II->replaceAllUsesWith(NoopCoroVoidPtr);
156  II->eraseFromParent();
157}
158
159// Prior to CoroSplit, calls to coro.begin needs to be marked as NoDuplicate,
160// as CoroSplit assumes there is exactly one coro.begin. After CoroSplit,
161// NoDuplicate attribute will be removed from coro.begin otherwise, it will
162// interfere with inlining.
163static void setCannotDuplicate(CoroIdInst *CoroId) {
164  for (User *U : CoroId->users())
165    if (auto *CB = dyn_cast<CoroBeginInst>(U))
166      CB->setCannotDuplicate();
167}
168
169void Lowerer::lowerEarlyIntrinsics(Function &F) {
170  CoroIdInst *CoroId = nullptr;
171  SmallVector<CoroFreeInst *, 4> CoroFrees;
172  bool HasCoroSuspend = false;
173  for (Instruction &I : llvm::make_early_inc_range(instructions(F))) {
174    auto *CB = dyn_cast<CallBase>(&I);
175    if (!CB)
176      continue;
177
178    switch (CB->getIntrinsicID()) {
179      default:
180        continue;
181      case Intrinsic::coro_free:
182        CoroFrees.push_back(cast<CoroFreeInst>(&I));
183        break;
184      case Intrinsic::coro_suspend:
185        // Make sure that final suspend point is not duplicated as CoroSplit
186        // pass expects that there is at most one final suspend point.
187        if (cast<CoroSuspendInst>(&I)->isFinal())
188          CB->setCannotDuplicate();
189        HasCoroSuspend = true;
190        break;
191      case Intrinsic::coro_end_async:
192      case Intrinsic::coro_end:
193        // Make sure that fallthrough coro.end is not duplicated as CoroSplit
194        // pass expects that there is at most one fallthrough coro.end.
195        if (cast<AnyCoroEndInst>(&I)->isFallthrough())
196          CB->setCannotDuplicate();
197        break;
198      case Intrinsic::coro_noop:
199        lowerCoroNoop(cast<IntrinsicInst>(&I));
200        break;
201      case Intrinsic::coro_id:
202        if (auto *CII = cast<CoroIdInst>(&I)) {
203          if (CII->getInfo().isPreSplit()) {
204            assert(F.isPresplitCoroutine() &&
205                   "The frontend uses Swtich-Resumed ABI should emit "
206                   "\"coroutine.presplit\" attribute for the coroutine.");
207            setCannotDuplicate(CII);
208            CII->setCoroutineSelf();
209            CoroId = cast<CoroIdInst>(&I);
210          }
211        }
212        break;
213      case Intrinsic::coro_id_retcon:
214      case Intrinsic::coro_id_retcon_once:
215      case Intrinsic::coro_id_async:
216        F.setPresplitCoroutine();
217        break;
218      case Intrinsic::coro_resume:
219        lowerResumeOrDestroy(*CB, CoroSubFnInst::ResumeIndex);
220        break;
221      case Intrinsic::coro_destroy:
222        lowerResumeOrDestroy(*CB, CoroSubFnInst::DestroyIndex);
223        break;
224      case Intrinsic::coro_promise:
225        lowerCoroPromise(cast<CoroPromiseInst>(&I));
226        break;
227      case Intrinsic::coro_done:
228        lowerCoroDone(cast<IntrinsicInst>(&I));
229        break;
230    }
231  }
232
233  // Make sure that all CoroFree reference the coro.id intrinsic.
234  // Token type is not exposed through coroutine C/C++ builtins to plain C, so
235  // we allow specifying none and fixing it up here.
236  if (CoroId)
237    for (CoroFreeInst *CF : CoroFrees)
238      CF->setArgOperand(0, CoroId);
239
240  // Coroutine suspention could potentially lead to any argument modified
241  // outside of the function, hence arguments should not have noalias
242  // attributes.
243  if (HasCoroSuspend)
244    for (Argument &A : F.args())
245      if (A.hasNoAliasAttr())
246        A.removeAttr(Attribute::NoAlias);
247}
248
249static bool declaresCoroEarlyIntrinsics(const Module &M) {
250  return coro::declaresIntrinsics(
251      M, {"llvm.coro.id", "llvm.coro.id.retcon", "llvm.coro.id.retcon.once",
252          "llvm.coro.id.async", "llvm.coro.destroy", "llvm.coro.done",
253          "llvm.coro.end", "llvm.coro.end.async", "llvm.coro.noop",
254          "llvm.coro.free", "llvm.coro.promise", "llvm.coro.resume",
255          "llvm.coro.suspend"});
256}
257
258PreservedAnalyses CoroEarlyPass::run(Module &M, ModuleAnalysisManager &) {
259  if (!declaresCoroEarlyIntrinsics(M))
260    return PreservedAnalyses::all();
261
262  Lowerer L(M);
263  for (auto &F : M)
264    L.lowerEarlyIntrinsics(F);
265
266  PreservedAnalyses PA;
267  PA.preserveSet<CFGAnalyses>();
268  return PA;
269}
270