[Coroutines] Also check lifetime intrinsic for local variable when build coroutine frame

[junparser]

Summary

Currently we move all allocas into the frame when build coroutine frame in CoroSplit pass. However, this can be relaxed.

Since CoroSplit pass run after Inline pass, we can use lifetime intrinsic to do such analysis: If the scope of lifetime intrinsic is
not across any suspend point, rather than move the allocas to frame, we can just move them to entry bb of corresponding function. This reduce the frame size.

More importantly, this also avoid data race in multithread environment. Consider one inline function by coroutine: it starts a thread which access local variables, while after inline the movement of allocs to frame also access them. cause data race.

TestPlan: check-llvm, cppcoro

[wenlei]

Nice change, thanks @junparser! The improved dataflow analysis using lifetime instead of allocation point and the resulting shrink-wrapping style opt for alloca would definitely be beneficial.

For data race however, I'm not sure if I understand it correctly, could you elaborate (perhaps with a concrete example..) why this is compiler's fault?

[junparser]

In D75664#1907758, @wenlei wrote:

Nice change, thanks @junparser! The improved dataflow analysis using lifetime instead of allocation point and the resulting shrink-wrapping style opt for alloca would definitely be beneficial.

For data race however, I'm not sure if I understand it correctly, could you elaborate (perhaps with a concrete example..) why this is compiler's fault?

For

auto await_suspend(std::experimental::coroutine_handle<> h) {
    auto& pr = h.promise();
    if (pr._ex) {
        // schedule to another thread
        pr._ex->schedule([&pr]() {
            pr._continuation.resume();
        });
    }
}

This case shows data race with inline, and does not when not inlined. It looks like both the current thread and schedule thread access some fields of std::function when I dig into the ir. I'm not sure whether this is caused by std::function constructor/destructor or move allocas to frame. I get such conclusion based on ir.

[junparser]

In D75664#1908825, @junparser wrote:

In D75664#1907758, @wenlei wrote:

Nice change, thanks @junparser! The improved dataflow analysis using lifetime instead of allocation point and the resulting shrink-wrapping style opt for alloca would definitely be beneficial.

For data race however, I'm not sure if I understand it correctly, could you elaborate (perhaps with a concrete example..) why this is compiler's fault?

For

auto await_suspend(std::experimental::coroutine_handle<> h) {
    auto& pr = h.promise();
    if (pr._ex) {
        // schedule to another thread
        pr._ex->schedule([&pr]() {
            pr._continuation.resume();
        });
    }
}

This case shows data race with inline, and does not when not inlined. It looks like both the current thread and schedule thread access some fields of std::function when I dig into the ir. I'm not sure whether this is caused by std::function constructor/destructor or move allocas to frame. I get such conclusion based on ir.

More specifically, the tsan gives 'use after free' error. So I think tsan can trace such access of frame block, but it may not be able to trace local allocas.

[junparser]

It seems that this patch also fix some workaround in folly, https://github.com/facebook/folly/blob/master/folly/Portability.h#L501, This is the same behavior as well as our coroutine library

[junparser]

In D75664#1912080, @junparser wrote:

It seems that this patch also fix some workaround in folly, https://github.com/facebook/folly/blob/master/folly/Portability.h#L501, This is the same behavior as well as our coroutine library

@wenlei， according to @lewissbaker 's comment on await object:

So within the await_suspend() method, once it’s possible for the coroutine to be resumed concurrently on another thread, you need to make sure that you avoid accessing this or the coroutine’s .promise() object because both could already be destroyed. In general, the only things that are safe to access after the operation is started and the coroutine is scheduled for resumption are local variables within await_suspend().

when local variable becomes frame variable, data race may happens.

[wenlei]

Thanks for clarification, LGTM!

[junparser]

@wenlei @modocache would you please review the updated patch?
I make the update for :

