[Coroutines] Part 10: Add coroutine promise support.

[GorNishanov]

Summary

1. CoroEarly now lowers llvm.coro.promise intrinsic that allows to obtain

a coroutine promise pointer from a coroutine frame and vice versa.

2. CoroFrame now interprets Promise argument of llvm.coro.begin to

place CoroutinPromise alloca at a deterministic offset from the coroutine frame.

Now, the coroutine promise example from docs\Coroutines.rst compiles and produces expected result (see test/Transform/Coroutines/ex4.ll).

[majnemer]

lib/Transforms/Coroutines/CoroEarly.cpp

66

Builder->getInt8Ty() is a little shorter.

68

auto *

lib/Transforms/Coroutines/CoroInstr.h

104–105

This is weird. There can be ptrtoint/inttoptr pairs, etc. Why do you have this restriction?

108

auto *

111–112

Don't use else after return: http://llvm.org/docs/CodingStandards.html#don-t-use-else-after-a-return

242–243

The alignment is signed?

[GorNishanov]

Addressed review feedback.

[GorNishanov]

lib/Transforms/Coroutines/CoroInstr.h

104–105

Mostly I wanted to observe what else could be used to convert alloca %promise-type into void* parameter, so that I can correctly move them after coro.begin (in case when they are used by something other than coro.begin).

coro.begin marks the point where we know the address of the coroutine frame, so anything touching the coroutine frame should appear after it. Since promise alloca always goes to coroutine frame.

I'd like to keep it for now and revisit it in patch 12 where I add code to deal with various instructions that may appear before coro.begin that would need be moved after it. That will give us more tools to deal with ptrtoint/inttoptr possibility.

[majnemer]

LGTM

lib/Transforms/Coroutines/CoroEarly.cpp

 //===- CoroEarly.cpp - Coroutine Early Function Pass ----------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 // This pass lowers coroutine intrinsics that hide the details of the exact
 // calling convention for coroutine resume and destroy functions and details of
 // the structure of the coroutine frame.
 //===----------------------------------------------------------------------===//
 
 #include "CoroInternal.h"
 #include "llvm/IR/CallSite.h"
+#include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Pass.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "coro-early"
 
 namespace {
 // Created on demand if CoroEarly pass has work to do.
 class Lowerer : public coro::LowererBase {
+  IRBuilder<> Builder;
+  PointerType *AnyResumeFnPtrTy;
+
   void lowerResumeOrDestroy(CallSite CS, CoroSubFnInst::ResumeKind);
+  void lowerCoroPromise(CoroPromiseInst *Intrin);
 
 public:
-  Lowerer(Module &M) : LowererBase(M) {}
+  Lowerer(Module &M)
+      : LowererBase(M), Builder(Context),
+        AnyResumeFnPtrTy(FunctionType::get(Type::getVoidTy(Context), Int8Ptr,
+                                           /*isVarArg=*/false)
+                             ->getPointerTo()) {}
   bool lowerEarlyIntrinsics(Function &F);
 };
 }
 
 // Replace a direct call to coro.resume or coro.destroy with an indirect call to
 // an address returned by coro.subfn.addr intrinsic. This is done so that
 // CGPassManager recognizes devirtualization when CoroElide pass replaces a call
 // to coro.subfn.addr with an appropriate function address.
 void Lowerer::lowerResumeOrDestroy(CallSite CS,
                                    CoroSubFnInst::ResumeKind Index) {
   Value *ResumeAddr =
       makeSubFnCall(CS.getArgOperand(0), Index, CS.getInstruction());
   CS.setCalledFunction(ResumeAddr);
   CS.setCallingConv(CallingConv::Fast);
 }
 
