contrib/llvm/lib/Transforms/Coroutines/CoroEarly.cpp

   1 //===- CoroEarly.cpp - Coroutine Early Function Pass ----------------------===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file is distributed under the University of Illinois Open Source
   6 // License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 // This pass lowers coroutine intrinsics that hide the details of the exact
  10 // calling convention for coroutine resume and destroy functions and details of
  11 // the structure of the coroutine frame.
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "CoroInternal.h"
  15 #include "llvm/IR/CallSite.h"
  16 #include "llvm/IR/IRBuilder.h"
  17 #include "llvm/IR/InstIterator.h"
  18 #include "llvm/IR/Module.h"
  19 #include "llvm/Pass.h"
  20
  21 using namespace llvm;
  22
  23 #define DEBUG_TYPE "coro-early"
  24
  25 namespace {
  26 // Created on demand if CoroEarly pass has work to do.
  27 class Lowerer : public coro::LowererBase {
  28   IRBuilder<> Builder;
  29   PointerType *const AnyResumeFnPtrTy;
  30
  31   void lowerResumeOrDestroy(CallSite CS, CoroSubFnInst::ResumeKind);
  32   void lowerCoroPromise(CoroPromiseInst *Intrin);
  33   void lowerCoroDone(IntrinsicInst *II);
  34
  35 public:
  36   Lowerer(Module &M)
  37       : LowererBase(M), Builder(Context),
  38         AnyResumeFnPtrTy(FunctionType::get(Type::getVoidTy(Context), Int8Ptr,
  39                                            /*isVarArg=*/false)
  40                              ->getPointerTo()) {}
  41   bool lowerEarlyIntrinsics(Function &F);
  42 };
  43 }
  44
  45 // Replace a direct call to coro.resume or coro.destroy with an indirect call to
  46 // an address returned by coro.subfn.addr intrinsic. This is done so that
  47 // CGPassManager recognizes devirtualization when CoroElide pass replaces a call
  48 // to coro.subfn.addr with an appropriate function address.
  49 void Lowerer::lowerResumeOrDestroy(CallSite CS,
  50                                    CoroSubFnInst::ResumeKind Index) {
  51   Value *ResumeAddr =
  52       makeSubFnCall(CS.getArgOperand(0), Index, CS.getInstruction());
  53   CS.setCalledFunction(ResumeAddr);
  54   CS.setCallingConv(CallingConv::Fast);
  55 }
  56
  57 // Coroutine promise field is always at the fixed offset from the beginning of
  58 // the coroutine frame. i8* coro.promise(i8*, i1 from) intrinsic adds an offset
  59 // to a passed pointer to move from coroutine frame to coroutine promise and
  60 // vice versa. Since we don't know exactly which coroutine frame it is, we build
  61 // a coroutine frame mock up starting with two function pointers, followed by a
  62 // properly aligned coroutine promise field.
  63 // TODO: Handle the case when coroutine promise alloca has align override.
  64 void Lowerer::lowerCoroPromise(CoroPromiseInst *Intrin) {
  65   Value *Operand = Intrin->getArgOperand(0);
  66   unsigned Alignement = Intrin->getAlignment();
  67   Type *Int8Ty = Builder.getInt8Ty();
  68
  69   auto *SampleStruct =
  70       StructType::get(Context, {AnyResumeFnPtrTy, AnyResumeFnPtrTy, Int8Ty});
  71   const DataLayout &DL = TheModule.getDataLayout();
  72   int64_t Offset = alignTo(
  73       DL.getStructLayout(SampleStruct)->getElementOffset(2), Alignement);
  74   if (Intrin->isFromPromise())
  75     Offset = -Offset;
  76
  77   Builder.SetInsertPoint(Intrin);
  78   Value *Replacement =
  79       Builder.CreateConstInBoundsGEP1_32(Int8Ty, Operand, Offset);
  80
  81   Intrin->replaceAllUsesWith(Replacement);
  82   Intrin->eraseFromParent();
  83 }
  84
  85 // When a coroutine reaches final suspend point, it zeros out ResumeFnAddr in
  86 // the coroutine frame (it is UB to resume from a final suspend point).
  87 // The llvm.coro.done intrinsic is used to check whether a coroutine is
  88 // suspended at the final suspend point or not.
  89 void Lowerer::lowerCoroDone(IntrinsicInst *II) {
  90   Value *Operand = II->getArgOperand(0);
  91
  92   // ResumeFnAddr is the first pointer sized element of the coroutine frame.
  93   auto *FrameTy = Int8Ptr;
  94   PointerType *FramePtrTy = FrameTy->getPointerTo();
  95
  96   Builder.SetInsertPoint(II);
  97   auto *BCI = Builder.CreateBitCast(Operand, FramePtrTy);
  98   auto *Gep = Builder.CreateConstInBoundsGEP1_32(FrameTy, BCI, 0);
  99   auto *Load = Builder.CreateLoad(Gep);
 100   auto *Cond = Builder.CreateICmpEQ(Load, NullPtr);
 101
 102   II->replaceAllUsesWith(Cond);
 103   II->eraseFromParent();
 104 }
 105
 106 // Prior to CoroSplit, calls to coro.begin needs to be marked as NoDuplicate,
 107 // as CoroSplit assumes there is exactly one coro.begin. After CoroSplit,
 108 // NoDuplicate attribute will be removed from coro.begin otherwise, it will
 109 // interfere with inlining.
 110 static void setCannotDuplicate(CoroIdInst *CoroId) {
 111   for (User *U : CoroId->users())
 112     if (auto *CB = dyn_cast<CoroBeginInst>(U))
 113       CB->setCannotDuplicate();
 114 }
 115
 116 bool Lowerer::lowerEarlyIntrinsics(Function &F) {
 117   bool Changed = false;
 118   CoroIdInst *CoroId = nullptr;
 119   SmallVector<CoroFreeInst *, 4> CoroFrees;
 120   for (auto IB = inst_begin(F), IE = inst_end(F); IB != IE;) {
 121     Instruction &I = *IB++;
 122     if (auto CS = CallSite(&I)) {
 123       switch (CS.getIntrinsicID()) {
 124       default:
 125         continue;
 126       case Intrinsic::coro_free:
 127         CoroFrees.push_back(cast<CoroFreeInst>(&I));
 128         break;
 129       case Intrinsic::coro_suspend:
 130         // Make sure that final suspend point is not duplicated as CoroSplit
 131         // pass expects that there is at most one final suspend point.
 132         if (cast<CoroSuspendInst>(&I)->isFinal())
 133           CS.setCannotDuplicate();
 134         break;
 135       case Intrinsic::coro_end:
 136         // Make sure that fallthrough coro.end is not duplicated as CoroSplit
 137         // pass expects that there is at most one fallthrough coro.end.
 138         if (cast<CoroEndInst>(&I)->isFallthrough())
 139           CS.setCannotDuplicate();
 140         break;
 141       case Intrinsic::coro_id:
 142         // Mark a function that comes out of the frontend that has a coro.id
 143         // with a coroutine attribute.
 144         if (auto *CII = cast<CoroIdInst>(&I)) {
 145           if (CII->getInfo().isPreSplit()) {
 146             F.addFnAttr(CORO_PRESPLIT_ATTR, UNPREPARED_FOR_SPLIT);
 147             setCannotDuplicate(CII);
 148             CII->setCoroutineSelf();
 149             CoroId = cast<CoroIdInst>(&I);
 150           }
 151         }
 152         break;
 153       case Intrinsic::coro_resume:
 154         lowerResumeOrDestroy(CS, CoroSubFnInst::ResumeIndex);
 155         break;
 156       case Intrinsic::coro_destroy:
 157         lowerResumeOrDestroy(CS, CoroSubFnInst::DestroyIndex);
 158         break;
 159       case Intrinsic::coro_promise:
 160         lowerCoroPromise(cast<CoroPromiseInst>(&I));
 161         break;
 162       case Intrinsic::coro_done:
 163         lowerCoroDone(cast<IntrinsicInst>(&I));
 164         break;
 165       }
 166       Changed = true;
 167     }
 168   }
 169   // Make sure that all CoroFree reference the coro.id intrinsic.
 170   // Token type is not exposed through coroutine C/C++ builtins to plain C, so
 171   // we allow specifying none and fixing it up here.
 172   if (CoroId)
 173     for (CoroFreeInst *CF : CoroFrees)
 174       CF->setArgOperand(0, CoroId);
 175   return Changed;
 176 }
 177
 178 //===----------------------------------------------------------------------===//
 179 //                              Top Level Driver
 180 //===----------------------------------------------------------------------===//
 181
 182 namespace {
 183
 184 struct CoroEarly : public FunctionPass {
 185   static char ID; // Pass identification, replacement for typeid.
 186   CoroEarly() : FunctionPass(ID) {
 187     initializeCoroEarlyPass(*PassRegistry::getPassRegistry());
 188   }
 189
 190   std::unique_ptr<Lowerer> L;
 191
 192   // This pass has work to do only if we find intrinsics we are going to lower
 193   // in the module.
 194   bool doInitialization(Module &M) override {
 195     if (coro::declaresIntrinsics(M, {"llvm.coro.id", "llvm.coro.destroy",
 196                                      "llvm.coro.done", "llvm.coro.end",
 197                                      "llvm.coro.free", "llvm.coro.promise",
 198                                      "llvm.coro.resume", "llvm.coro.suspend"}))
 199       L = llvm::make_unique<Lowerer>(M);
 200     return false;
 201   }
 202
 203   bool runOnFunction(Function &F) override {
 204     if (!L)
 205       return false;
 206
 207     return L->lowerEarlyIntrinsics(F);
 208   }
 209
 210   void getAnalysisUsage(AnalysisUsage &AU) const override {
 211     AU.setPreservesCFG();
 212   }
 213   StringRef getPassName() const override {
 214     return "Lower early coroutine intrinsics";
 215   }
 216 };
 217 }
 218
 219 char CoroEarly::ID = 0;
 220 INITIALIZE_PASS(CoroEarly, "coro-early", "Lower early coroutine intrinsics",
 221                 false, false)
 222
 223 Pass *llvm::createCoroEarlyPass() { return new CoroEarly(); }