1. remove assertion since we use lifetime intrinsic as def
2. support multiple lifetime intrinsic
3. it invalidates one case in cppcoro (large number of synchronous completions doesn't result in stack-overflow) since we do not move local variables into frame now @lewissbaker would please udpate it?

[modocache]

Nice, thanks for this improvement! Am I right in thinking patch is dependent upon D76119? My understanding is that, without that patch, then we wouldn't be able to rely on lifetime intrinsics when using -O0. You may want to use Phabricator's "Edit Related Objects" link to specify that this patch depends on that one.

This LGTM assuming we can rely on lifetime intrinsics being generated at -O0. I'll comment on D76119 on that point.

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1470

Nit-pick: I think the LLVM coding style would have you remove these for loop braces. There's only a single statement in the for loop, the if statement above, so you technically don't need these here.

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

572

Wow, I'm surprised coro-split-02.ll was the only test case you ended up needing to modify. I would have thought that this one change would have broken a lot more tests' assertions.

584

Awesome! Nice to remove this TODO.

[junparser]

@modocache Thanks !

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1470

Thanks, I'll update this one.

[junparser]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1470

we should keep the braces, since there is another else below

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

struct SuspendCrossingInfo {
   void dump(StringRef Label, BitVector const &BV) const;
 
   SuspendCrossingInfo(Function &F, coro::Shape &Shape);
 
   bool hasPathCrossingSuspendPoint(BasicBlock *DefBB, BasicBlock *UseBB) const {
     size_t const DefIndex = Mapping.blockToIndex(DefBB);
     size_t const UseIndex = Mapping.blockToIndex(UseBB);
 
-    assert(Block[UseIndex].Consumes[DefIndex] && "use must consume def");
     bool const Result = Block[UseIndex].Kills[DefIndex];
     LLVM_DEBUG(dbgs() << UseBB->getName() << " => " << DefBB->getName()
                       << " answer is " << Result << "\n");
     return Result;
   }
 
   bool isDefinitionAcrossSuspend(BasicBlock *DefBB, User *U) const {
     auto *I = cast<Instruction>(U);
if (Spills.empty())
       break;
 
     // Rewrite materializable instructions to be materialized at the use point.
     LLVM_DEBUG(dump("Materializations", Spills));
     rewriteMaterializableInstructions(Builder, Spills);
     Spills.clear();
   }
 
+  // Collect lifetime.start info for each alloca.
+  using LifetimeStart = SmallPtrSet<Instruction *, 2>;
+  llvm::DenseMap<Instruction *, std::unique_ptr<LifetimeStart>> LifetimeMap;
+  for (Instruction &I : instructions(F)) {
+    auto *II = dyn_cast<IntrinsicInst>(&I);
+    if (!II || II->getIntrinsicID() != Intrinsic::lifetime_start)
+      continue;
+
+    if (auto *OpInst = dyn_cast<BitCastInst>(I.getOperand(1)))
+      if (auto *AI = dyn_cast<AllocaInst>(OpInst->getOperand(0))) {
+
+        if (LifetimeMap.find(AI) == LifetimeMap.end())
+          LifetimeMap[AI] = std::make_unique<LifetimeStart>();
+
+        LifetimeMap[AI]->insert(OpInst);
+      }
+  }
+
   // Collect the spills for arguments and other not-materializable values.
   for (Argument &A : F.args())
     for (User *U : A.users())
       if (Checker.isDefinitionAcrossSuspend(A, U))
         Spills.emplace_back(&A, U);
 
   for (Instruction &I : instructions(F)) {
     // Values returned from coroutine structure intrinsics should not be part
if (auto AI = dyn_cast<CoroAllocaAllocInst>(&I)) {
       continue;
     }
 
     // Ignore alloca.get; we process this as part of coro.alloca.alloc.
     if (isa<CoroAllocaGetInst>(I)) {
       continue;
     }
 
-    for (User *U : I.users())
-      if (Checker.isDefinitionAcrossSuspend(I, U)) {
+    auto Iter = LifetimeMap.find(&I);
+    for (User *U : I.users()) {
+      bool NeedSpill = false;
+
+      // Check against lifetime.start if the instruction has the info.
+      if (Iter != LifetimeMap.end())
+        for (auto *S : *Iter->second) {
+          if ((NeedSpill = Checker.isDefinitionAcrossSuspend(*S, U)))
+            break;
+        }
+      else
+        NeedSpill = Checker.isDefinitionAcrossSuspend(I, U);
+
+      if (NeedSpill) {
         // 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);
       }
     }
+  }
   LLVM_DEBUG(dump("Spills", Spills));
   Shape.FrameTy = buildFrameType(F, Shape, Spills);
   Shape.FramePtr = insertSpills(Spills, Shape);
   lowerLocalAllocas(LocalAllocas, DeadInstructions);
 
   for (auto I : DeadInstructions)
     I->eraseFromParent();
 }

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

 
 void CoroCloner::replaceEntryBlock() {
   // In the original function, the AllocaSpillBlock is a block immediately
   // following the allocation of the frame object which defines GEPs for
   // all the allocas that have been moved into the frame, and it ends by
   // branching to the original beginning of the coroutine.  Make this 
   // the entry block of the cloned function.
   auto *Entry = cast<BasicBlock>(VMap[Shape.AllocaSpillBlock]);
+  auto *OldEntry = &NewF->getEntryBlock();
   Entry->setName("entry" + Suffix);
-  Entry->moveBefore(&NewF->getEntryBlock());
+  Entry->moveBefore(OldEntry);
   Entry->getTerminator()->eraseFromParent();
 
   // Clear all predecessors of the new entry block.  There should be
   // exactly one predecessor, which we created when splitting out
   // AllocaSpillBlock to begin with.
   assert(Entry->hasOneUse());
   auto BranchToEntry = cast<BranchInst>(Entry->user_back());
   assert(BranchToEntry->isUnconditional());
   Builder.SetInsertPoint(BranchToEntry);
   Builder.CreateUnreachable();
   BranchToEntry->eraseFromParent();
 
-  // TODO: move any allocas into Entry that weren't moved into the frame.
-  // (Currently we move all allocas into the frame.)
+  // Move any allocas into Entry that weren't moved into the frame.
+  for (auto IT = OldEntry->begin(), End = OldEntry->end(); IT != End;) {
+    Instruction &I = *IT++;
+    if (!isa<AllocaInst>(&I) || I.getNumUses() == 0)
+      continue;
+
+    I.moveBefore(*Entry, Entry->getFirstInsertionPt());
+  }
 
   // Branch from the entry to the appropriate place.
   Builder.SetInsertPoint(Entry);
   switch (Shape.ABI) {
   case coro::ABI::Switch: {
     // In switch-lowering, we built a resume-entry block in the original
     // function.  Make the entry block branch to this.
     auto *SwitchBB =

llvm/test/Transforms/Coroutines/coro-split-02.ll

 ; Tests that coro-split can handle the case when a code after coro.suspend uses
 ; a value produces between coro.save and coro.suspend (%Result.i19)
 ; and checks whether stray coro.saves are properly removed
 ; RUN: opt < %s -coro-split -S | FileCheck %s
 ; RUN: opt < %s -passes=coro-split -S | FileCheck %s
 
 %"struct.std::coroutine_handle" = type { i8* }
 %"struct.std::coroutine_handle.0" = type { %"struct.std::coroutine_handle" }
 %"struct.lean_future<int>::Awaiter" = type { i32, %"struct.std::coroutine_handle.0" }
 
 declare i8* @malloc(i64)
 declare void @print(i32)
 
 define void @a() "coroutine.presplit"="1" {
 entry:
   %ref.tmp7 = alloca %"struct.lean_future<int>::Awaiter", align 8
+  %testval = alloca i32
   %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
   %alloc = call i8* @malloc(i64 16) #3
   %vFrame = call noalias nonnull i8* @llvm.coro.begin(token %id, i8* %alloc)
 
   %save = call token @llvm.coro.save(i8* null)
   %Result.i19 = getelementptr inbounds %"struct.lean_future<int>::Awaiter", %"struct.lean_future<int>::Awaiter"* %ref.tmp7, i64 0, i32 0
   %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
   switch i8 %suspend, label %exit [
     i8 0, label %await.ready
     i8 1, label %exit
   ]
 await.ready:
   %StrayCoroSave = call token @llvm.coro.save(i8* null)
   %val = load i32, i32* %Result.i19
+  %cast = bitcast i32* %testval to i8*
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* %cast)
+  call void @llvm.lifetime.end.p0i8(i64 4, i8*  %cast)
   call void @print(i32 %val)
   br label %exit
 exit:
   call i1 @llvm.coro.end(i8* null, i1 false)
   ret void
 }
 
 ; CHECK-LABEL: @a.resume(
+; CHECK:         %testval = alloca i32
 ; CHECK:         getelementptr inbounds %a.Frame
 ; CHECK-NEXT:    getelementptr inbounds %"struct.lean_future<int>::Awaiter"
 ; CHECK-NOT:     call token @llvm.coro.save(i8* null)
 ; CHECK-NEXT:    %val = load i32, i32* %Result
+; CHECK-NEXT:    %cast = bitcast i32* %testval to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 4, i8* %cast)
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 4, i8* %cast)
 ; CHECK-NEXT:    call void @print(i32 %val)
 ; CHECK-NEXT:    ret void
 
 declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*)
 declare i1 @llvm.coro.alloc(token) #3
 declare noalias nonnull i8* @"\01??2@YAPEAX_K@Z"(i64) local_unnamed_addr
 declare i64 @llvm.coro.size.i64() #5
 declare i8* @llvm.coro.begin(token, i8* writeonly) #3
 declare void @"\01?puts@@YAXZZ"(...)
 declare token @llvm.coro.save(i8*) #3
 declare i8* @llvm.coro.frame() #5
 declare i8 @llvm.coro.suspend(token, i1) #3
 declare void @"\01??3@YAXPEAX@Z"(i8*) local_unnamed_addr #10
 declare i8* @llvm.coro.free(token, i8* nocapture readonly) #2
 declare i1 @llvm.coro.end(i8*, i1) #3
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #4
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #4