[SelectOpti] Fix lifetime intrinsic bug

[apostolakis]

Summary

When a select is converted to a branch and load instructions are sinked to the true/false blocks,
lifetime intrinsics (if present) could be made unsound if not moved.

This conservatively moves all lifetime intrinsics in a transformed BB to the end block to ensure
preserved lifetime semantics.

[apostolakis]

Add test

[apostolakis]

Move lifetime intrinsics only if there is sinking

[davidxl]

llvm/lib/CodeGen/SelectOptimize.cpp

432

is it better to track the life end variable and inserted after the target instruction?

[davidxl]

It is unclear that the more precise life time tracking can make a big difference given the code region (post transformation) is small, so this handling LGTM.

[apostolakis]

llvm/lib/CodeGen/SelectOptimize.cpp

432

Yes as you said, the potential gains are small since the code region is just a single split basic block.

Also note that, to begin with, currently-emitted lifetimes intrinsics are not very tight to the actual lifetime of objects.
This is probably because other optimizations are also conservative when placing them and moving them .
For example, for the case that caused the bug, the lifetime ended 4 instructions after the final use of the variable.

llvm/lib/CodeGen/SelectOptimize.cpp

 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TargetSchedule.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instruction.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/ProfDataUtils.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/ScaledNumber.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/SizeOpts.h"
 #include <algorithm>
 #include <memory>
for (SelectGroup &ASI : ProfSIGroups) {
     BasicBlock *EndBlock = StartBlock->splitBasicBlock(SplitPt, "select.end");
     BFI->setBlockFreq(EndBlock, BFI->getBlockFreq(StartBlock).getFrequency());
     // Delete the unconditional branch that was just created by the split.
     StartBlock->getTerminator()->eraseFromParent();
 
     // Move any debug/pseudo instructions that were in-between the select
     // group to the newly-created end block.
     SmallVector<Instruction *, 2> DebugPseudoINS;
     auto DIt = SI->getIterator();
     while (&*DIt != LastSI) {
       if (DIt->isDebugOrPseudoInst())
         DebugPseudoINS.push_back(&*DIt);
       DIt++;
     }
     for (auto *DI : DebugPseudoINS) {
       DI->moveBefore(&*EndBlock->getFirstInsertionPt());
     }
if (TrueBlock == FalseBlock) {
       assert(TrueBlock == nullptr &&
              "Unexpected basic block transform while optimizing select");
 
       FalseBlock = BasicBlock::Create(SI->getContext(), "select.false",
                                       EndBlock->getParent(), EndBlock);
       auto *FalseBranch = BranchInst::Create(EndBlock, FalseBlock);
       FalseBranch->setDebugLoc(SI->getDebugLoc());
     }
+    // If there was sinking, move any lifetime-end intrinsic calls found in the
+    // StartBlock to the newly-created end block to ensure sound lifetime info
+    // after sinking (in case any use is sinked).
+    else {
+      SmallVector<Instruction *, 2> EndLifetimeCalls;
+      for (Instruction &I : *StartBlock) {
+        if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I)) {
+          if (II->getIntrinsicID() == Intrinsic::lifetime_end) {
+            EndLifetimeCalls.push_back(&I);
+          }
+        }
+      }
+      for (auto *LC : EndLifetimeCalls) {
+        LC->moveBefore(&*EndBlock->getFirstInsertionPt());
+      }
+    }
 
     // Insert the real conditional branch based on the original condition.
     // If we did not create a new block for one of the 'true' or 'false' paths
     // of the condition, it means that side of the branch goes to the end block
     // directly and the path originates from the start block from the point of
     // view of the new PHI.
     BasicBlock *TT, *FT;
     if (TrueBlock == nullptr) {

llvm/test/CodeGen/X86/select-optimize.ll

 ;
   %load = load i32, ptr %a, align 8
   %sel = select i1 %cmp, i32 %x, i32 %load, !prof !17
   ret i32 %sel
 }
 
 ; Cold value operand with load in its one-use dependence slice shoud result
 ; into a branch with sinked dependence slice.
+; The lifetime-end intrinsics should be soundly preserved.
 define i32 @expensive_val_operand3(ptr nocapture %a, i32 %b, i32 %y, i1 %cmp) {
 ; CHECK-LABEL: @expensive_val_operand3(
 ; CHECK-NEXT:    [[SEL_FROZEN:%.*]] = freeze i1 [[CMP:%.*]]
 ; CHECK-NEXT:    br i1 [[SEL_FROZEN]], label [[SELECT_TRUE_SINK:%.*]], label [[SELECT_END:%.*]], !prof [[PROF18]]
 ; CHECK:       select.true.sink:
 ; CHECK-NEXT:    [[LOAD:%.*]] = load i32, ptr [[A:%.*]], align 8
 ; CHECK-NEXT:    [[X:%.*]] = add i32 [[LOAD]], [[B:%.*]]
 ; CHECK-NEXT:    br label [[SELECT_END]]
 ; CHECK:       select.end:
 ; CHECK-NEXT:    [[SEL:%.*]] = phi i32 [ [[X]], [[SELECT_TRUE_SINK]] ], [ [[Y:%.*]], [[TMP0:%.*]] ]
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 2, ptr nonnull [[A]])
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %load = load i32, ptr %a, align 8
+  call void @llvm.lifetime.end.p0(i64 2, ptr nonnull %a)
   %x = add i32 %load, %b
   %sel = select i1 %cmp, i32 %x, i32 %y, !prof !17
   ret i32 %sel
 }
 
 ; Multiple uses of the load value operand.
 define i32 @expensive_val_operand4(i32 %a, ptr nocapture %b, i32 %x, i1 %cmp) {
 ; CHECK-LABEL: @expensive_val_operand4(
for.body:                                         ; preds = %for.body.preheader, %for.body
   %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
   %exitcond.not = icmp eq i64 %indvars.iv.next, %wide.trip.count
   br i1 %exitcond.not, label %for.cond.cleanup.loopexit, label %for.body
 }
 
 ; Function Attrs: nounwind readnone speculatable willreturn
 declare void @llvm.dbg.value(metadata, metadata, metadata)
 
+; Function Attrs: argmemonly mustprogress nocallback nofree nosync nounwind willreturn
+declare void @llvm.lifetime.end.p0(i64 immarg, ptr nocapture)
+
 !llvm.module.flags = !{!0, !26, !27}
 !0 = !{i32 1, !"ProfileSummary", !1}
 !1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
 !2 = !{!"ProfileFormat", !"InstrProf"}
 !3 = !{!"TotalCount", i64 10000}
 !4 = !{!"MaxCount", i64 10}
 !5 = !{!"MaxInternalCount", i64 1}
 !6 = !{!"MaxFunctionCount", i64 1000}