Handle lifetime intrinsics when sinking calls with dead writes

[reames]

Summary

Follow up to D116200, specifically adding the lifetime case.

I'm not terribly happy with this patch, and hoping someone else has a better idea. The tricky case is when the call is writeonly. If it read, we could rely on the block scan done for the read case just below, but we end up having to duplicate a bunch of logic just to handle the writeonly case.

A couple of thoughts:

1. Maybe we just shouldn't handle this case? e.g. allow readwrite calls to be sunk with lifetimes, but bailtout with a comment explaining why?
2. Maybe comesBefore (and potentially numbering the source block) isn't that bad afterall?
3. Maybe we introduce a ScanBlockSuffix flag, reuse the existing code, and accept the false negative of their being unrelated writes in the block suffix? (Given this only causes additional false negatives for writeonly, maybe that's okay?)

[reames]

No review progress, and not currently blocking anything for me.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

if (I->mayWriteToMemory()) {
     if (!AI)
       // TODO: allow malloc?
       return false;
     // TODO: allow memory access dominated by move point?  Note that since AI
     // could have a reference to itself captured by the call, we would need to
     // account for cycles in doing so.
     SmallVector<const User *> AllocaUsers;
     SmallPtrSet<const User *, 4> Visited;
+    SmallPtrSet<const Instruction*, 4> LifetimeBlockers;
     auto pushUsers = [&](const Instruction &I) {
       for (const User *U : I.users()) {
         if (Visited.insert(U).second)
           AllocaUsers.push_back(U);
       }
     };
     pushUsers(*AI);
     while (!AllocaUsers.empty()) {
       auto *UserI = cast<Instruction>(AllocaUsers.pop_back_val());
       if (isa<BitCastInst>(UserI) || isa<GetElementPtrInst>(UserI) ||
           isa<AddrSpaceCastInst>(UserI)) {
         pushUsers(*UserI);
         continue;
       }
       if (UserI == CB)
         continue;
-      // TODO: support lifetime.start/end here
+      // We can't move the call past a lifetime end without introducing UB.
+      if (UserI->isLifetimeStartOrEnd() &&
+          DestBlock->getUniquePredecessor() == I->getParent()) {
+        if (UserI->getParent() == I->getParent())
+          LifetimeBlockers.insert(UserI);
+        continue;
+      }
+
+      // Unhandled case
+      return false;
+    }
+    if (!LifetimeBlockers.empty()) {
+      // If we have lifetime intrinsics in the block, make sure they
+      // preceed I.  We could have used comesBefore above, but this avoids
+      // needing to number the entire, potentially large, source block.
+      for (BasicBlock::iterator Scan = std::next(I->getIterator()),
+                                E = I->getParent()->end();
+           Scan != E; ++Scan) {
+        if (LifetimeBlockers.contains(&*Scan))
           return false;
       }
     }
+  }
 
   // We can only sink load instructions if there is nothing between the load and
   // the end of block that could change the value.
   if (I->mayReadFromMemory()) {
     // We don't want to do any sophisticated alias analysis, so we only check
     // the instructions after I in I's parent block if we try to sink to its
     // successor block.
     if (DestBlock->getUniquePredecessor() != I->getParent())

llvm/test/Transforms/InstCombine/sink_sideeffecting_instruction.ll

 }
 
 define i32 @sink_lifetime1(i1 %c) {
 ; CHECK-LABEL: @sink_lifetime1(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[VAR:%.*]] = alloca i32, align 4
 ; CHECK-NEXT:    [[BITCAST:%.*]] = bitcast i32* [[VAR]] to i8*
 ; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull [[BITCAST]])
-; CHECK-NEXT:    [[VAR3:%.*]] = call i32 @unknown(i32* nonnull [[VAR]]) #[[ATTR1]]
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[EARLY_RETURN:%.*]], label [[USE_BLOCK:%.*]]
 ; CHECK:       early_return:
 ; CHECK-NEXT:    ret i32 0
 ; CHECK:       use_block:
+; CHECK-NEXT:    [[VAR3:%.*]] = call i32 @unknown(i32* nonnull [[VAR]]) #[[ATTR1]]
 ; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull [[BITCAST]])
 ; CHECK-NEXT:    ret i32 [[VAR3]]
 ;
 entry:
   %var = alloca i32, align 4
   %bitcast = bitcast i32* %var to i8*
   call void @llvm.lifetime.start.p0i8(i64 4, i8* %bitcast)
   %var3 = call i32 @unknown(i32* %var) argmemonly nounwind willreturn
   br i1 %c, label %early_return, label %use_block
 
 early_return:
   ret i32 0
 
 use_block:
   call void @llvm.lifetime.end.p0i8(i64 4, i8* %bitcast)
   ret i32 %var3
 }
 
 define i32 @sink_lifetime2(i1 %c) {
 ; CHECK-LABEL: @sink_lifetime2(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[VAR:%.*]] = alloca i32, align 4
 ; CHECK-NEXT:    [[BITCAST:%.*]] = bitcast i32* [[VAR]] to i8*
 ; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull [[BITCAST]])
