[llvm] ddbf196 - [Coroutines] Ignore partial lifetime markers refer of an alloca

[Chuanqi Xu via llvm-commits]

Author: Chuanqi Xu
Date: 2021-10-22T09:49:50+08:00
New Revision: ddbf1961941b86a61e824fb6bab1ab11f8b3ea90

URL: https://github.com/llvm/llvm-project/commit/ddbf1961941b86a61e824fb6bab1ab11f8b3ea90
DIFF: https://github.com/llvm/llvm-project/commit/ddbf1961941b86a61e824fb6bab1ab11f8b3ea90.diff

LOG: [Coroutines] Ignore partial lifetime markers refer of an alloca

When I playing with Coroutines, I found that it is possible to generate
following IR:
```
%struct = alloca ...
%sub.element = getelementptr %struct, i64 0, i64 index ; index is not
%zero
lifetime.marker.start(%sub.element)
% use of %sub.element
lifetime.marker.end(%sub.element)
store %struct to xxx ;  %struct is escaping!

<suspend points>
```

Then the AllocaUseVisitor would collect the lifetime marker for
sub.element and treat it as the lifetime markers of the alloca! So it
judges that the alloca could be put on the stack instead of the frame by
judging the lifetime markers only.
The root cause for the bug is that AllocaUseVisitor collects wrong
lifetime markers.

This patch fixes this.

Reviewed By: lxfind

Differential Revision: https://reviews.llvm.org/D112216

Added: 
    llvm/test/Transforms/Coroutines/coro-alloca-09.ll

Modified: 
    llvm/lib/Transforms/Coroutines/CoroFrame.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
index 47b39fca3582c..21e78bda73d8d 100644
--- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
@@ -1349,7 +1349,11 @@ struct AllocaUseVisitor : PtrUseVisitor<AllocaUseVisitor> {
   }
 
   void visitIntrinsicInst(IntrinsicInst &II) {
-    if (II.getIntrinsicID() != Intrinsic::lifetime_start)
+    // When we found the lifetime markers refers to a
+    // subrange of the original alloca, ignore the lifetime
+    // markers to avoid misleading the analysis.
+    if (II.getIntrinsicID() != Intrinsic::lifetime_start || !IsOffsetKnown ||
+        !Offset.isZero())
       return Base::visitIntrinsicInst(II);
     LifetimeStarts.insert(&II);
   }

diff  --git a/llvm/test/Transforms/Coroutines/coro-alloca-09.ll b/llvm/test/Transforms/Coroutines/coro-alloca-09.ll
new file mode 100644
index 0000000000000..bd848eb733e7d
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-alloca-09.ll
@@ -0,0 +1,57 @@
+; RUN: opt < %s -passes='cgscc(coro-split),simplifycfg,early-cse' -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)
+
+%i8.array = type { [100 x i8] }
+declare void @consume.i8(i8*)
+
+; The testval lives across suspend point so that it should be put on the frame.
+; However, part of testval has lifetime marker which indicates the part
+; wouldn't live across suspend point.
+; This test whether or not %testval would be put on the frame by ignoring the
+; partial lifetime markers.
+define void @foo(%i8.array** %to_store) "coroutine.presplit"="1" {
+entry:
+  %testval = alloca %i8.array
+  %subrange = getelementptr inbounds %i8.array, %i8.array* %testval, i64 0, i32 0, i64 50
+  %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 50, i8* %subrange)
+  call void @consume.i8(i8* %subrange)
+  call void @llvm.lifetime.end.p0i8(i64 50, i8*  %subrange)
+  store %i8.array* %testval, %i8.array** %to_store
+
+  %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)
+  br label %exit
+exit:
+  call i1 @llvm.coro.end(i8* null, i1 false)
+  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*)*, %i8.array, i1 }
+
+declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*)
+declare i1 @llvm.coro.alloc(token) #3
+declare i64 @llvm.coro.size.i64() #5
+declare i8* @llvm.coro.begin(token, i8* writeonly) #3
+declare token @llvm.coro.save(i8*) #3
+declare i8* @llvm.coro.frame() #5
+declare i8 @llvm.coro.suspend(token, i1) #3
+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