[Coroutines] Allow FramePtr to be an Argument

[nikic]

Summary

With opaque pointers, after splitRetconCoroutine() the FramePtr may be an Argument rather than an Instruction. With typed pointers, this currently doesn't happen because the FramePtr would be a bitcast instruction.

Fix this by making FramePtr a Value and adding a helper for the "after FramePtr" insertion pointer, which would be the start of the function in the Argument case.

[ChuanqiXu]

LGTM except a nit.

llvm/lib/Transforms/Coroutines/CoroInternal.h

270

I prefer getInsertPtAfterFramePtr.

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

SmallVector<Metadata *, 32> RetainedNodesVec(RetainedNodes.begin(),
                                                  RetainedNodes.end());
     RetainedNodesVec.push_back(FrameDIVar);
     SubProgram->replaceOperandWith(
         7, (MDTuple::get(F.getContext(), RetainedNodesVec)));
   }
 
   DBuilder.insertDeclare(Shape.FramePtr, FrameDIVar,
                          DBuilder.createExpression(), DILoc,
-                         Shape.FramePtr->getNextNode());
+                         Shape.getInsertPtAfterFramePtr());
 }
 
 // Build a struct that will keep state for an active coroutine.
 //   struct f.frame {
 //     ResumeFnTy ResumeFnAddr;
 //     ResumeFnTy DestroyFnAddr;
 //     int ResumeIndex;
 //     ... promise (if present) ...

 //    whatever
 //
 //
 static void insertSpills(const FrameDataInfo &FrameData, coro::Shape &Shape) {
   auto *CB = Shape.CoroBegin;
   LLVMContext &C = CB->getContext();
   IRBuilder<> Builder(C);
   StructType *FrameTy = Shape.FrameTy;
-  Instruction *FramePtr = Shape.FramePtr;
+  Value *FramePtr = Shape.FramePtr;
   DominatorTree DT(*CB->getFunction());
   SmallDenseMap<llvm::Value *, llvm::AllocaInst *, 4> DbgPtrAllocaCache;
 
   // Create a GEP with the given index into the coroutine frame for the original
   // value Orig. Appends an extra 0 index for array-allocas, preserving the
   // original type.
   auto GetFramePointer = [&](Value *Orig) -> Value * {
     FieldIDType Index = FrameData.getFieldIndex(Orig);
for (auto const &E : FrameData.Spills) {
     // Create a store instruction storing the value into the
     // coroutine frame.
     Instruction *InsertPt = nullptr;
     Type *ByValTy = nullptr;
     if (auto *Arg = dyn_cast<Argument>(Def)) {
       // For arguments, we will place the store instruction right after
       // the coroutine frame pointer instruction, i.e. bitcast of
       // coro.begin from i8* to %f.frame*.
-      InsertPt = FramePtr->getNextNode();
+      InsertPt = Shape.getInsertPtAfterFramePtr();
 
       // If we're spilling an Argument, make sure we clear 'nocapture'
       // from the coroutine function.
       Arg->getParent()->removeParamAttr(Arg->getArgNo(), Attribute::NoCapture);
 
       if (Arg->hasByValAttr())
         ByValTy = Arg->getParamByValType();
     } else if (auto *CSI = dyn_cast<AnyCoroSuspendInst>(Def)) {
       // Don't spill immediately after a suspend; splitting assumes
       // that the suspend will be followed by a branch.
       InsertPt = CSI->getParent()->getSingleSuccessor()->getFirstNonPHI();
     } else {
       auto *I = cast<Instruction>(Def);
       if (!DT.dominates(CB, I)) {
         // If it is not dominated by CoroBegin, then spill should be
         // inserted immediately after CoroFrame is computed.
-        InsertPt = FramePtr->getNextNode();
+        InsertPt = Shape.getInsertPtAfterFramePtr();
       } else if (auto *II = dyn_cast<InvokeInst>(I)) {
         // If we are spilling the result of the invoke instruction, split
         // the normal edge and insert the spill in the new block.
         auto *NewBB = SplitEdge(II->getParent(), II->getNormalDest());
         InsertPt = NewBB->getTerminator();
       } else if (isa<PHINode>(I)) {
         // Skip the PHINodes and EH pads instructions.
         BasicBlock *DefBlock = I->getParent();
for (auto *U : E.second) {
       }
 
       // Replace all uses of CurrentValue in the current instruction with
       // reload.
       U->replaceUsesOfWith(Def, CurrentReload);
     }
   }
 
-  BasicBlock *FramePtrBB = FramePtr->getParent();
+  BasicBlock *FramePtrBB = Shape.getInsertPtAfterFramePtr()->getParent();
 
-  auto SpillBlock =
-      FramePtrBB->splitBasicBlock(FramePtr->getNextNode(), "AllocaSpillBB");
+  auto SpillBlock = FramePtrBB->splitBasicBlock(
+      Shape.getInsertPtAfterFramePtr(), "AllocaSpillBB");
   SpillBlock->splitBasicBlock(&SpillBlock->front(), "PostSpill");
   Shape.AllocaSpillBlock = SpillBlock;
 
   // retcon and retcon.once lowering assumes all uses have been sunk.
   if (Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce ||
       Shape.ABI == coro::ABI::Async) {
     // If we found any allocas, replace all of their remaining uses with Geps.
     Builder.SetInsertPoint(&SpillBlock->front());
for (const auto &A : FrameData.Allocas) {
     SmallVector<DbgVariableIntrinsic *, 4> DIs;
     findDbgUsers(DIs, Alloca);
     for (auto *DVI : DIs)
       DVI->replaceUsesOfWith(Alloca, G);
 
     for (Instruction *I : UsersToUpdate)
       I->replaceUsesOfWith(Alloca, G);
   }
-  Builder.SetInsertPoint(FramePtr->getNextNode());
+  Builder.SetInsertPoint(Shape.getInsertPtAfterFramePtr());
   for (const auto &A : FrameData.Allocas) {
     AllocaInst *Alloca = A.Alloca;
     if (A.MayWriteBeforeCoroBegin) {
       // isEscaped really means potentially modified before CoroBegin.
       if (Alloca->isArrayAllocation())
         report_fatal_error(
             "Coroutines cannot handle copying of array allocas yet");
 

llvm/lib/Transforms/Coroutines/CoroInternal.h

struct SwitchFieldIndex {
     };
   };
 
   coro::ABI ABI;
 
   StructType *FrameTy;
   Align FrameAlign;
   uint64_t FrameSize;
-  Instruction *FramePtr;
+  Value *FramePtr;
   BasicBlock *AllocaSpillBlock;
 
   /// This would only be true if optimization are enabled.
   bool OptimizeFrame;
 
   struct SwitchLoweringStorage {
     SwitchInst *ResumeSwitch;
     AllocaInst *PromiseAlloca;
struct LLVM_LIBRARY_VISIBILITY Shape {
   }
 
   AllocaInst *getPromiseAlloca() const {
     if (ABI == coro::ABI::Switch)
       return SwitchLowering.PromiseAlloca;
     return nullptr;
   }
 
+  Instruction *getInsertPtAfterFramePtr() const {
+    if (auto *I = dyn_cast<Instruction>(FramePtr))
+      return I->getNextNode();
+    return &cast<Argument>(FramePtr)->getParent()->getEntryBlock().front();
+  }
+
   /// Allocate memory according to the rules of the active lowering.
   ///
   /// \param CG - if non-null, will be updated for the new call
   Value *emitAlloc(IRBuilder<> &Builder, Value *Size, CallGraph *CG) const;
 
   /// Deallocate memory according to the rules of the active lowering.
   ///
   /// \param CG - if non-null, will be updated for the new call

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

static void setCoroInfo(Function &F, coro::Shape &Shape,
   Shape.getSwitchCoroId()->setInfo(BC);
 }
 
 // Store addresses of Resume/Destroy/Cleanup functions in the coroutine frame.
 static void updateCoroFrame(coro::Shape &Shape, Function *ResumeFn,
                             Function *DestroyFn, Function *CleanupFn) {
   assert(Shape.ABI == coro::ABI::Switch);
 
-  IRBuilder<> Builder(Shape.FramePtr->getNextNode());
+  IRBuilder<> Builder(Shape.getInsertPtAfterFramePtr());
+
   auto *ResumeAddr = Builder.CreateStructGEP(
       Shape.FrameTy, Shape.FramePtr, coro::Shape::SwitchFieldIndex::Resume,
       "resume.addr");
   Builder.CreateStore(ResumeFn, ResumeAddr);
 
   Value *DestroyOrCleanupFn = DestroyFn;
 
   CoroIdInst *CoroId = Shape.getSwitchCoroId();
static void splitAsyncCoroutine(Function &F, coro::Shape &Shape,
   FramePtr = Builder.CreateBitOrPointerCast(FramePtr, Int8PtrTy);
   FramePtr = Builder.CreateConstInBoundsGEP1_32(
       Type::getInt8Ty(Context), FramePtr, Shape.AsyncLowering.FrameOffset,
       "async.ctx.frameptr");
 
   // Map all uses of llvm.coro.begin to the allocated frame pointer.
   {
     // Make sure we don't invalidate Shape.FramePtr.
-    TrackingVH<Instruction> Handle(Shape.FramePtr);
+    TrackingVH<Value> Handle(Shape.FramePtr);
     Shape.CoroBegin->replaceAllUsesWith(FramePtr);
     Shape.FramePtr = Handle.getValPtr();
   }
 
   // Create all the functions in order after the main function.
   auto NextF = std::next(F.getIterator());
 
   // Create a continuation function for each of the suspend points.
if (Shape.RetconLowering.IsFrameInlineInStorage) {
     auto Dest = Builder.CreateBitCast(Id->getStorage(),
                                       RawFramePtr->getType()->getPointerTo());
     Builder.CreateStore(RawFramePtr, Dest);
   }
 
   // Map all uses of llvm.coro.begin to the allocated frame pointer.
   {
     // Make sure we don't invalidate Shape.FramePtr.
-    TrackingVH<Instruction> Handle(Shape.FramePtr);
+    TrackingVH<Value> Handle(Shape.FramePtr);
     Shape.CoroBegin->replaceAllUsesWith(RawFramePtr);
     Shape.FramePtr = Handle.getValPtr();
   }
 
   // Create a unique return block.
   BasicBlock *ReturnBB = nullptr;
   SmallVector<PHINode *, 4> ReturnPHIs;
 

llvm/test/Transforms/Coroutines/coro-retcon-opaque-ptr.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -enable-coroutines -passes='default<O2>' -opaque-pointers -S | FileCheck %s
+
+; Same test as coro-retcon.ll, but with opaque pointers enabled.
+
+define ptr @f(ptr %buffer, i32 %n) {
+; CHECK-LABEL: @f(
+; CHECK-NEXT:  coro.return:
+; CHECK-NEXT:    store i32 [[N:%.*]], ptr [[BUFFER:%.*]], align 4
+; CHECK-NEXT:    tail call void @print(i32 [[N]])
+; CHECK-NEXT:    ret ptr @f.resume.0
+;
+entry:
+  %id = call token @llvm.coro.id.retcon(i32 8, i32 4, ptr %buffer, ptr @prototype, ptr @allocate, ptr @deallocate)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
+  br label %loop
+
+loop:                                             ; preds = %resume, %entry
+  %n.val = phi i32 [ %n, %entry ], [ %inc, %resume ]
+  call void @print(i32 %n.val)
+  %unwind0 = call i1 (...) @llvm.coro.suspend.retcon.i1()
+  br i1 %unwind0, label %cleanup, label %resume
+
+resume:                                           ; preds = %loop
+  %inc = add i32 %n.val, 1
+  br label %loop
+
+cleanup:                                          ; preds = %loop
+  %0 = call i1 @llvm.coro.end(ptr %hdl, i1 false)
+  unreachable
+}
+
+define i32 @main() {
+; CHECK-LABEL: @main(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = alloca [8 x i8], align 4
+; CHECK-NEXT:    store i32 4, ptr [[TMP0]], align 4
+; CHECK-NEXT:    call void @print(i32 4)
+; CHECK-NEXT:    call void @llvm.experimental.noalias.scope.decl(metadata [[META0:![0-9]+]])
+; CHECK-NEXT:    [[N_VAL_RELOAD_I:%.*]] = load i32, ptr [[TMP0]], align 4, !alias.scope !0
+; CHECK-NEXT:    [[INC_I:%.*]] = add i32 [[N_VAL_RELOAD_I]], 1
+; CHECK-NEXT:    store i32 [[INC_I]], ptr [[TMP0]], align 4, !alias.scope !0
+; CHECK-NEXT:    call void @print(i32 [[INC_I]]), !noalias !0
+; CHECK-NEXT:    call void @llvm.experimental.noalias.scope.decl(metadata [[META3:![0-9]+]])
+; CHECK-NEXT:    [[N_VAL_RELOAD_I1:%.*]] = load i32, ptr [[TMP0]], align 4, !alias.scope !3
+; CHECK-NEXT:    [[INC_I2:%.*]] = add i32 [[N_VAL_RELOAD_I1]], 1
+; CHECK-NEXT:    call void @print(i32 [[INC_I2]]), !noalias !3
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  %0 = alloca [8 x i8], align 4
+  %prepare = call ptr @llvm.coro.prepare.retcon(ptr @f)
+  %cont0 = call ptr %prepare(ptr %0, i32 4)
+  %cont1 = call ptr %cont0(ptr %0, i1 zeroext false)
+  %cont2 = call ptr %cont1(ptr %0, i1 zeroext false)
+  %1 = call ptr %cont2(ptr %0, i1 zeroext true)
+  ret i32 0
+}
+
+define hidden { ptr, ptr } @g(ptr %buffer, ptr %ptr) {
+; CHECK-LABEL: @g(
+; CHECK-NEXT:  coro.return:
+; CHECK-NEXT:    [[TMP0:%.*]] = tail call ptr @allocate(i32 8)
+; CHECK-NEXT:    store ptr [[TMP0]], ptr [[BUFFER:%.*]], align 8
+; CHECK-NEXT:    store ptr [[PTR:%.*]], ptr [[TMP0]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = insertvalue { ptr, ptr } { ptr @g.resume.0, ptr undef }, ptr [[PTR]], 1
+; CHECK-NEXT:    ret { ptr, ptr } [[TMP1]]
+;
+entry:
+  %id = call token @llvm.coro.id.retcon(i32 8, i32 4, ptr %buffer, ptr @g_prototype, ptr @allocate, ptr @deallocate)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
+  br label %loop
+
+loop:                                             ; preds = %resume, %entry
+  %unwind0 = call i1 (...) @llvm.coro.suspend.retcon.i1(ptr %ptr)
+  br i1 %unwind0, label %cleanup, label %resume
+
+resume:                                           ; preds = %loop
+  br label %loop
+
+cleanup:                                          ; preds = %loop
+  %0 = call i1 @llvm.coro.end(ptr %hdl, i1 false)
+  unreachable
+}
+
+declare token @llvm.coro.id.retcon(i32, i32, i8*, i8*, i8*, i8*)
+declare i8* @llvm.coro.begin(token, i8*)
+declare i1 @llvm.coro.suspend.retcon.i1(...)
+declare i1 @llvm.coro.end(i8*, i1)
+declare i8* @llvm.coro.prepare.retcon(i8*)
+
+declare i8* @prototype(i8*, i1 zeroext)
+declare {i8*,i8*} @g_prototype(i8*, i1 zeroext)
+
+declare noalias i8* @allocate(i32 %size)
+declare void @deallocate(i8* %ptr)
+
+declare void @print(i32)