[coroutines] Always place escaped allocas into the coroutine frame

[GorNishanov]

Summary

An alloca that escaped, might be accessed after resume via code external to the coroutine.
Therefore even if we do not see a use of an alloca after resume in the coroutine body,
we still need to place it into the coroutine frame.

In the follwoing example, %sneaky needs to go into the coroutine frame.

  %sneaky = alloca i32
...
  %hideit = ptrtoint i32* %sneaky to i64
  call void @escape(i64 %hideit)
...
  %tok = call i8 @llvm.coro.suspend(token none, i1 false)

  ... not used after suspend ...

[GorNishanov]

gentle ping

[lewissbaker]

lib/Transforms/Coroutines/CoroFrame.cpp

889

Does this mean that passing the address of any allocad memory to a function call will cause it to be spilled even if its lifetime does not extend past a suspension point?

I'm thinking of cases where we have:

{
  co_await something;
  {
    int x = 123;
    escapes(&x);
  }
  co_await somethingElse;
}

Am I correct in understanding that with the proposed logic that x would be spilled to the frame even thought its lifetime does not cross a suspension-point? Can we avoid that?

Should we be checking for llvm.lifetime.start / llvm.lifetime.end intrinsics and if found only considering the result of allocaEscapes(&I) if the lifetime start/end span contains a suspension point?

[lewissbaker]

lib/Transforms/Coroutines/CoroFrame.cpp

889

You're calling allocaEscapes(&I) for each User in I.users() here but the call doesn't seem to depend on the user.

Can the call to allocaEscapes(&I) be moved out of the loop and just called once?

lib/Transforms/Coroutines/CoroFrame.cpp

 // GEP + load from the coroutine frame. At the point of the definition we spill
 // the value into the coroutine frame.
 //
 // TODO: pack values tightly using liveness info.
 //===----------------------------------------------------------------------===//
 
 #include "CoroInternal.h"
 #include "llvm/ADT/BitVector.h"
+#include "llvm/Analysis/PtrUseVisitor.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/circular_raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
if (auto *PN = dyn_cast<PHINode>(E.user())) {
       continue;
     }
 
     // Replace all uses of CurrentValue in the current instruction with reload.
     E.user()->replaceUsesOfWith(CurrentValue, CurrentReload);
   }
 
   BasicBlock *FramePtrBB = FramePtr->getParent();
+
+  // Create AllocaSpillBlock that will become the new entry block for resume
+  // and destroy parts of the coroutine. All the alloca spills go here.
   Shape.AllocaSpillBlock =
       FramePtrBB->splitBasicBlock(FramePtr->getNextNode(), "AllocaSpillBB");
   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) {
for (auto const &E : Spills) {
     }
   }
 
   Instruction *InsertPt = CoroBegin->getNextNode();
   for (Instruction *I : NeedsMoving)
     I->moveBefore(InsertPt);
 }
 
