[Coroutine] Properly use lifetime intrinsics to analyze allocas

[lxfind]

Summary

In this patch, we traverse from each lifetime.start intrinsics, follow the control flow. If we never reach a coro.suspend before reaching a lifetime.end intrinsic for the same alloca, we can safely assume that the alloca
will not live across suspension points. This should still achieve similar effect as before but is precise and correct.
Added a test case.

[lxfind]

fix clang-tidy warnings

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

for (BasicBlock *DomBB : DomSet) {
         }
 
         break;
       }
     }
   }
 }
 
+// Check whether the lifetime of an alloca might cross a coroutine suspension
+// point.
+static bool checkAllocaLifetimeMayCrossingSuspend(const AllocaInst &AI) {
+  llvm::SmallVector<const Instruction *, 16> WorkList;
+  SmallPtrSet<const Instruction *, 2> LifetimeEnds;
+  // Collect all lifetime.start and lifetime.end intrinsics for AI.
+  // All lifetime.start will be put in WorkList to traverse.
+  for (const auto &I : instructions(AI.getFunction())) {
+    if (!I.isLifetimeStartOrEnd())
+      continue;
+    auto *II = cast<IntrinsicInst>(&I);
+    if (const auto *OpInst = dyn_cast<Instruction>(I.getOperand(1)))
+      if (const auto *AII = dyn_cast<AllocaInst>(OpInst->stripPointerCasts()))
+        if (AII == &AI) {
+          if (II->getIntrinsicID() == Intrinsic::lifetime_start)
+            WorkList.push_back(II);
+          else
+            LifetimeEnds.insert(&I);
+        }
+  }
+  // If there is no lifetime information, we assume conservatively that the
+  // lifetime of alloca might cross suspension points.
+  if (WorkList.empty())
+    return true;
+
+  // From each lifetime.start, if control flow ever reaches a suspension point
+  // before seen a lifetime.end, we can conclude the alloca's lifetime crosses
+  // coroutine suspensions.
+  llvm::SmallPtrSet<const Instruction *, 16> Visited;
+  while (!WorkList.empty()) {
+    const auto *I = WorkList.back();
+    WorkList.pop_back();
+    if (!Visited.insert(I).second)
+      continue;
+
+    while (!I->isTerminator()) {
+      if (LifetimeEnds.count(I))
+        break;
+      if (isa<AnyCoroSuspendInst>(I))
+        return true;
+      I = I->getNextNode();
+    }
+
+    if (I->isTerminator())
+      for (unsigned Index = 0; Index < I->getNumSuccessors(); ++Index)
+        WorkList.push_back(&I->getSuccessor(Index)->front());
+  }
+
+  return false;
+}
+
 static void collectFrameAllocas(Function &F, coro::Shape &Shape,
                                 const SuspendCrossingInfo &Checker,
                                 SmallVectorImpl<AllocaInfo> &Allocas) {
   for (Instruction &I : instructions(F)) {
     auto *AI = dyn_cast<AllocaInst>(&I);
     if (!AI)
       continue;
     // The PromiseAlloca will be specially handled since it needs to be in a
     // fixed position in the frame.
-    if (AI == Shape.SwitchLowering.PromiseAlloca) {
+    if (AI == Shape.SwitchLowering.PromiseAlloca)
       continue;
-    }
+
+    // If lifetime information guarantees the alloca will not live across
+    // suspensions, it's safe to put it on the stack.
+    if (!checkAllocaLifetimeMayCrossingSuspend(*AI))
+      continue;
+
     DominatorTree DT(F);
     AllocaUseVisitor Visitor{F.getParent()->getDataLayout(), DT,
                              *Shape.CoroBegin, Checker};
     Visitor.visitPtr(*AI);
     if (!Visitor.getShouldLiveOnFrame())
       continue;
     Allocas.emplace_back(AI, Visitor.getAliasesCopy(),
                          Visitor.getMayWriteBeforeCoroBegin());

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

-; Corresponding to coro-split-sink-lifetime-01.ll. This file tests that whether the CoroFrame
-; pass knows the operand of lifetime.start intrinsic may be GEP as well.
 ; 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)
 
 %i8.array = type { [100 x i8] }
-declare void @consume.i8.array(%i8.array* nocapture)
+declare void @consume.i8.array(%i8.array*)
 
-define void @a.gep() "coroutine.presplit"="1" {
+; The lifetime of testval starts and ends before coro.suspend. Even though consume.i8.array
+; might capture it, we can safely say it won't live across suspension.
+define void @foo() "coroutine.presplit"="1" {
 entry:
-  %ref.tmp7 = alloca %"struct.lean_future<int>::Awaiter", align 8
   %testval = alloca %i8.array
   %cast = getelementptr inbounds %i8.array, %i8.array* %testval, i64 0, i32 0, i64 0
-  ; lifetime of %testval starts here, but not used until await.ready.
-  call void @llvm.lifetime.start.p0i8(i64 100, i8* %cast)
   %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)
 
+  call void @llvm.lifetime.start.p0i8(i64 100, i8* %cast)
+  call void @consume.i8.array(%i8.array* %testval)
+  call void @llvm.lifetime.end.p0i8(i64 100, i8*  %cast)
+
   %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
+  br label %exit
+exit:
+  call i1 @llvm.coro.end(i8* null, i1 false)
+  ret void
+}
+
+; The lifetime of testval starts after coro.suspend. So it will never live across suspension
+; points.
+define void @bar() "coroutine.presplit"="1" {
+entry:
+  %testval = alloca %i8.array
+  %cast = getelementptr inbounds %i8.array, %i8.array* %testval, i64 0, i32 0, i64 0
+  %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)
+  %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)
+
+  call void @llvm.lifetime.start.p0i8(i64 100, i8* %cast)
   call void @consume.i8.array(%i8.array* %testval)
   call void @llvm.lifetime.end.p0i8(i64 100, 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.gep.resume(
-; CHECK:         %testval = alloca %i8.array
-; CHECK-NEXT:    getelementptr inbounds %a.gep.Frame
-; CHECK-NEXT:    %0 = bitcast %i8.array* %testval to i8*
-; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 100, i8* %0)
-; CHECK-NEXT:    getelementptr inbounds %"struct.lean_future<int>::Awaiter"
-; CHECK-NEXT:    getelementptr inbounds %i8.array, %i8.array* %testval
-; CHECK-NEXT:    %val = load i32, i32* %Result
-; CHECK-NEXT:    call void @consume.i8.array(%i8.array* %testval)
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 100, i8* %cast1)
-; CHECK-NEXT:    call void @print(i32 %val)
-; CHECK-NEXT:    ret void
+
+; Verify that for both foo and bar, testval isn't put on the frame.
+; CHECK: %foo.Frame = type { void (%foo.Frame*)*, void (%foo.Frame*)*, i1 }
+; CHECK: %bar.Frame = type { void (%bar.Frame*)*, void (%bar.Frame*)*, i1 }
 
 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

llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-03.ll

 ; Corresponding to coro-split-sink-lifetime-01.ll. This file tests that whether the CoroFrame
 ; pass knows the operand of lifetime.start intrinsic may be GEP as well.
 ; 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)
 
 %i8.array = type { [100 x i8] }
-declare void @consume.i8.array(%i8.array* nocapture)
+declare void @consume.i8.array(%i8.array*)
 
 define void @a.gep() "coroutine.presplit"="1" {
 entry:
   %ref.tmp7 = alloca %"struct.lean_future<int>::Awaiter", align 8
   %testval = alloca %i8.array
   %cast = getelementptr inbounds %i8.array, %i8.array* %testval, i64 0, i32 0, i64 0
   ; lifetime of %testval starts here, but not used until await.ready.
   call void @llvm.lifetime.start.p0i8(i64 100, i8* %cast)