Finish adding support for lifetime intrinsics to SROA. Fixes PR10121!

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@136008 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 249 insertions, 1 deletions

diff --git a/lib/Transforms/Scalar/ScalarReplAggregates.cpp b/lib/Transforms/Scalar/ScalarReplAggregates.cpp
index 4b1c392cd9..fcc5f1985b 100644
--- a/lib/Transforms/Scalar/ScalarReplAggregates.cpp
+++ b/lib/Transforms/Scalar/ScalarReplAggregates.cpp
@@ -145,6 +145,9 @@ namespace {
                         SmallVector<AllocaInst*, 32> &NewElts);
     void RewriteGEP(GetElementPtrInst *GEPI, AllocaInst *AI, uint64_t Offset,
                     SmallVector<AllocaInst*, 32> &NewElts);
+    void RewriteLifetimeIntrinsic(IntrinsicInst *II, AllocaInst *AI,
+                                  uint64_t Offset,
+                                  SmallVector<AllocaInst*, 32> &NewElts);
     void RewriteMemIntrinUserOfAlloca(MemIntrinsic *MI, Instruction *Inst,
                                       AllocaInst *AI,
                                       SmallVector<AllocaInst*, 32> &NewElts);
@@ -508,7 +511,8 @@ bool ConvertToScalarInfo::CanConvertToScalar(Value *V, uint64_t Offset) {
     }
 
     if (BitCastInst *BCI = dyn_cast<BitCastInst>(User)) {
-      IsNotTrivial = true;  // Can't be mem2reg'd.
+      if (!onlyUsedByLifetimeMarkers(BCI))
+        IsNotTrivial = true;  // Can't be mem2reg'd.
       if (!CanConvertToScalar(BCI, Offset))
         return false;
       continue;
@@ -566,6 +570,14 @@ bool ConvertToScalarInfo::CanConvertToScalar(Value *V, uint64_t Offset) {
       continue;
     }
 
+    // If this is a lifetime intrinsic, we can handle it.
+    if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(User)) {
+      if (II->getIntrinsicID() == Intrinsic::lifetime_start ||
+          II->getIntrinsicID() == Intrinsic::lifetime_end) {
+        continue;
+      }
+    }
+
     // Otherwise, we cannot handle this!
     return false;
   }
@@ -709,6 +721,16 @@ void ConvertToScalarInfo::ConvertUsesToScalar(Value *Ptr, AllocaInst *NewAI,
       continue;
     }
 
+    if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(User)) {
+      if (II->getIntrinsicID() == Intrinsic::lifetime_start ||
+          II->getIntrinsicID() == Intrinsic::lifetime_end) {
+        // There's no need to preserve these, as the resulting alloca will be
+        // converted to a register anyways.
+        II->eraseFromParent();
+        continue;
+      }
+    }
+
     llvm_unreachable("Unsupported operation!");
   }
 }
@@ -1349,6 +1371,13 @@ static bool tryToMakeAllocaBePromotable(AllocaInst *AI, const TargetData *TD) {
       continue;
     }
     
+    if (BitCastInst *BCI = dyn_cast<BitCastInst>(U)) {
+      if (onlyUsedByLifetimeMarkers(BCI)) {
+        InstsToRewrite.insert(BCI);
+        continue;
+      }
+    }
+    
     return false;
   }
 
