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

Alan Zhao via llvm-commits llvm-commits at lists.llvm.org
Mon Apr 15 15:25:22 PDT 2024


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

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

>From d99badcf1e2a2f160dba9c9e215e3d21993a4b54 Mon Sep 17 00:00:00 2001
From: Alan Zhao <ayzhao at google.com>
Date: Mon, 15 Apr 2024 14:57:22 -0700
Subject: [PATCH 1/2] [nfc][coroutines] Precommit a test for allocas without
 lifetime end annotations

---
 .../Coroutines/coro-alloca-no-lifetime-end.ll | 55 +++++++++++++++++++
 1 file changed, 55 insertions(+)
 create mode 100644 llvm/test/Transforms/Coroutines/coro-alloca-no-lifetime-end.ll

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..db3c2491418208
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-alloca-no-lifetime-end.ll
@@ -0,0 +1,55 @@
+; 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:    [[TESTVAL:%.*]] = alloca [[I8_ARRAY:%.*]], align 8
+; 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:    call void @llvm.lifetime.start.p0(i64 100, ptr [[TESTVAL]])
+; CHECK-NEXT:    call void @consume.i8.array(ptr [[TESTVAL]])
+; CHECK-NEXT:    [[INDEX_ADDR1:%.*]] = getelementptr inbounds [[FOO_FRAME]], ptr [[VFRAME]], i32 0, i32 2
+; CHECK-NEXT:    store i1 false, ptr [[INDEX_ADDR1]], 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

>From f2dc10b98d8595174007ed1b47256bd629b0b0a5 Mon Sep 17 00:00:00 2001
From: Alan Zhao <ayzhao at google.com>
Date: Mon, 15 Apr 2024 15:11:28 -0700
Subject: [PATCH 2/2] [coroutines][CoroSplit] Store allocas on the frame if
 there is no explicit lifetime.end

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 #86580
---
 llvm/lib/Transforms/Coroutines/CoroFrame.cpp               | 7 +++++++
 .../Transforms/Coroutines/coro-alloca-no-lifetime-end.ll   | 7 +++----
 2 files changed, 10 insertions(+), 4 deletions(-)

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
index db3c2491418208..c9fdd4fd3d2bab 100644
--- a/llvm/test/Transforms/Coroutines/coro-alloca-no-lifetime-end.ll
+++ b/llvm/test/Transforms/Coroutines/coro-alloca-no-lifetime-end.ll
@@ -11,17 +11,16 @@ declare void @consume.i8.array(ptr)
 define void @foo() presplitcoroutine {
 ; CHECK-LABEL: define void @foo() {
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TESTVAL:%.*]] = alloca [[I8_ARRAY:%.*]], align 8
 ; 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:    call void @llvm.lifetime.start.p0(i64 100, ptr [[TESTVAL]])
-; CHECK-NEXT:    call void @consume.i8.array(ptr [[TESTVAL]])
 ; CHECK-NEXT:    [[INDEX_ADDR1:%.*]] = getelementptr inbounds [[FOO_FRAME]], ptr [[VFRAME]], i32 0, i32 2
-; CHECK-NEXT:    store i1 false, ptr [[INDEX_ADDR1]], align 1
+; 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:



More information about the llvm-commits mailing list