-; CHECK-NEXT:    [[VAR3:%.*]] = call i32 @unknown(i32* nonnull [[VAR]]) #[[ATTR1]]
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[MERGE:%.*]], label [[USE_BLOCK:%.*]]
 ; CHECK:       merge:
-; CHECK-NEXT:    [[RET:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[VAR3]], [[USE_BLOCK]] ]
+; CHECK-NEXT:    [[RET:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[VAR3:%.*]], [[USE_BLOCK]] ]
 ; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull [[BITCAST]])
 ; CHECK-NEXT:    ret i32 [[RET]]
 ; CHECK:       use_block:
+; CHECK-NEXT:    [[VAR3]] = call i32 @unknown(i32* nonnull [[VAR]]) #[[ATTR1]]
 ; CHECK-NEXT:    br label [[MERGE]]
 ;
 entry:
   %var = alloca i32, align 4
   %bitcast = bitcast i32* %var to i8*
   call void @llvm.lifetime.start.p0i8(i64 4, i8* %bitcast)
   %var3 = call i32 @unknown(i32* %var) argmemonly nounwind willreturn
   br i1 %c, label %merge, label %use_block

 
 early_return:
   ret i32 0
 
 use_block:
   ret i32 %var3
 }
 
-define i32 @sink_lifetime4(i1 %c) {
-; CHECK-LABEL: @sink_lifetime4(
+define i32 @sink_lifetime4a(i1 %c) {
+; CHECK-LABEL: @sink_lifetime4a(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[VAR:%.*]] = alloca i32, align 4
 ; CHECK-NEXT:    [[BITCAST:%.*]] = bitcast i32* [[VAR]] to i8*
 ; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull [[BITCAST]])
 ; CHECK-NEXT:    [[VAR3:%.*]] = call i32 @unknown(i32* nonnull [[VAR]]) #[[ATTR1]]
 ; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull [[BITCAST]])
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[EARLY_RETURN:%.*]], label [[USE_BLOCK:%.*]]
 ; CHECK:       early_return:

 
 early_return:
   ret i32 0
 
 use_block:
   ret i32 %var3
 }
 
+; Version which only writes to var, and thus can't rely on may-read scan for
+; clobbers to prevent the transform
+define i32 @sink_lifetime4b(i1 %c) {
+; CHECK-LABEL: @sink_lifetime4b(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[VAR:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[BITCAST:%.*]] = bitcast i32* [[VAR]] to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull [[BITCAST]])
+; CHECK-NEXT:    [[VAR3:%.*]] = call i32 @unknown(i32* nonnull writeonly [[VAR]]) #[[ATTR1]]
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull [[BITCAST]])
+; CHECK-NEXT:    br i1 [[C:%.*]], label [[EARLY_RETURN:%.*]], label [[USE_BLOCK:%.*]]
+; CHECK:       early_return:
+; CHECK-NEXT:    ret i32 0
+; CHECK:       use_block:
+; CHECK-NEXT:    ret i32 [[VAR3]]
+;
+entry:
+  %var = alloca i32, align 4
+  %bitcast = bitcast i32* %var to i8*
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* %bitcast)
+  %var3 = call i32 @unknown(i32* writeonly %var) argmemonly nounwind willreturn
+  call void @llvm.lifetime.end.p0i8(i64 4, i8* %bitcast)
+  br i1 %c, label %early_return, label %use_block
+
+early_return:
+  ret i32 0
+
+use_block:
+  ret i32 %var3
+}
 ; Mostly checking that trying to sink a non-call doesn't crash (i.e. prior bug)
 define i32 @sink_atomicrmw_to_use(i1 %c) {
 ; CHECK-LABEL: @sink_atomicrmw_to_use(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[VAR:%.*]] = alloca i32, align 4
 ; CHECK-NEXT:    store i32 0, i32* [[VAR]], align 4
 ; CHECK-NEXT:    [[VAR3:%.*]] = atomicrmw add i32* [[VAR]], i32 1 seq_cst, align 4
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[EARLY_RETURN:%.*]], label [[USE_BLOCK:%.*]]