[coroutines] Store an address of destroy OR cleanup part in the coroutine frame.

[GorNishanov]

Summary

If heap allocation of a coroutine is elided, we need to make sure that we will update an address stored in the coroutine frame from f.destroy to f.cleanup.
Before this change, CoroSplit synthesized these stores after coro.begin:

store void (%f.Frame*)* @f.resume, void (%f.Frame*)** %resume.addr
store void (%f.Frame*)* @f.destroy, void (%f.Frame*)** %destroy.addr

In those cases where we did heap elision, but were not able to devirtualize all indirect calls, destroy call will attempt to "free" the coroutine frame stored on the stack. Oops.

Now we use select to put an appropriate coroutine subfunction in the destroy slot. As bellow:

store void (%f.Frame*)* @f.resume, void (%f.Frame*)** %resume.addr
%0 = select i1 %need.alloc, void (%f.Frame*)* @f.destroy, void (%f.Frame*)* @f.cleanup
store void (%f.Frame*)* %0, void (%f.Frame*)** %destroy.addr

[GorNishanov]

ran through clang-format

[GorNishanov]

Removed unused function that sneaked in here

[GorNishanov]

Remove unused #include InstIterator.h

[majnemer]

LGTM

lib/Transforms/Coroutines/CoroInstr.h

86–87

Couldn't you do:

if (auto *CA = dyn_cast<CoroAllocInst>(II))
  return CA;

[GorNishanov]

lib/Transforms/Coroutines/CoroInstr.h

86–87

I can and I will!

:)

Thank you for the review!

[GorNishanov]

Simplifed getCoroAllocInstr in CoroInstr.h and preparing to land

[GorNishanov]

