llvm/lib/Transforms/Coroutines/CoroEarly.cpp

   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
  17 using namespace llvm;
  18
  19 #define DEBUG_TYPE "coro-early"
  20
  21 namespace {
  22 // Created on demand if the coro-early pass has work to do.
  23 class 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
  33 public:
  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.
  47 void 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.
  61 void 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.
  86 void 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
 104 static 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
 123 void 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.
 163 static void setCannotDuplicate(CoroIdInst *CoroId) {
 164   for (User *U : CoroId->users())
 165     if (auto *CB = dyn_cast<CoroBeginInst>(U))
 166       CB->setCannotDuplicate();
 167 }
 168
 169 void 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                    "\"presplitcoroutine\" 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
 249 static 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
 258 PreservedAnalyses 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 }