[Coroutine] Move all used local allocas to the .resume function

[lxfind]

Summary

Prior to D89768, any alloca that's used after suspension points will be put on to the coroutine frame, and hence they will always be reloaded in the resume function.
However D89768 introduced a more precise way to determine whether an alloca should live on the frame. Allocas that are only used within one suspension region (hence does not need to live across suspension points) will not be put on the frame. They will remain local to the resume function.
When creating the new entry for the .resume function, the existing logic only moved all the allocas from the old entry to the new entry. This covers every alloca from the old entry. However allocas that's defined afer coro.begin are put into a separate basic block during CoroSplit (the PostSpill basic block). We need to make sure these allocas are moved to the new entry as well if they are used.
This patch walks through all allocas, and check if they are still used but are not reachable from the new entry, if so, we move them to the new entry.
This should fix the bug reported in D89768

[lxfind]

update description

[lxfind]

update description

[lxfind]

rebase

[ben-clayton]

I can confirm that this fixes the crash we've been seeing - thank you!

[lxfind]

In D90977#2383607, @ben-clayton wrote:

I can confirm that this fixes the crash we've been seeing - thank you!

Cool. I also noticed that your code have quite lots of allocas, which makes it a great example to benchmark the effect of memory savings from D89768.
I was wondering, comparing to before D89768 and now, how much reduction do you see in the coroutine frame size in your codebase (i.e. the eventual argument to @coroutine_alloc_frame)? I would really appreciate if you could share some insights on the magnitude of that.

[junparser]

LGTM, Thanks!

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

void CoroCloner::replaceEntryBlock() {
   // AllocaSpillBlock to begin with.
   assert(Entry->hasOneUse());
   auto BranchToEntry = cast<BranchInst>(Entry->user_back());
   assert(BranchToEntry->isUnconditional());
   Builder.SetInsertPoint(BranchToEntry);
   Builder.CreateUnreachable();
   BranchToEntry->eraseFromParent();
 
-  // 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.use_empty())
-      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 =
       cast<BasicBlock>(VMap[Shape.SwitchLowering.ResumeEntryBlock]);
assert((Shape.ABI == coro::ABI::Async &&
             isa<CoroSuspendRetconInst>(ActiveSuspend)));
     auto *MappedCS = cast<AnyCoroSuspendInst>(VMap[ActiveSuspend]);
     auto Branch = cast<BranchInst>(MappedCS->getNextNode());
     assert(Branch->isUnconditional());
     Builder.CreateBr(Branch->getSuccessor(0));
     break;
   }
   }
+
+  // Any alloca that's still being used but not reachable from the new entry
+  // needs to be moved to the new entry.
+  Function *F = OldEntry->getParent();
+  DominatorTree DT{*F};
+  for (auto IT = inst_begin(F), End = inst_end(F); IT != End;) {
+    Instruction &I = *IT++;
+    if (!isa<AllocaInst>(&I) || I.use_empty())
+      continue;
+    if (DT.isReachableFromEntry(I.getParent()))
+      continue;
+    I.moveBefore(*Entry, Entry->getFirstInsertionPt());
+  }
 }
 
 /// Derive the value of the new frame pointer.
 Value *CoroCloner::deriveNewFramePointer() {
   // Builder should be inserting to the front of the new entry block.
 
   switch (Shape.ABI) {
   // In switch-lowering, the argument is the frame pointer.

llvm/test/Transforms/Coroutines/coro-alloca-05.ll

+; Tests that allocas after coro.begin are properly that do not need to
+; live on the frame are properly moved to the .resume function.
+; RUN: opt < %s -coro-split -S | FileCheck %s
+; RUN: opt < %s -passes=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)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call i8* @malloc(i32 %size)
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
+  %x = alloca i32
+  %sp1 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %sp1, label %suspend [i8 0, label %resume
+  i8 1, label %cleanup]
+resume:
+  %x.value = load i32, i32* %x
+  call void @print(i32 %x.value)
+  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
+}
+
+; CHECK-LABEL: @f.resume(
+; CHECK-NEXT:  entry.resume:
+; CHECK-NEXT:    [[VFRAME:%.*]] = bitcast %f.Frame* [[FRAMEPTR:%.*]] to i8*
+; CHECK-NEXT:    [[X:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[X_VALUE:%.*]] = load i32, i32* [[X]], align 4
+; CHECK-NEXT:    call void @print(i32 [[X_VALUE]])
+; CHECK-NEXT:    call void @free(i8* [[VFRAME]])
+; CHECK-NEXT:    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 i1 @llvm.coro.alloc(token)
+declare i8* @llvm.coro.begin(token, i8*)
+declare i1 @llvm.coro.end(i8*, i1)
+
+declare void @print(i32)
+declare noalias i8* @malloc(i32)
+declare void @free(i8*)