Closed by commit: [llvm] r283625 - [coroutines] Store an address of destroy OR cleanup part in the coroutine frame.

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 *getCoroAlloc() {
+  CoroAllocInst *getCoroAlloc() {
     for (User *U : users())
       if (auto *II = dyn_cast<IntrinsicInst>(U))
         if (II->getIntrinsicID() == Intrinsic::coro_alloc)
-          return II;
+          return cast<CoroAllocInst>(II);
     return nullptr;
   }
 
   IntrinsicInst *getCoroBegin() {
     for (User *U : users())
       if (auto *II = dyn_cast<IntrinsicInst>(U))
         if (II->getIntrinsicID() == Intrinsic::coro_begin)
           return II;

lib/Transforms/Coroutines/CoroSplit.cpp

 // coroutine subfunctions we extracted before proceeding to the caller of the
 // coroutine.
 //===----------------------------------------------------------------------===//
 
 #include "CoroInternal.h"
 #include "llvm/Analysis/CallGraphSCCPass.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/InstIterator.h"
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
 
 using namespace llvm;
for (CoroSuspendInst *CS : Shape.CoroSuspends) {
     MappedCS->eraseFromParent();
   }
 
   // Remove coro.end intrinsics.
   replaceFallthroughCoroEnd(Shape.CoroEnds.front(), VMap);
   // FIXME: coming in upcoming patches:
   // replaceUnwindCoroEnds(Shape.CoroEnds, VMap);
 
-  // We only store resume(0) and destroy(1) addresses in the coroutine frame.
-  // The cleanup(2) clone is only used during devirtualization when coroutine is
-  // eligible for heap elision and thus does not participate in indirect calls
-  // and does not need its address to be stored in the coroutine frame.
-  if (FnIndex < 2) {
-    // Store the address of this clone in the coroutine frame.
-    Builder.SetInsertPoint(Shape.FramePtr->getNextNode());
-    auto *G = Builder.CreateConstInBoundsGEP2_32(Shape.FrameTy, Shape.FramePtr,
-                                                 0, FnIndex, "fn.addr");
-    Builder.CreateStore(NewF, G);
-  }
-
   // Eliminate coro.free from the clones, replacing it with 'null' in cleanup,
   // to suppress deallocation code.
   coro::replaceCoroFree(cast<CoroIdInst>(VMap[Shape.CoroBegin->getId()]),
                         /*Elide=*/FnIndex == 2);
 
   NewF->setCallingConv(CallingConv::Fast);
 
   return NewF;
auto *GV = new GlobalVariable(*M, ConstVal->getType(), /*isConstant=*/true,
                                 F.getName() + Twine(".resumers"));
 
   // Update coro.begin instruction to refer to this constant.
   LLVMContext &C = F.getContext();
   auto *BC = ConstantExpr::getPointerCast(GV, Type::getInt8PtrTy(C));
   CoroBegin->getId()->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) {
+
+  IRBuilder<> Builder(Shape.FramePtr->getNextNode());
+  auto *ResumeAddr = Builder.CreateConstInBoundsGEP2_32(
+      Shape.FrameTy, Shape.FramePtr, 0, coro::Shape::ResumeField,
+      "resume.addr");
+  Builder.CreateStore(ResumeFn, ResumeAddr);
+
+  Value *DestroyOrCleanupFn = DestroyFn;
+
+  CoroIdInst *CoroId = Shape.CoroBegin->getId();
+  if (CoroAllocInst *CA = CoroId->getCoroAlloc()) {
+    // If there is a CoroAlloc and it returns false (meaning we elide the
+    // allocation, use CleanupFn instead of DestroyFn).
+    DestroyOrCleanupFn = Builder.CreateSelect(CA, DestroyFn, CleanupFn);
+  }
+
+  auto *DestroyAddr = Builder.CreateConstInBoundsGEP2_32(
+      Shape.FrameTy, Shape.FramePtr, 0, coro::Shape::DestroyField,
+      "destroy.addr");
+  Builder.CreateStore(DestroyOrCleanupFn, DestroyAddr);
+}
+
 static void postSplitCleanup(Function &F) {
   removeUnreachableBlocks(F);
   llvm::legacy::FunctionPassManager FPM(F.getParent());
 
   FPM.add(createVerifierPass());
   FPM.add(createSCCPPass());
   FPM.add(createCFGSimplificationPass());
   FPM.add(createEarlyCSEPass());
static void splitCoroutine(Function &F, CallGraph &CG, CallGraphSCC &SCC) {
   // We no longer need coro.end in F.
   removeCoroEnds(Shape);
 
   postSplitCleanup(F);
   postSplitCleanup(*ResumeClone);
   postSplitCleanup(*DestroyClone);
   postSplitCleanup(*CleanupClone);
 
+  // Store addresses resume/destroy/cleanup functions in the coroutine frame.
+  updateCoroFrame(Shape, ResumeClone, DestroyClone, CleanupClone);
+
+  // Create a constant array referring to resume/destroy/clone functions pointed
+  // by the last argument of @llvm.coro.info, so that CoroElide pass can
+  // determined correct function to call.
   setCoroInfo(F, Shape.CoroBegin, {ResumeClone, DestroyClone, CleanupClone});
+
+  // Update call graph and add the functions we created to the SCC.
   coro::updateCallGraph(F, {ResumeClone, DestroyClone, CleanupClone}, CG, SCC);
 }
 
 // When we see the coroutine the first time, we insert an indirect call to a
 // devirt trigger function and mark the coroutine that it is now ready for
 // split.
 static void prepareForSplit(Function &F, CallGraph &CG) {
   Module &M = *F.getParent();

test/Transforms/Coroutines/coro-split-00.ll

 ; Tests that coro-split pass splits the coroutine into f, f.resume and f.destroy
 ; RUN: opt < %s -coro-split -S | FileCheck %s
 
 define i8* @f() "coroutine.presplit"="1" {
 entry:
   %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
+  %need.alloc = call i1 @llvm.coro.alloc(token %id)
+  br i1 %need.alloc, label %dyn.alloc, label %begin
+
+dyn.alloc:  
   %size = call i32 @llvm.coro.size.i32()
   %alloc = call i8* @malloc(i32 %size)
-  %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
+  br label %begin
+
+begin:
+  %phi = phi i8* [ null, %entry ], [ %alloc, %dyn.alloc ]
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* %phi)
   call void @print(i32 0)
   %0 = call i8 @llvm.coro.suspend(token none, i1 false)
   switch i8 %0, label %suspend [i8 0, label %resume 
                                 i8 1, label %cleanup]
 resume:
   call void @print(i32 1)
   br 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 0)  
   ret i8* %hdl
 }
 
 ; CHECK-LABEL: @f(
 ; CHECK: call i8* @malloc
+; CHECK: @llvm.coro.begin(token %id, i8* %phi)
+; CHECK: store void (%f.Frame*)* @f.resume, void (%f.Frame*)** %resume.addr
+; CHECK: %[[SEL:.+]] = select i1 %need.alloc, void (%f.Frame*)* @f.destroy, void (%f.Frame*)* @f.cleanup
+; CHECK: store void (%f.Frame*)* %[[SEL]], void (%f.Frame*)** %destroy.addr
 ; CHECK: call void @print(i32 0)
 ; CHECK-NOT: call void @print(i32 1)
 ; CHECK-NOT: call void @free(
 ; CHECK: ret i8* %hdl
 
 ; CHECK-LABEL: @f.resume(
 ; CHECK-NOT: call i8* @malloc
 ; CHECK-NOT: call void @print(i32 0)
 ; CHECK: call void @print(i32 1)
 ; CHECK-NOT: call void @print(i32 0)
 ; CHECK: call void @free(
 ; CHECK: ret void
 
 ; CHECK-LABEL: @f.destroy(
 ; CHECK-NOT: call i8* @malloc
 ; CHECK-NOT: call void @print(
 ; CHECK: call void @free(
 ; CHECK: ret void
 
+; CHECK-LABEL: @f.cleanup(
+; CHECK-NOT: call i8* @malloc
+; CHECK-NOT: call void @print(
+; CHECK-NOT: call void @free(
+; CHECK: ret void
+
 declare i8* @llvm.coro.free(token, i8*)
 declare i32 @llvm.coro.size.i32()
 declare i8  @llvm.coro.suspend(token, i1)
 declare void @llvm.coro.resume(i8*)
 declare void @llvm.coro.destroy(i8*)
 
 declare token @llvm.coro.id(i32, i8*, i8*, i8*)
-declare i8* @llvm.coro.alloc(token)
+declare i1 @llvm.coro.alloc(token)
 declare i8* @llvm.coro.begin(token, i8*)
 declare void @llvm.coro.end(i8*, i1) 
 
 declare noalias i8* @malloc(i32)
 declare void @print(i32)
 declare void @free(i8*)