[CodeExtractor] Do not lift lifetime.end markers for region inputs

[vsk]

Summary

If a lifetime.end marker occurs along one path through the extraction
region, but not another, then it's still incorrect to lift the marker,
because there is some path through the extracted function which would
ordinarily not reach the marker. If the call to the extracted function
is in a loop, unrolling can cause inputs to the function to become
optimized out as undef after the first iteration.

To prevent incorrect stack slot merging in the calling function, it
should be sufficient to lift lifetime.start markers for region inputs.
I've tested this theory out by doing a stage2 check-all with randomized
splitting enabled.

This is a follow-up to r353973, and there's additional context for this
change in https://reviews.llvm.org/D57834.

rdar://47896986

[kachkov98]

LGTM

llvm/lib/Transforms/Utils/CodeExtractor.cpp

if (Instruction *I = dyn_cast<Instruction>(Users[i]))
       if (I->isTerminator() && !Blocks.count(I->getParent()) &&
           I->getParent()->getParent() == oldFunction)
         I->replaceUsesOfWith(header, newHeader);
 
   return newFunction;
 }
 
 /// Erase lifetime.start markers which reference inputs to the extraction
-/// region, and insert the referenced memory into \p LifetimesStart. Do the same
-/// with lifetime.end markers (but insert them into \p LifetimesEnd).
+/// region, and insert the referenced memory into \p LifetimesStart.
 ///
 /// The extraction region is defined by a set of blocks (\p Blocks), and a set
 /// of allocas which will be moved from the caller function into the extracted
 /// function (\p SunkAllocas).
 static void eraseLifetimeMarkersOnInputs(const SetVector<BasicBlock *> &Blocks,
                                          const SetVector<Value *> &SunkAllocas,
-                                         SetVector<Value *> &LifetimesStart,
-                                         SetVector<Value *> &LifetimesEnd) {
+                                         SetVector<Value *> &LifetimesStart) {
   for (BasicBlock *BB : Blocks) {
     for (auto It = BB->begin(), End = BB->end(); It != End;) {
       auto *II = dyn_cast<IntrinsicInst>(&*It);
       ++It;
       if (!II || !II->isLifetimeStartOrEnd())
         continue;
 
       // Get the memory operand of the lifetime marker. If the underlying
       // object is a sunk alloca, or is otherwise defined in the extraction
       // region, the lifetime marker must not be erased.
       Value *Mem = II->getOperand(1)->stripInBoundsOffsets();
       if (SunkAllocas.count(Mem) || definedInRegion(Blocks, Mem))
         continue;
 
       if (II->getIntrinsicID() == Intrinsic::lifetime_start)
         LifetimesStart.insert(Mem);
-      else
-        LifetimesEnd.insert(Mem);
       II->eraseFromParent();
     }
   }
 }
 
 /// Insert lifetime start/end markers surrounding the call to the new function
 /// for objects defined in the caller.
 static void insertLifetimeMarkersSurroundingCall(
Function *CodeExtractor::extractCodeRegion() {
   if (!HoistingCands.empty()) {
     auto *HoistToBlock = findOrCreateBlockForHoisting(CommonExit);
     Instruction *TI = HoistToBlock->getTerminator();
     for (auto *II : HoistingCands)
       cast<Instruction>(II)->moveBefore(TI);
   }
 
   // Collect objects which are inputs to the extraction region and also
-  // referenced by lifetime start/end markers within it. The effects of these
+  // referenced by lifetime start markers within it. The effects of these
   // markers must be replicated in the calling function to prevent the stack
   // coloring pass from merging slots which store input objects.
-  ValueSet LifetimesStart, LifetimesEnd;
-  eraseLifetimeMarkersOnInputs(Blocks, SinkingCands, LifetimesStart,
-                               LifetimesEnd);
+  ValueSet LifetimesStart;
+  eraseLifetimeMarkersOnInputs(Blocks, SinkingCands, LifetimesStart);
 
   // Construct new function based on inputs/outputs & add allocas for all defs.
   Function *newFunction =
       constructFunction(inputs, outputs, header, newFuncRoot, codeReplacer,
                         oldFunction, oldFunction->getParent());
 
   // Update the entry count of the function.
   if (BFI) {
     auto Count = BFI->getProfileCountFromFreq(EntryFreq.getFrequency());
     if (Count.hasValue())
       newFunction->setEntryCount(
           ProfileCount(Count.getValue(), Function::PCT_Real)); // FIXME
     BFI->setBlockFreq(codeReplacer, EntryFreq.getFrequency());
   }
 
   CallInst *TheCall =
       emitCallAndSwitchStatement(newFunction, codeReplacer, inputs, outputs);
 
   moveCodeToFunction(newFunction);
 
   // Replicate the effects of any lifetime start/end markers which referenced
   // input objects in the extraction region by placing markers around the call.
-  insertLifetimeMarkersSurroundingCall(oldFunction->getParent(),
-                                       LifetimesStart.getArrayRef(),
-                                       LifetimesEnd.getArrayRef(), TheCall);
+  insertLifetimeMarkersSurroundingCall(
+      oldFunction->getParent(), LifetimesStart.getArrayRef(), {}, TheCall);
 
   // Propagate personality info to the new function if there is one.
   if (oldFunction->hasPersonalityFn())
     newFunction->setPersonalityFn(oldFunction->getPersonalityFn());
 
   // Update the branch weights for the exit block.
   if (BFI && NumExitBlocks > 1)
     calculateNewCallTerminatorWeights(codeReplacer, ExitWeights, BPI);

llvm/test/Transforms/CodeExtractor/PartialInlineAlloca4.ll

 ; RUN: opt < %s -partial-inliner -skip-partial-inlining-cost-analysis -S | FileCheck  %s
 ; RUN: opt < %s -passes=partial-inliner -skip-partial-inlining-cost-analysis -S | FileCheck   %s
 
 %"class.base" = type { %"struct.base"* }
 %"struct.base" = type opaque
 
 @g = external local_unnamed_addr global i32, align 4
 
 ; CHECK-LABEL: define{{.*}}@caller(
 ; CHECK: call void @llvm.lifetime.start.p0i8(i64 -1, i8* %tmp.i)
 ; CHECK-NEXT: call void @callee_unknown_use1.{{.*}}(i8* %tmp.i
-; CHECK-NEXT: call void @llvm.lifetime.end.p0i8(i64 -1, i8* %tmp.i)
 
 define i32 @callee_unknown_use1(i32 %arg) local_unnamed_addr #0 {
 ; CHECK-LABEL:define{{.*}}@callee_unknown_use1.{{[0-9]}}
 ; CHECK-NOT: alloca
 bb:
   %tmp = alloca  i8, align 4
   %tmp2 = load i32, i32* @g, align 4, !tbaa !2
   %tmp3 = add nsw i32 %tmp2, 1

llvm/test/Transforms/HotColdSplit/lifetime-markers-on-inputs-1.ll

normalPath:
   ret void
 
 ; CHECK-LABEL: codeRepl:
 ; CHECK: [[local1_cast:%.*]] = bitcast i256* %local1 to i8*
 ; CHECK-NEXT: call void @llvm.lifetime.start.p0i8(i64 -1, i8* [[local1_cast]])
 ; CHECK-NEXT: [[local2_cast:%.*]] = bitcast i256* %local2 to i8*
 ; CHECK-NEXT: call void @llvm.lifetime.start.p0i8(i64 -1, i8* [[local2_cast]])
 ; CHECK-NEXT: call i1 @foo.cold.1(i8* %local1_cast, i8* %local2_cast)
-; CHECK-NEXT: call void @llvm.lifetime.end.p0i8(i64 -1, i8* [[local1_cast]])
-; CHECK-NEXT: call void @llvm.lifetime.end.p0i8(i64 -1, i8* [[local2_cast]])
 ; CHECK-NEXT: br i1
 
 outlinedPath:
   ; These two uses of stack slots are overlapping. This should prevent
   ; merging of stack slots. CodeExtractor must replicate the effects of
   ; these markers in the caller to inhibit stack coloring.
   %gep1 = getelementptr inbounds i8, i8* %local1_cast, i64 1
   call void @llvm.lifetime.start.p0i8(i64 1, i8* %gep1)

llvm/test/Transforms/HotColdSplit/lifetime-markers-on-inputs-2.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -hotcoldsplit -hotcoldsplit-threshold=0 < %s 2>&1 | FileCheck %s
 
 declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture)
 
 declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture)
 
 declare void @cold_use(i8*) cold
 
+declare void @use(i8*)
+
 ; In this CFG, splitting will extract the blocks extract{1,2}. I.e., it will
 ; extract a lifetime.start marker, but not the corresponding lifetime.end
 ; marker. Make sure that a lifetime.start marker is emitted before the call to
 ; the split function, and *only* that marker.
 ;
 ;            entry
 ;          /       \
 ;      extract1  no-extract1

 
 exit:
   ; lt.end
   call void @llvm.lifetime.end.p0i8(i64 1, i8* %local1_cast)
   ret void
 }
 
 ; In this CFG, splitting will extract the block extract1. I.e., it will extract
-; a lifetime.end marker, but not the corresponding lifetime.start marker. Make
-; sure that a lifetime.end marker is emitted after the call to the split
-; function, and *only* that marker.
+; a lifetime.end marker, but not the corresponding lifetime.start marker. Do
+; not emit a lifetime.end marker after the call to the split function.
 ;
 ;            entry
 ;         (lt.start)
 ;        /          \
 ;   no-extract1  extract1
 ;    (lt.end)    (lt.end)
 ;        \         /
 ;            exit
 ;
 ; After splitting, we should see:
 ;
 ;            entry
 ;         (lt.start)
 ;        /          \
 ;   no-extract1  codeRepl
-;    (lt.end)    (lt.end)
+;    (lt.end)
 ;        \         /
 ;            exit
 define void @only_lifetime_end_is_cold() {
 ; CHECK-LABEL: @only_lifetime_end_is_cold(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[LOCAL1:%.*]] = alloca i256
 ; CHECK-NEXT:    [[LOCAL1_CAST:%.*]] = bitcast i256* [[LOCAL1]] to i8*
 ; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 1, i8* [[LOCAL1_CAST]])
 ; CHECK-NEXT:    br i1 undef, label [[NO_EXTRACT1:%.*]], label [[CODEREPL:%.*]]
 ; CHECK:       no-extract1:
 ; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 1, i8* [[LOCAL1_CAST]])
 ; CHECK-NEXT:    br label [[EXIT:%.*]]
 ; CHECK:       codeRepl:
-; CHECK-NEXT:    [[LT_CAST:%.*]] = bitcast i256* [[LOCAL1]] to i8*
 ; CHECK-NEXT:    call void @only_lifetime_end_is_cold.cold.1(i8* [[LOCAL1_CAST]]) #3
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 -1, i8* [[LT_CAST]])
 ; CHECK-NEXT:    br label [[EXIT]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   ; lt.start
   %local1 = alloca i256
   %local1_cast = bitcast i256* %local1 to i8*
extract1:
   ; lt.end
   call void @cold_use(i8* %local1_cast)
   call void @llvm.lifetime.end.p0i8(i64 1, i8* %local1_cast)
   br label %exit
 
 exit:
   ret void
 }
+
+; In this CFG, splitting will extract the blocks extract{1,2,3}. Lifting the
+; lifetime.end marker would be a miscompile.
+define void @do_not_lift_lifetime_end() {
+; CHECK-LABEL: @do_not_lift_lifetime_end(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[LOCAL1:%.*]] = alloca i256
+; CHECK-NEXT:    [[LOCAL1_CAST:%.*]] = bitcast i256* [[LOCAL1]] to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 1, i8* [[LOCAL1_CAST]])
+; CHECK-NEXT:    br label [[HEADER:%.*]]
+; CHECK:       header:
+; CHECK-NEXT:    call void @use(i8* [[LOCAL1_CAST]])
+; CHECK-NEXT:    br i1 undef, label [[EXIT:%.*]], label [[CODEREPL:%.*]]
+; CHECK:       codeRepl:
+; CHECK-NEXT:    [[TARGETBLOCK:%.*]] = call i1 @do_not_lift_lifetime_end.cold.1(i8* [[LOCAL1_CAST]]) #3
+; CHECK-NEXT:    br i1 [[TARGETBLOCK]], label [[HEADER]], label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  ; lt.start
+  %local1 = alloca i256
+  %local1_cast = bitcast i256* %local1 to i8*
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %local1_cast)
+  br label %header
+
+header:
+  ; If the lifetime.end marker is lifted, this use becomes dead the second time
+  ; the header block is executed.
+  call void @use(i8* %local1_cast)
+  br i1 undef, label %exit, label %extract1
+
+extract1:
+  call void @cold_use(i8* %local1_cast)
+  br i1 undef, label %extract2, label %extract3
+
+extract2:
+  ; Backedge.
+  br label %header
+
+extract3:
+  ; lt.end
+  call void @llvm.lifetime.end.p0i8(i64 1, i8* %local1_cast)
+  br label %exit
+
+exit:
+  ret void
+}