+// Coroutine promise field is always at the fixed offset from the beginning of
+// the coroutine frame. i8* coro.promise(i8*, i1 from) intrinsic adds an offset
+// to a passed pointer to move from coroutine frame to coroutine promise and
+// vice versa. Since we don't know exactly which coroutine frame it is, we build
+// a coroutine frame mock up starting with two function pointers, followed by a
+// properly aligned coroutine promise field.
+// TODO: Handle the case when coroutine promise alloca has align override.
+void Lowerer::lowerCoroPromise(CoroPromiseInst *Intrin) {
+  Value *Operand = Intrin->getArgOperand(0);
+  unsigned Alignement = Intrin->getAlignment();
+  Type *Int8Ty = Builder.getInt8Ty();
+
+  auto *SampleStruct =
+      StructType::get(Context, {AnyResumeFnPtrTy, AnyResumeFnPtrTy, Int8Ty});
+  const DataLayout &DL = TheModule.getDataLayout();
+  int64_t Offset = alignTo(
+      DL.getStructLayout(SampleStruct)->getElementOffset(2), Alignement);
+  if (Intrin->isFromPromise())
+    Offset = -Offset;
+
+  Builder.SetInsertPoint(Intrin);
+  Value *Replacement =
+      Builder.CreateConstInBoundsGEP1_32(Int8Ty, Operand, Offset);
+
+  Intrin->replaceAllUsesWith(Replacement);
+  Intrin->eraseFromParent();
+}
+
 // Prior to CoroSplit, calls to coro.begin needs to be marked as NoDuplicate,
 // as CoroSplit assumes there is exactly one coro.begin. After CoroSplit,
 // NoDuplicate attribute will be removed from coro.begin otherwise, it will
 // interfere with inlining.
 static void setCannotDuplicate(CoroIdInst *CoroId) {
   for (User *U : CoroId->users())
     if (auto *CB = dyn_cast<CoroBeginInst>(U))
       CB->setCannotDuplicate();
if (auto CS = CallSite(&I)) {
         }
         break;
       case Intrinsic::coro_resume:
         lowerResumeOrDestroy(CS, CoroSubFnInst::ResumeIndex);
         break;
       case Intrinsic::coro_destroy:
         lowerResumeOrDestroy(CS, CoroSubFnInst::DestroyIndex);
         break;
+      case Intrinsic::coro_promise:
+        lowerCoroPromise(cast<CoroPromiseInst>(&I));
+        break;
       }
       Changed = true;
     }
   }
   return Changed;
 }
 
 //===----------------------------------------------------------------------===//

lib/Transforms/Coroutines/CoroFrame.cpp

