[llvm] [coroutines][CoroSplit] Store allocas on the frame if there is no exp… …licit lifetime.end (PR #88806)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Alan Zhao (alanzhao1)

<details>
<summary>Changes</summary>

Certain passes such as SimplifyCFG may remove existing lifetime.end
intrinsics on presplit coroutines. CoroSplitPass then incorrectly
assumes that these allocas may not persist across the coroutine
suspension point. To fix this, computeShouldLiveOnFrame() now saves
allocas on the coroutine frame object if their addresses are escaped and
those allocas don't have a corresponding lifetime.end.

Fixes https://github.com/llvm/llvm-project/issues/86580

---
Full diff: https://github.com/llvm/llvm-project/pull/88806.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Coroutines/CoroFrame.cpp (+7) 
- (added) llvm/test/Transforms/Coroutines/coro-alloca-no-lifetime-end.ll (+54) 


``````````diff
diff --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
index 08a4522e3fac6d..8bd7d0dc36d532 100644
--- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
@@ -1550,6 +1550,8 @@ struct AllocaUseVisitor : PtrUseVisitor<AllocaUseVisitor> {
   }
 
   void visitIntrinsicInst(IntrinsicInst &II) {
+    if (II.getIntrinsicID() == Intrinsic::lifetime_end)
+      HasExplicitLifetimeEnd = true;
     // When we found the lifetime markers refers to a
     // subrange of the original alloca, ignore the lifetime
     // markers to avoid misleading the analysis.
@@ -1596,6 +1598,7 @@ struct AllocaUseVisitor : PtrUseVisitor<AllocaUseVisitor> {
   SmallPtrSet<IntrinsicInst *, 2> LifetimeStarts{};
   bool MayWriteBeforeCoroBegin{false};
   bool ShouldUseLifetimeStartInfo{true};
+  bool HasExplicitLifetimeEnd{false};
 
   mutable std::optional<bool> ShouldLiveOnFrame{};
 
@@ -1614,6 +1617,10 @@ struct AllocaUseVisitor : PtrUseVisitor<AllocaUseVisitor> {
       // suspend point between lifetime markers. This should also cover the
       // case of a single lifetime.start intrinsic in a loop with suspend point.
       if (PI.isEscaped()) {
+        // If there is no explicit lifetime.end, then assume the address can
+        // cross suspension points.
+        if (!HasExplicitLifetimeEnd)
+          return true;
         for (auto *A : LifetimeStarts) {
           for (auto *B : LifetimeStarts) {
             if (Checker.hasPathOrLoopCrossingSuspendPoint(A->getParent(),
diff --git a/llvm/test/Transforms/Coroutines/coro-alloca-no-lifetime-end.ll b/llvm/test/Transforms/Coroutines/coro-alloca-no-lifetime-end.ll
new file mode 100644
index 00000000000000..c9fdd4fd3d2bab
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-alloca-no-lifetime-end.ll
@@ -0,0 +1,54 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt < %s -passes='cgscc(coro-split),simplifycfg,early-cse' -S | FileCheck %s
+
+declare ptr @malloc(i64)
+
+%i8.array = type { [100 x i8] }
+declare void @consume.i8.array(ptr)
+
+; testval does not contain an explicit lifetime end. We must assume that it may
+; live across suspension.
+define void @foo() presplitcoroutine {
+; CHECK-LABEL: define void @foo() {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @foo.resumers)
+; CHECK-NEXT:    [[ALLOC:%.*]] = call ptr @malloc(i64 16)
+; CHECK-NEXT:    [[VFRAME:%.*]] = call noalias nonnull ptr @llvm.coro.begin(token [[ID]], ptr [[ALLOC]])
+; CHECK-NEXT:    store ptr @foo.resume, ptr [[VFRAME]], align 8
+; CHECK-NEXT:    [[DESTROY_ADDR:%.*]] = getelementptr inbounds [[FOO_FRAME:%.*]], ptr [[VFRAME]], i32 0, i32 1
+; CHECK-NEXT:    store ptr @foo.destroy, ptr [[DESTROY_ADDR]], align 8
+; CHECK-NEXT:    [[INDEX_ADDR1:%.*]] = getelementptr inbounds [[FOO_FRAME]], ptr [[VFRAME]], i32 0, i32 2
+; CHECK-NEXT:    call void @consume.i8.array(ptr [[INDEX_ADDR1]])
+; CHECK-NEXT:    [[INDEX_ADDR2:%.*]] = getelementptr inbounds [[FOO_FRAME]], ptr [[VFRAME]], i32 0, i32 3
+; CHECK-NEXT:    store i1 false, ptr [[INDEX_ADDR2]], align 1
+; CHECK-NEXT:    ret void
+;
+entry:
+  %testval = alloca %i8.array
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %alloc = call ptr @malloc(i64 16) #3
+  %vFrame = call noalias nonnull ptr @llvm.coro.begin(token %id, ptr %alloc)
+
+  call void @llvm.lifetime.start.p0(i64 100, ptr %testval)
+  call void @consume.i8.array(ptr %testval)
+
+  %save = call token @llvm.coro.save(ptr 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(ptr null)
+  br label %exit
+exit:
+  ret void
+}
+
+declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr)
+declare ptr @llvm.coro.begin(token, ptr writeonly) #3
+declare token @llvm.coro.save(ptr) #3
+declare ptr @llvm.coro.frame() #5
+declare i8 @llvm.coro.suspend(token, i1) #3
+declare i1 @llvm.coro.end(ptr, i1, token) #3
+declare void @llvm.lifetime.start.p0(i64, ptr nocapture) #4

``````````

</details>


https://github.com/llvm/llvm-project/pull/88806