2 files changed, 81 insertions, 5 deletions

diff --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index c8b0ea8c992..d64c5d24fc8 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -788,15 +788,33 @@ bool DSE::handleEndBlock(BasicBlock &BB) {
 
   const DataLayout &DL = BB.getModule()->getDataLayout();
 
+  // becomes false once lifetime intrinsics are observable or useful for stack
+  // coloring
+  bool canRemoveLifetimeIntrinsics = true;
+
   // Scan the basic block backwards
   for (BasicBlock::iterator BBI = BB.end(); BBI != BB.begin(); ){
     --BBI;
 
-    // If we find a store, check to see if it points into a dead stack value.
-    if (hasMemoryWrite(BBI, *TLI) && isRemovable(BBI)) {
+    Value *V = nullptr;
+    if (canRemoveLifetimeIntrinsics)
+      if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(BBI))
+        switch (II->getIntrinsicID()) {
+          default: break;
+          case Intrinsic::lifetime_start:
+          case Intrinsic::lifetime_end:
+            V = II->getArgOperand(1);
+            break;
+        }
+
+    if (!V && hasMemoryWrite(BBI, *TLI) && isRemovable(BBI))
+      V = getStoredPointerOperand(BBI);
+
+    // If we found a store, check to see if it points into a dead stack value.
+    if (V) {
       // See through pointer-to-pointer bitcasts
       SmallVector<Value *, 4> Pointers;
-      GetUnderlyingObjects(getStoredPointerOperand(BBI), Pointers, DL);
+      GetUnderlyingObjects(V, Pointers, DL);
 
       // Stores to stack values are valid candidates for removal.
       bool AllDead = true;
@@ -844,6 +862,15 @@ bool DSE::handleEndBlock(BasicBlock &BB) {
       continue;
     }
 
+    if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(BBI))
+      if (II->getIntrinsicID() == Intrinsic::lifetime_start) {
+        // We found a lifetime start for a live object, which we could not
+        // remove. So we must stop removing lifetime intrinsics from this block
+        // because they're useful for stack coloring again
+        canRemoveLifetimeIntrinsics = false;
+        continue;
+      }
+
     if (auto CS = CallSite(BBI)) {
       // Remove allocation function calls from the list of dead stack objects; 
       // there can't be any references before the definition.
diff --git a/llvm/test/Transforms/DeadStoreElimination/lifetime.ll b/llvm/test/Transforms/DeadStoreElimination/lifetime.ll
index 305c916dc02..d68e9c0f51c 100644
--- a/llvm/test/Transforms/DeadStoreElimination/lifetime.ll
+++ b/llvm/test/Transforms/DeadStoreElimination/lifetime.ll
@@ -5,14 +5,16 @@ target datalayout = "E-p:64:64:64-a0:0:8-f32:32:32-f64:64:64-i1:8:8-i8:8:8-i16:1
 declare void @llvm.lifetime.start(i64, i8* nocapture) nounwind
 declare void @llvm.lifetime.end(i64, i8* nocapture) nounwind
 declare void @llvm.memset.p0i8.i8(i8* nocapture, i8, i8, i32, i1) nounwind
+declare void @callee(i8*)
 
 define void @test1() {
 ; CHECK-LABEL: @test1(
   %A = alloca i8
 
   store i8 0, i8* %A  ;; Written to by memset
+; CHECK-NOT: store
   call void @llvm.lifetime.end(i64 1, i8* %A)
-; CHECK: lifetime.end
+; CHECK-NOT: lifetime.end
 
   call void @llvm.memset.p0i8.i8(i8* %A, i8 0, i8 -1, i32 0, i1 false)
 ; CHECK-NOT: memset
@@ -22,7 +24,7 @@ define void @test1() {
 }
 
 define void @test2(i32* %P) {
-; CHECK: test2
+; CHECK-LABEL: test2
   %Q = getelementptr i32, i32* %P, i32 1
   %R = bitcast i32* %Q to i8*
   call void @llvm.lifetime.start(i64 4, i8* %R)
@@ -34,4 +36,51 @@ define void @test2(i32* %P) {
   ret void
 }
 
+define void @test3(i8*) {
+; CHECK-LABEL: test3
+  %a = alloca i8
+  call void @llvm.lifetime.start(i64 1, i8* %a)
+; CHECK-NOT: lifetime.start
+  call void @llvm.lifetime.end(i64 1, i8* %a)
+; CHECK-NOT: lifetime.end
+  call void @llvm.lifetime.start(i64 1, i8* undef)
+; CHECK-NOT: lifetime.start
+  call void @llvm.lifetime.end(i64 1, i8* undef)
+; CHECK-NOT: lifetime.end
+  ret void
+}
+
+define void @test4(i8*) {
+; CHECK-LABEL: test4
+  %a = alloca i8
+  call void @llvm.lifetime.start(i64 1, i8* %a)
+; CHECK: lifetime.start
+  call void @llvm.lifetime.end(i64 1, i8* %a)
+; CHECK: lifetime.end
+  call void @llvm.lifetime.start(i64 1, i8* %0)
+; CHECK: lifetime.start
+  call void @llvm.lifetime.end(i64 1, i8* %0)
+; CHECK: lifetime.end
+  call void @llvm.lifetime.start(i64 1, i8* %a)
+; CHECK-NOT: lifetime.start
+  call void @llvm.lifetime.end(i64 1, i8* %a)
+; CHECK-NOT: lifetime.end
+  ret void
+}
 
+define void @test5() {
+; CHECK-LABEL: test5
+  %a = alloca i8
+  %b = alloca i8
+  call void @llvm.lifetime.start(i64 1, i8* %a)
+; CHECK: lifetime.start
+  call void @llvm.lifetime.end(i64 1, i8* %a)
+; CHECK: lifetime.end
+  call void @llvm.lifetime.start(i64 1, i8* %b)
+; CHECK: lifetime.start
+  call void @callee(i8* %b)
+; CHECK: call void @callee
+  call void @llvm.lifetime.end(i64 1, i8* %b)
+; CHECK-NOT: lifetime.end
+  ret void
+}