static StructType *buildFrameType(Function &F, coro::Shape &Shape,
   StructType *FrameTy = StructType::create(C, Name);
   auto *FramePtrTy = FrameTy->getPointerTo();
   auto *FnTy = FunctionType::get(Type::getVoidTy(C), FramePtrTy,
                                  /*IsVarArgs=*/false);
   auto *FnPtrTy = FnTy->getPointerTo();
 
   // Figure out how wide should be an integer type storing the suspend index.
   unsigned IndexBits = std::max(1U, Log2_64_Ceil(Shape.CoroSuspends.size()));
-
-  SmallVector<Type *, 8> Types{FnPtrTy, FnPtrTy, Type::getIntNTy(C, IndexBits)};
+  Type *PromiseType = Shape.PromiseAlloca
+                          ? Shape.PromiseAlloca->getType()->getElementType()
+                          : Type::getInt1Ty(C);
+  SmallVector<Type *, 8> Types{FnPtrTy, FnPtrTy, PromiseType,
+                               Type::getIntNTy(C, IndexBits)};
   Value *CurrentDef = nullptr;
 
   // Create an entry for every spilled value.
   for (auto const &S : Spills) {
     if (CurrentDef == S.def())
       continue;
 
     CurrentDef = S.def();
+    // PromiseAlloca was already added to Types array earlier.
+    if (CurrentDef == Shape.PromiseAlloca)
+      continue;
 
     Type *Ty = nullptr;
     if (auto *AI = dyn_cast<AllocaInst>(CurrentDef))
       Ty = AI->getAllocatedType();
     else
       Ty = CurrentDef->getType();
 
     Types.push_back(Ty);
static Instruction *insertSpills(SpillInfo &Spills, coro::Shape &Shape) {
   unsigned Index = coro::Shape::LastKnownField;
 
   // We need to keep track of any allocas that need "spilling"
   // since they will live in the coroutine frame now, all access to them
   // need to be changed, not just the access across suspend points
   // we remember allocas and their indices to be handled once we processed
   // all the spills.
   SmallVector<std::pair<AllocaInst *, unsigned>, 4> Allocas;
+  // Promise alloca (if present) has a fixed field number (Shape::PromiseField)
+  if (Shape.PromiseAlloca)
+    Allocas.emplace_back(Shape.PromiseAlloca, coro::Shape::PromiseField);
 
   // Create a load instruction to reload the spilled value from the coroutine
   // frame.
   auto CreateReload = [&](Instruction *InsertBefore) {
     Builder.SetInsertPoint(InsertBefore);
     auto *G = Builder.CreateConstInBoundsGEP2_32(FrameTy, FramePtr, 0, Index,
                                                  CurrentValue->getName() +
                                                      Twine(".reload.addr"));
     return isa<AllocaInst>(CurrentValue)
                ? G
                : Builder.CreateLoad(G,
                                     CurrentValue->getName() + Twine(".reload"));
   };
 
   for (auto const &E : Spills) {
     // If we have not seen the value, generate a spill.
     if (CurrentValue != E.def()) {
       CurrentValue = E.def();
       CurrentBlock = nullptr;
       CurrentReload = nullptr;
 
       ++Index;
 
       if (auto *AI = dyn_cast<AllocaInst>(CurrentValue)) {
-        // Spiled AllocaInst will be replaced with GEP from the coroutine frame
+        // Spilled AllocaInst will be replaced with GEP from the coroutine frame
         // there is no spill required.
         Allocas.emplace_back(AI, Index);
         if (!AI->isStaticAlloca())
           report_fatal_error("Coroutines cannot handle non static allocas yet");
       } else {
         // Otherwise, create a store instruction storing the value into the
         // coroutine frame. For, argument, we will place the store instruction
         // right after the coroutine frame pointer instruction, i.e. bitcase of
static Instruction *insertSpills(SpillInfo &Spills, coro::Shape &Shape) {
   Shape.AllocaSpillBlock->splitBasicBlock(&Shape.AllocaSpillBlock->front(),
                                           "PostSpill");
 
   Builder.SetInsertPoint(&Shape.AllocaSpillBlock->front());
   // If we found any allocas, replace all of their remaining uses with Geps.
   for (auto &P : Allocas) {
     auto *G =
         Builder.CreateConstInBoundsGEP2_32(FrameTy, FramePtr, 0, P.second);
-    ReplaceInstWithInst(P.first, cast<Instruction>(G));
+    // We are not using ReplaceInstWithInst(P.first, cast<Instruction>(G)) here,
+    // as we are changing location of the instruction.
+    G->takeName(P.first);
+    P.first->replaceAllUsesWith(G);
+    P.first->eraseFromParent();
   }
   return FramePtr;
 }
 
 static void rewritePHIs(BasicBlock &BB) {
   // For every incoming edge we will create a block holding all
   // incoming values in a single PHI nodes.
   //

 // Split above and below a particular instruction so that it
 // will be all alone by itself in a block.
 static void splitAround(Instruction *I, const Twine &Name) {
   splitBlockIfNotFirst(I, Name);
   splitBlockIfNotFirst(I->getNextNode(), "After" + Name);
 }
 
 void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
+  Shape.PromiseAlloca = Shape.CoroBegin->getId()->getPromise();
+  if (Shape.PromiseAlloca) {
+    Shape.CoroBegin->getId()->clearPromise();
+  }
 
   // Make sure that all coro.saves and the fallthrough coro.end are in their
   // own block to simplify the logic of building up SuspendCrossing data.
   for (CoroSuspendInst *CSI : Shape.CoroSuspends)
     splitAround(CSI->getCoroSave(), "CoroSave");
 
   // Put fallthrough CoroEnd into its own block. Note: Shape::buildFrom places
   // the fallthrough coro.end as the first element of CoroEnds array.
for (Instruction &I : instructions(F)) {
     // parameter to .resume and .cleanup parts and should not go into coroutine
     // frame.
     if (isa<CoroBeginInst>(&I))
       continue;
     // A token returned CoroIdInst is used to tie together structural intrinsics
     // in a coroutine. It should not be saved to the coroutine frame.
     if (isa<CoroIdInst>(&I))
       continue;
+    // The Coroutine Promise always included into coroutine frame, no need to
+    // check for suspend crossing.
+    if (Shape.PromiseAlloca == &I)
+      continue;
 
     for (User *U : I.users())
       if (Checker.isDefinitionAcrossSuspend(I, U)) {
         // We cannot spill a token.
         if (I.getType()->isTokenTy())
           report_fatal_error(
               "token definition is separated from the use by a suspend point");
         assert(!materializable(I) &&

lib/Transforms/Coroutines/CoroInstr.h

public:
   }
 };
 
 /// This represents the llvm.coro.alloc instruction.
 class LLVM_LIBRARY_VISIBILITY CoroIdInst : public IntrinsicInst {
   enum { AlignArg, PromiseArg, CoroutineArg, InfoArg };
 
 public:
+  IntrinsicInst *getCoroBegin() {
+    for (User *U : users())
+      if (auto *II = dyn_cast<IntrinsicInst>(U))
+        if (II->getIntrinsicID() == Intrinsic::coro_begin)
+          return II;
+    llvm_unreachable("no coro.begin associated with coro.id");
+  }
+
+  AllocaInst *getPromise() const {
+    Value *Arg = getArgOperand(PromiseArg);
+    return isa<ConstantPointerNull>(Arg)
+               ? nullptr
+               : cast<AllocaInst>(Arg->stripPointerCasts());
+  }
+
+  void clearPromise() {
+    Value *Arg = getArgOperand(PromiseArg);
+    setArgOperand(PromiseArg,
+                  ConstantPointerNull::get(Type::getInt8PtrTy(getContext())));
+    if (isa<AllocaInst>(Arg))
+      return;
+    assert((isa<BitCastInst>(Arg) || isa<GetElementPtrInst>(Arg)) &&
+           "unexpected instruction designating the promise");
+    // TODO: Add a check that any remaining users of Inst are after coro.begin
+    // or add code to move the users after coro.begin.
+    auto *Inst = cast<Instruction>(Arg);
+    if (Inst->use_empty()) {
+      Inst->eraseFromParent();
+      return;
+    }
+    Inst->moveBefore(getCoroBegin()->getNextNode());
+  }
+
   // Info argument of coro.id is
   //   fresh out of the frontend: null ;
   //   outlined                 : {Init, Return, Susp1, Susp2, ...} ;
   //   postsplit                : [resume, destroy, cleanup] ;
   //
   // If parts of the coroutine were outlined to protect against undesirable
   // code motion, these functions will be stored in a struct literal referred to
   // by the Info parameter. Note: this is only needed before coroutine is split.
public:
   static inline bool classof(const IntrinsicInst *I) {
     return I->getIntrinsicID() == Intrinsic::coro_save;
   }
   static inline bool classof(const Value *V) {
     return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
   }
 };
 
+/// This represents the llvm.coro.promise instruction.
+class LLVM_LIBRARY_VISIBILITY CoroPromiseInst : public IntrinsicInst {
+  enum { FrameArg, AlignArg, FromArg };
+
+public:
+  bool isFromPromise() const {
+    return cast<Constant>(getArgOperand(FromArg))->isOneValue();
+  }
+  unsigned getAlignment() const {
+    return cast<ConstantInt>(getArgOperand(AlignArg))->getZExtValue();
+  }
+
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static inline bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_promise;
+  }
+  static inline bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
 /// This represents the llvm.coro.suspend instruction.
 class LLVM_LIBRARY_VISIBILITY CoroSuspendInst : public IntrinsicInst {
   enum { SaveArg, FinalArg };
 
 public:
   CoroSaveInst *getCoroSave() const {
     Value *Arg = getArgOperand(SaveArg);
     if (auto *SI = dyn_cast<CoroSaveInst>(Arg))

lib/Transforms/Coroutines/CoroInternal.h

struct LLVM_LIBRARY_VISIBILITY Shape {
   SmallVector<CoroEndInst *, 4> CoroEnds;
   SmallVector<CoroSizeInst *, 2> CoroSizes;
   SmallVector<CoroSuspendInst *, 4> CoroSuspends;
 
   // Field Indexes for known coroutine frame fields.
   enum {
     ResumeField,
     DestroyField,
+    PromiseField,
     IndexField,
     LastKnownField = IndexField
   };
 
   StructType *FrameTy;
   Instruction *FramePtr;
-  BasicBlock* AllocaSpillBlock;
-  SwitchInst* ResumeSwitch;
+  BasicBlock *AllocaSpillBlock;
+  SwitchInst *ResumeSwitch;
+  AllocaInst *PromiseAlloca;
   bool HasFinalSuspend;
 
-  IntegerType* getIndexType() const {
+  IntegerType *getIndexType() const {
     assert(FrameTy && "frame type not assigned");
     return cast<IntegerType>(FrameTy->getElementType(IndexField));
   }
   ConstantInt *getIndex(uint64_t Value) const {
     return ConstantInt::get(getIndexType(), Value);
   }
 
   Shape() = default;
   explicit Shape(Function &F) { buildFrom(F); }
   void buildFrom(Function &F);
 };
 
-void buildCoroutineFrame(Function& F, Shape& Shape);
+void buildCoroutineFrame(Function &F, Shape &Shape);
 
 } // End namespace coro.
 } // End namespace llvm
 
 #endif

lib/Transforms/Coroutines/Coroutines.cpp

static void clear(coro::Shape &Shape) {
   Shape.CoroEnds.clear();
   Shape.CoroSizes.clear();
   Shape.CoroSuspends.clear();
 
   Shape.FrameTy = nullptr;
   Shape.FramePtr = nullptr;
   Shape.AllocaSpillBlock = nullptr;
   Shape.ResumeSwitch = nullptr;
+  Shape.PromiseAlloca = nullptr;
   Shape.HasFinalSuspend = false;
 }
 
 static CoroSaveInst *createCoroSave(CoroBeginInst *CoroBegin,
                                     CoroSuspendInst *SuspendInst) {
   Module *M = SuspendInst->getModule();
   auto *Fn = Intrinsic::getDeclaration(M, Intrinsic::coro_save);
   auto *SaveInst =

test/Transforms/Coroutines/ex4.ll

+; Fourth example from Doc/Coroutines.rst (coroutine promise)
+; RUN: opt < %s -O2 -enable-coroutines -S | FileCheck %s
+
+define i8* @f(i32 %n) {
+entry:
+  %promise = alloca i32
+  %pv = bitcast i32* %promise to i8*
+  %id = call token @llvm.coro.id(i32 0, i8* %pv, i8* null, i8* null)
+  %need.dyn.alloc = call i1 @llvm.coro.alloc(token %id)
+  br i1 %need.dyn.alloc, label %dyn.alloc, label %coro.begin
+dyn.alloc:
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call i8* @malloc(i32 %size)
+  br label %coro.begin
+coro.begin:
+  %phi = phi i8* [ null, %entry ], [ %alloc, %dyn.alloc ]
+  %hdl = call noalias i8* @llvm.coro.begin(token %id, i8* %phi)
+  br label %loop
+loop:
+  %n.val = phi i32 [ %n, %coro.begin ], [ %inc, %loop ]
+  %inc = add nsw i32 %n.val, 1
+  store i32 %n.val, i32* %promise
+  %0 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %0, label %suspend [i8 0, label %loop
+                                i8 1, label %cleanup]
+cleanup:
+  %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
+  call void @free(i8* %mem)
+  br label %suspend
+suspend:
+  call void @llvm.coro.end(i8* %hdl, i1 false)
+  ret i8* %hdl
+}
+
+; CHECK-LABEL: @main
+define i32 @main() {
+entry:
+  %hdl = call i8* @f(i32 4)
+  %promise.addr.raw = call i8* @llvm.coro.promise(i8* %hdl, i32 4, i1 false)
+  %promise.addr = bitcast i8* %promise.addr.raw to i32*
+  %val0 = load i32, i32* %promise.addr
+  call void @print(i32 %val0)
+  call void @llvm.coro.resume(i8* %hdl)
+  %val1 = load i32, i32* %promise.addr
+  call void @print(i32 %val1)
+  call void @llvm.coro.resume(i8* %hdl)
+  %val2 = load i32, i32* %promise.addr
+  call void @print(i32 %val2)
+  call void @llvm.coro.destroy(i8* %hdl)
+  ret i32 0
+; CHECK:      call void @print(i32 4)
+; CHECK-NEXT: call void @print(i32 5)
+; CHECK-NEXT: call void @print(i32 6)
+; CHECK:      ret i32 0
+}
+
+declare i8* @llvm.coro.promise(i8*, i32, i1)
+declare i8* @malloc(i32)
+declare void @free(i8*)
+declare void @print(i32)
+
+declare token @llvm.coro.id(i32, i8*, i8*, i8*)
+declare i1 @llvm.coro.alloc(token)
+declare i32 @llvm.coro.size.i32()
+declare i8* @llvm.coro.begin(token, i8*)
+declare i8 @llvm.coro.suspend(token, i1)
+declare i8* @llvm.coro.free(token, i8*)
+declare void @llvm.coro.end(i8*, i1)
+
+declare void @llvm.coro.resume(i8*)
+declare void @llvm.coro.destroy(i8*)