+// Checks if alloca escapes.
+static bool allocaEscapes(Instruction *I) {
+  auto *AI = dyn_cast<AllocaInst>(I);
+  if (!AI)
+    return false;
+
+  class AllocaEscapeChecker : public PtrUseVisitor<AllocaEscapeChecker> {
+    friend class PtrUseVisitor<AllocaEscapeChecker>;
+    friend class InstVisitor<AllocaEscapeChecker>;
+
+  public:
+    AllocaEscapeChecker(AllocaInst &AI)
+        : PtrUseVisitor<AllocaEscapeChecker>(AI.getModule()->getDataLayout()),
+          PtrI(visitPtr(AI)) {}
+
+    bool isEscapedOrAborted() const {
+      return PtrI.isEscaped() || PtrI.isAborted();
+    }
+
+  private:
+    void visitPHINode(PHINode &PN) { enqueueUsers(PN); }
+    void visitSelectInst(SelectInst &SI) { enqueueUsers(SI); }
+
+    PtrInfo PtrI;
+  };
+
+  AllocaEscapeChecker Checker(*AI);
+  return Checker.isEscapedOrAborted();
+}
+
 // Splits the block at a particular instruction unless it is the first
 // instruction in the block with a single predecessor.
 static BasicBlock *splitBlockIfNotFirst(Instruction *I, const Twine &Name) {
   auto *BB = I->getParent();
   if (&BB->front() == I) {
     if (BB->getSinglePredecessor()) {
       BB->setName(Name);
       return BB;
for (CoroSuspendInst *CSI : Shape.CoroSuspends) {
     splitAround(CSI->getCoroSave(), "CoroSave");
     splitAround(CSI, "CoroSuspend");
   }
 
   // Put CoroEnds into their own blocks.
   for (CoroEndInst *CE : Shape.CoroEnds)
     splitAround(CE, "CoroEnd");
 
+  // Split the block after CoroBegin so that spill reload code will never try
+  // to insert a reload before CoroBegin itself.
+  Shape.CoroBegin->getParent()->splitBasicBlock(Shape.CoroBegin->getNextNode(),
+                                                "AfterCoroBegin");
+
   // Transforms multi-edge PHI Nodes, so that any value feeding into a PHI will
   // never has its definition separated from the PHI by the suspend point.
   rewritePHIs(F);
 
   // Build suspend crossing info.
   SuspendCrossingInfo Checker(F, Shape);
 
   IRBuilder<> Builder(F.getContext());
for (Instruction &I : instructions(F)) {
     if (isCoroutineStructureIntrinsic(I) || &I == Shape.CoroBegin)
       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)) {
+      if (Checker.isDefinitionAcrossSuspend(I, U) || allocaEscapes(&I)) {
         // We cannot spill a token.
         if (I.getType()->isTokenTy())
           report_fatal_error(
               "token definition is separated from the use by a suspend point");
         Spills.emplace_back(&I, U);
       }
   }
   DEBUG(dump("Spills", Spills));
   moveSpillUsesAfterCoroBegin(F, Spills, Shape.CoroBegin);
   Shape.FrameTy = buildFrameType(F, Shape, Spills);
   Shape.FramePtr = insertSpills(Spills, Shape);
 }

test/Transforms/Coroutines/coro-spill-escape.ll

+; Check that if alloca escapes it is always put into coroutine frame.
+; Even if it is not accessed beyound suspend points.
+; RUN: opt < %s -coro-split -S | FileCheck %s
+
+define i8* @escaped.case() "coroutine.presplit"="1" {
+entry:
+  %sneaky = alloca i32
+  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call i8* @malloc(i32 %size)
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
+
+  %hideit = ptrtoint i32* %sneaky to i64
+  call void @escape(i64 %hideit)
+
+  %tok = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %tok, label %suspend [i8 0, label %resume
+                                i8 1, label %cleanup]
+resume:
+  br label %cleanup
+
+cleanup:
+  %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
+  call void @free(i8* %mem)
+  br label %suspend
+suspend:
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+  ret i8* %hdl
+}
+
+; Verify that we spilled alloca due to escape.
+; CHECK: %escaped.case.Frame = type { void (%escaped.case.Frame*)*, void (%escaped.case.Frame*)*, i1, i1, i32 }
+
+; CHECK-LABEL: @escaped.case(
+; CHECK:   %[[sneaky:.+]] = getelementptr inbounds %escaped.case.Frame, %escaped.case.Frame* %FramePtr, i32 0, i32 4
+; CHECK:   %[[hideit:.+]]  = ptrtoint i32* %[[sneaky]] to i64
+; CHECK:   call void @escape(i64 %[[hideit]])
+; CHECK:   ret i8* %hdl
+
+declare i8* @llvm.coro.free(token, i8*)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+
+declare token @llvm.coro.id(i32, i8*, i8*, i8*)
+declare i1 @llvm.coro.alloc(token)
+declare i8* @llvm.coro.begin(token, i8*)
+declare i1 @llvm.coro.end(i8*, i1)
+
+declare noalias i8* @malloc(i32)
+declare void @escape(i64)
+declare void @free(i8*)