[coro async] Don't use lifetime.start based alloca localization for ABI.Async/ABI.Retcon

[aschwaighofer]

Summary

Infinite loops can lead to an IR representation where the lifetime.end
intrinsice is missing. The code to do lifetime based optimization then
fails to see that an address escapes (is life) accross a supspend.

Eventually, we could detect such situations and disable it under more narrow
circumstances. For now, do the correct thing.

rdar://83635953

[rjmccall]

Same question as for https://reviews.llvm.org/D110953: is there a reason this isn't broken for the switch lowering as well?

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

 // offset is unknown (e.g. when you have a PHINode that takes in different
 // offset values). We cannot handle unknown offsets and will assert. This is the
 // potential issue left out. An ideal solution would likely require a
 // significant redesign.
 namespace {
 struct AllocaUseVisitor : PtrUseVisitor<AllocaUseVisitor> {
   using Base = PtrUseVisitor<AllocaUseVisitor>;
   AllocaUseVisitor(const DataLayout &DL, const DominatorTree &DT,
-                   const CoroBeginInst &CB, const SuspendCrossingInfo &Checker)
-      : PtrUseVisitor(DL), DT(DT), CoroBegin(CB), Checker(Checker) {}
+                   const CoroBeginInst &CB, const SuspendCrossingInfo &Checker,
+                   bool ShouldUseLifetimeStartInfo)
+      : PtrUseVisitor(DL), DT(DT), CoroBegin(CB), Checker(Checker),
+        ShouldUseLifetimeStartInfo(ShouldUseLifetimeStartInfo) {}
 
   void visit(Instruction &I) {
     Users.insert(&I);
     Base::visit(I);
     // If the pointer is escaped prior to CoroBegin, we have to assume it would
     // be written into before CoroBegin as well.
     if (PI.isEscaped() && !DT.dominates(&CoroBegin, PI.getEscapingInst())) {
       MayWriteBeforeCoroBegin = true;
private:
   const SuspendCrossingInfo &Checker;
   // All alias to the original AllocaInst, created before CoroBegin and used
   // after CoroBegin. Each entry contains the instruction and the offset in the
   // original Alloca. They need to be recreated after CoroBegin off the frame.
   DenseMap<Instruction *, llvm::Optional<APInt>> AliasOffetMap{};
   SmallPtrSet<Instruction *, 4> Users{};
   SmallPtrSet<IntrinsicInst *, 2> LifetimeStarts{};
   bool MayWriteBeforeCoroBegin{false};
+  bool ShouldUseLifetimeStartInfo{true};
 
   mutable llvm::Optional<bool> ShouldLiveOnFrame{};
 
   bool computeShouldLiveOnFrame() const {
     // If lifetime information is available, we check it first since it's
     // more precise. We look at every pair of lifetime.start intrinsic and
     // every basic block that uses the pointer to see if they cross suspension
     // points. The uses cover both direct uses as well as indirect uses.
-    if (!LifetimeStarts.empty()) {
+    if (ShouldUseLifetimeStartInfo && !LifetimeStarts.empty()) {
       for (auto *I : Users)
         for (auto *S : LifetimeStarts)
           if (Checker.isDefinitionAcrossSuspend(*S, I))
             return true;
       return false;
     }
     // FIXME: Ideally the isEscaped check should come at the beginning.
     // However there are a few loose ends that need to be fixed first before
for (Instruction &I : instructions(F)) {
     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) {
       continue;
     }
     DominatorTree DT(F);
+    // The code that uses lifetime.start intrinsic does not work for functions
+    // with loops without exit. Disable it on ABIs we know to generate such
+    // code.
+    bool ShouldUseLifetimeStartInfo =
+        (Shape.ABI != coro::ABI::Async && Shape.ABI != coro::ABI::Retcon &&
+         Shape.ABI != coro::ABI::RetconOnce);
     AllocaUseVisitor Visitor{F.getParent()->getDataLayout(), DT,
-                             *Shape.CoroBegin, Checker};
+                             *Shape.CoroBegin, Checker,
+                             ShouldUseLifetimeStartInfo};
     Visitor.visitPtr(*AI);
     if (!Visitor.getShouldLiveOnFrame())
       continue;
     Allocas.emplace_back(AI, Visitor.getAliasesCopy(),
                          Visitor.getMayWriteBeforeCoroBegin());
   }
 }
 

llvm/test/Transforms/Coroutines/coro-async-addr-lifetime-infinite-loop-bug.ll

+; RUN: opt < %s -enable-coroutines -passes='default<O0>' -S | FileCheck --check-prefixes=CHECK %s
+; RUN: opt < %s -enable-coroutines -O0 -S | FileCheck --check-prefixes=CHECK %s
+
+target datalayout = "p:64:64:64"
+
+%async.task = type { i64 }
+%async.actor = type { i64 }
+%async.fp = type <{ i32, i32 }>
+
+%async.ctxt = type { i8*, void (i8*, %async.task*, %async.actor*)* }
+
+; The async callee.
+@my_other_async_function_fp = external global <{ i32, i32 }>
+declare void @my_other_async_function(i8* %async.ctxt)
+
+@my_async_function_fp = constant <{ i32, i32 }>
+  <{ i32 trunc ( ; Relative pointer to async function
+       i64 sub (
+         i64 ptrtoint (void (i8*)* @my_async_function to i64),
+         i64 ptrtoint (i32* getelementptr inbounds (<{ i32, i32 }>, <{ i32, i32 }>* @my_async_function_fp, i32 0, i32 1) to i64)
+       )
+     to i32),
+     i32 128    ; Initial async context size without space for frame
+}>
+
+define swiftcc void @my_other_async_function_fp.apply(i8* %fnPtr, i8* %async.ctxt) {
+  %callee = bitcast i8* %fnPtr to void(i8*)*
+  tail call swiftcc void %callee(i8* %async.ctxt)
+  ret void
+}
+
+declare void @escape(i64*)
+declare void @store_resume(i8*)
+define i8* @resume_context_projection(i8* %ctxt) {
+entry:
+  %resume_ctxt_addr = bitcast i8* %ctxt to i8**
+  %resume_ctxt = load i8*, i8** %resume_ctxt_addr, align 8
+  ret i8* %resume_ctxt
+}
+
+; The address of alloca escapes but the analysis based on lifetimes fails to see
+; that it can't localize this alloca.
+; CHECK: define swiftcc void @my_async_function(i8* swiftasync %async.ctxt) {
+; CHECK: entry:
+; CHECK-NOT: ret
+; CHECK-NOT:   [[ESCAPED_ADDR:%.*]] = alloca i64, align 8
+; CHECK: ret
+define swiftcc void @my_async_function(i8* swiftasync %async.ctxt) {
+entry:
+  %escaped_addr = alloca i64
+
+  %id = call token @llvm.coro.id.async(i32 128, i32 16, i32 0,
+          i8* bitcast (<{i32, i32}>* @my_async_function_fp to i8*))
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* null)
+  %ltb = bitcast i64* %escaped_addr to i8*
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* %ltb)
+  call void @escape(i64* %escaped_addr)
+  br label %callblock
+
+
+callblock:
+
+  %callee_context = call i8* @context_alloc()
+
+  %resume.func_ptr = call i8* @llvm.coro.async.resume()
+  call void @store_resume(i8* %resume.func_ptr)
+  %resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
+  %callee = bitcast void(i8*)* @asyncSuspend to i8*
+  %res = call {i8*, i8*, i8*} (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 0,
+                                                  i8* %resume.func_ptr,
+                                                  i8* %resume_proj_fun,
+                                                  void (i8*, i8*)* @my_other_async_function_fp.apply,
+                                                  i8* %callee, i8* %callee_context)
+  br label %callblock
+}
+
+declare { i8*, i8*, i8*, i8* } @llvm.coro.suspend.async.sl_p0i8p0i8p0i8p0i8s(i32, i8*, i8*, ...)
+declare i8* @llvm.coro.prepare.async(i8*)
+declare token @llvm.coro.id.async(i32, i32, i32, i8*)
+declare i8* @llvm.coro.begin(token, i8*)
+declare i1 @llvm.coro.end.async(i8*, i1, ...)
+declare i1 @llvm.coro.end(i8*, i1)
+declare {i8*, i8*, i8*} @llvm.coro.suspend.async(i32, i8*, i8*, ...)
+declare i8* @context_alloc()
+declare void @llvm.coro.async.context.dealloc(i8*)
+declare swiftcc void @asyncSuspend(i8*)
+declare i8* @llvm.coro.async.resume()
+declare void @llvm.coro.async.size.replace(i8*, i8*)
+declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture) #0
+declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture) #0
+attributes #0 = { argmemonly nofree nosync nounwind willreturn }