@@ -1360,6 +1389,18 @@ static bool tryToMakeAllocaBePromotable(AllocaInst *AI, const TargetData *TD) {
   // If we have instructions that need to be rewritten for this to be promotable
   // take care of it now.
   for (unsigned i = 0, e = InstsToRewrite.size(); i != e; ++i) {
+    if (BitCastInst *BCI = dyn_cast<BitCastInst>(InstsToRewrite[i])) {
+      // This could only be a bitcast used by nothing but lifetime intrinsics.
+      for (BitCastInst::use_iterator I = BCI->use_begin(), E = BCI->use_end();
+           I != E;) {
+        Use &U = I.getUse();
+        ++I;
+        cast<Instruction>(U.getUser())->eraseFromParent();
+      }
+      BCI->eraseFromParent();
+      continue;
+    }
+
     if (SelectInst *SI = dyn_cast<SelectInst>(InstsToRewrite[i])) {
       // Selects in InstsToRewrite only have load uses.  Rewrite each as two
       // loads with a new select.
@@ -1692,6 +1733,10 @@ void SROA::isSafeForScalarRepl(Instruction *I, uint64_t Offset,
       isSafeMemAccess(Offset, TD->getTypeAllocSize(SIType),
                       SIType, true, Info, SI, true /*AllowWholeAccess*/);
       Info.hasALoadOrStore = true;
+    } else if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(User)) {
+      if (II->getIntrinsicID() != Intrinsic::lifetime_start &&
+          II->getIntrinsicID() != Intrinsic::lifetime_end)
+        return MarkUnsafe(Info, User);
     } else if (isa<PHINode>(User) || isa<SelectInst>(User)) {
       isSafePHISelectUseForScalarRepl(User, Offset, Info);
     } else {
@@ -1929,6 +1974,14 @@ void SROA::RewriteForScalarRepl(Instruction *I, AllocaInst *AI, uint64_t Offset,
       // address operand will be updated, so nothing else needs to be done.
       continue;
     }
+
+    if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(User)) {
+      if (II->getIntrinsicID() == Intrinsic::lifetime_start ||
+          II->getIntrinsicID() == Intrinsic::lifetime_end) {
+        RewriteLifetimeIntrinsic(II, AI, Offset, NewElts);
+      }
+      continue;
+    }
     
     if (LoadInst *LI = dyn_cast<LoadInst>(User)) {
       Type *LIType = LI->getType();
@@ -2095,6 +2148,62 @@ void SROA::RewriteGEP(GetElementPtrInst *GEPI, AllocaInst *AI, uint64_t Offset,
   DeadInsts.push_back(GEPI);
 }
 
+/// RewriteLifetimeIntrinsic - II is a lifetime.start/lifetime.end. Rewrite it
+/// to mark the lifetime of the scalarized memory.
+void SROA::RewriteLifetimeIntrinsic(IntrinsicInst *II, AllocaInst *AI,
+                                    uint64_t Offset,
+                                    SmallVector<AllocaInst*, 32> &NewElts) {
+  ConstantInt *OldSize = cast<ConstantInt>(II->getArgOperand(0));
+  // Put matching lifetime markers on everything from Offset up to
+  // Offset+OldSize.
+  Type *AIType = AI->getAllocatedType();
+  uint64_t NewOffset = Offset;
+  Type *IdxTy;
+  uint64_t Idx = FindElementAndOffset(AIType, NewOffset, IdxTy);
+
+  IRBuilder<> Builder(II);
+  uint64_t Size = OldSize->getLimitedValue();
+
+  if (NewOffset) {
+    // Splice the first element and index 'NewOffset' bytes in.  SROA will
+    // split the alloca again later.
+    Value *V = Builder.CreateBitCast(NewElts[Idx], Builder.getInt8PtrTy());
+    V = Builder.CreateGEP(V, Builder.getInt64(NewOffset));
+
+    IdxTy = NewElts[Idx]->getAllocatedType();
+    uint64_t EltSize = TD->getTypeAllocSize(IdxTy) - NewOffset;
+    if (EltSize > Size) {
+      EltSize = Size;
+      Size = 0;
+    } else {
+      Size -= EltSize;
+    }
+    if (II->getIntrinsicID() == Intrinsic::lifetime_start)
+      Builder.CreateLifetimeStart(V, Builder.getInt64(EltSize));
+    else
+      Builder.CreateLifetimeEnd(V, Builder.getInt64(EltSize));
+    ++Idx;
+  }
+
+  for (; Idx != NewElts.size() && Size; ++Idx) {
+    IdxTy = NewElts[Idx]->getAllocatedType();
+    uint64_t EltSize = TD->getTypeAllocSize(IdxTy);
+    if (EltSize > Size) {
+      EltSize = Size;
+      Size = 0;
+    } else {
+      Size -= EltSize;
+    }
+    if (II->getIntrinsicID() == Intrinsic::lifetime_start)
+      Builder.CreateLifetimeStart(NewElts[Idx],
+                                  Builder.getInt64(EltSize));
+    else
+      Builder.CreateLifetimeEnd(NewElts[Idx],
+                                Builder.getInt64(EltSize));
+  }
+  DeadInsts.push_back(II);
+}
+
 /// RewriteMemIntrinUserOfAlloca - MI is a memcpy/memset/memmove from or to AI.
 /// Rewrite it to copy or set the elements of the scalarized memory.
 void SROA::RewriteMemIntrinUserOfAlloca(MemIntrinsic *MI, Instruction *Inst,
diff --git a/test/Transforms/ScalarRepl/lifetime.ll b/test/Transforms/ScalarRepl/lifetime.ll
new file mode 100644
index 0000000000..3f558a1c37
--- /dev/null
+++ b/test/Transforms/ScalarRepl/lifetime.ll
@@ -0,0 +1,139 @@
+; RUN: opt -scalarrepl -S < %s | FileCheck %s
+
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64"
+target triple = "x86_64-unknown-linux-gnu"
+
+declare void @llvm.lifetime.start(i64, i8*)
+declare void @llvm.lifetime.end(i64, i8*)
+
+%t1 = type {i32, i32, i32}
+
+define void @test1() {
+; CHECK: @test1
+  %A = alloca %t1
+  %A1 = getelementptr %t1* %A, i32 0, i32 0
+  %A2 = getelementptr %t1* %A, i32 0, i32 1
+  %A3 = getelementptr %t1* %A, i32 0, i32 2
+  %B = bitcast i32* %A1 to i8*
+  store i32 0, i32* %A1
+  call void @llvm.lifetime.start(i64 -1, i8* %B)
+  ret void
+; CHECK-NEXT: ret void
+}
+
+define void @test2() {
+; CHECK: @test2
+  %A = alloca %t1
+  %A1 = getelementptr %t1* %A, i32 0, i32 0
+  %A2 = getelementptr %t1* %A, i32 0, i32 1
+  %A3 = getelementptr %t1* %A, i32 0, i32 2
+  %B = bitcast i32* %A2 to i8*
+  store i32 0, i32* %A2
+  call void @llvm.lifetime.start(i64 -1, i8* %B)
+  %C = load i32* %A2
+  ret void
+; CHECK: ret void
+}
+
+define void @test3() {
+; CHECK: @test3
+  %A = alloca %t1
+  %A1 = getelementptr %t1* %A, i32 0, i32 0
+  %A2 = getelementptr %t1* %A, i32 0, i32 1
+  %A3 = getelementptr %t1* %A, i32 0, i32 2
+  %B = bitcast i32* %A2 to i8*
+  store i32 0, i32* %A2
+  call void @llvm.lifetime.start(i64 6, i8* %B)
+  %C = load i32* %A2
+  ret void
+; CHECK-NEXT: ret void
+}
+
+define void @test4() {
+; CHECK: @test4
+  %A = alloca %t1
+  %A1 = getelementptr %t1* %A, i32 0, i32 0
+  %A2 = getelementptr %t1* %A, i32 0, i32 1
+  %A3 = getelementptr %t1* %A, i32 0, i32 2
+  %B = bitcast i32* %A2 to i8*
+  store i32 0, i32* %A2
+  call void @llvm.lifetime.start(i64 1, i8* %B)
+  %C = load i32* %A2
+  ret void
+; CHECK-NEXT: ret void
+}
+
+%t2 = type {i32, [4 x i8], i32}
+
+define void @test5() {
+; CHECK: @test5
+  %A = alloca %t2
+; CHECK: alloca{{.*}}i8
+; CHECK: alloca{{.*}}i8
+; CHECK: alloca{{.*}}i8
+
+  %A21 = getelementptr %t2* %A, i32 0, i32 1, i32 0
+  %A22 = getelementptr %t2* %A, i32 0, i32 1, i32 1
+  %A23 = getelementptr %t2* %A, i32 0, i32 1, i32 2
+  %A24 = getelementptr %t2* %A, i32 0, i32 1, i32 3
+; CHECK-NOT: store i8 1
+  store i8 1, i8* %A21
+  store i8 2, i8* %A22
+  store i8 3, i8* %A23
+  store i8 4, i8* %A24
+
+  %A1 = getelementptr %t2* %A, i32 0, i32 0
+  %A2 = getelementptr %t2* %A, i32 0, i32 1, i32 1
+  %A3 = getelementptr %t2* %A, i32 0, i32 2
+  store i8 0, i8* %A2
+  call void @llvm.lifetime.start(i64 5, i8* %A2)
+; CHECK: llvm.lifetime{{.*}}i64 1
+; CHECK: llvm.lifetime{{.*}}i64 1
+; CHECK: llvm.lifetime{{.*}}i64 1
+  %C = load i8* %A2
+  ret void
+}
+
+%t3 = type {[4 x i16], [4 x i8]}
+
+define void @test6() {
+; CHECK: @test6
+  %A = alloca %t3
+; CHECK: alloca i8
+; CHECK: alloca i8
+; CHECK: alloca i8
+
+  %A11 = getelementptr %t3* %A, i32 0, i32 0, i32 0
+  %A12 = getelementptr %t3* %A, i32 0, i32 0, i32 1
+  %A13 = getelementptr %t3* %A, i32 0, i32 0, i32 2
+  %A14 = getelementptr %t3* %A, i32 0, i32 0, i32 3
+  store i16 11, i16* %A11
+  store i16 12, i16* %A12
+  store i16 13, i16* %A13
+  store i16 14, i16* %A14
+; CHECK-NOT: store i16 11
+; CHECK-NOT: store i16 12
+; CHECK-NOT: store i16 13
+; CHECK-NOT: store i16 14
+
+  %A21 = getelementptr %t3* %A, i32 0, i32 1, i32 0
+  %A22 = getelementptr %t3* %A, i32 0, i32 1, i32 1
+  %A23 = getelementptr %t3* %A, i32 0, i32 1, i32 2
+  %A24 = getelementptr %t3* %A, i32 0, i32 1, i32 3
+  store i8 21, i8* %A21
+  store i8 22, i8* %A22
+  store i8 23, i8* %A23
+  store i8 24, i8* %A24
+; CHECK: store i8 21
+; CHECK: store i8 22
+; CHECK: store i8 23
+; CHECK-NOT: store i8 24
+
+  %B = bitcast i16* %A13 to i8*
+  call void @llvm.lifetime.start(i64 7, i8* %B)
+; CHECK: lifetime.start{{.*}}i64 1
+; CHECK: lifetime.start{{.*}}i64 1
+; CHECK: lifetime.start{{.*}}i64 1
+
+  ret void
+}