[CodeExtractor] Do not extract unsafe lifetime markers

[vsk]

Summary

Lifetime markers which reference inputs to the extraction region are not
safe to extract. Example ('rhs' will be extracted):

  		entry:
  	      x = alloca
  	      y = alloca
         /              \
       lhs:             rhs:
lifetime_start(x)   lifetime_start(x)
use(x)    	      lifetime_start(y)
lifetime_end(x)     use(x, y)
lifetime_start(y)   lifetime_end(y)
use(y)              lifetime_end(x)
lifetime_end(y)     ...
...

Prior to extraction, the stack coloring pass sees that the slots for 'x'
and 'y' are in-use at the same time. After extraction, the coloring pass
infers that 'x' and 'y' are *not* in-use concurrently, because markers
from 'rhs' are no longer available to help decide otherwise.

This leads to a miscompile, because the stack slots actually are in-use
concurrently in the extracted function.

Fix this by moving lifetime start/end markers for memory regions defined
in the calling function around the call to the extracted function.

Fixes llvm.org/PR39671 (rdar://45939472).

[davidxl]

llvm/test/Transforms/CodeExtractor/PartialInlineAlloca4.ll

12

The purpose of the extraction is to shrink the alloca call into the outlined function and thus reducing call overhead of the outline function. This fix breaks that.

[vsk]

llvm/test/Transforms/CodeExtractor/PartialInlineAlloca4.ll

12

Here, the alloca remains in 'caller' both without and with this patch. Are you suggesting that lifetime markers have an affect on codegen even when the referenced alloca is in a different function? I don't see why this would be, but if so, would the right fix be to only move lifetime markers out of outlined code when there are conflicting markers in the caller?

[davidxl]

llvm/test/Transforms/CodeExtractor/PartialInlineAlloca4.ll

12

Sorry I mis-read the test case.

[davidxl]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

1188

document the second parameter?

1198

Is there a common utility function to detect lifetime marker intrinsics?

1199

Is there a better API to access the mem operand instead of using hard coded operand number?

[kachkov98]

I'm concerned about skipping processing markers for outputs. Currently HotColdSplit and partial inlining do not outline function if number of outputs is not zero (by calling findInputsOutputs()), but this check is unreliable because there is some processing of region before extracting (splitReturnBlocks, splitPHINodesOfEntry, splitPHINodesOfExit) and input/output parameters can be changed (their number may become non-zero). So another question is why this check should be done at all (does it bypass some known problem?) Seems that CodeExtractor does not impose such restrictions.

[vsk]

(Marking this WIP while I address @davidxl 's feedback.)

@kachkov98 -- The |outputs| > 0 case had known issues, but I believe they were all resolved by D55018. I have a follow-up planned to enable hot/cold splitting when |outputs| > 0.

Secondly, I think you're correct, it does seem beneficial to add lifetime markers around the reloads for output values. However, I don't think it's needed for correctness, because it's impossible for there to be lifetime markers for output allocas in the caller function prior to extraction. For that reason I'd prefer to do this as a follow-up.

[vsk]

• Address feedback from @davidxl

llvm/lib/Transforms/Utils/CodeExtractor.cpp

1198

I'll rebase this on top of D56019.

1199

Unfortunately not. I'll leave a better comment to explain what is happening here.

[vsk]

Ping.

[kachkov98]

LGTM

[davidxl]

lgtm

llvm/include/llvm/Transforms/Utils/CodeExtractor.h

 #include <limits>
 
 namespace llvm {
 
 class BasicBlock;
 class BlockFrequency;
 class BlockFrequencyInfo;
 class BranchProbabilityInfo;
+class CallInst;
 class DominatorTree;
 class Function;
 class Instruction;
 class Loop;
 class Module;
 class Type;
 class Value;
 
private:
 
     void moveCodeToFunction(Function *newFunction);
 
     void calculateNewCallTerminatorWeights(
         BasicBlock *CodeReplacer,
         DenseMap<BasicBlock *, BlockFrequency> &ExitWeights,
         BranchProbabilityInfo *BPI);
 
-    void emitCallAndSwitchStatement(Function *newFunction,
+    CallInst *emitCallAndSwitchStatement(Function *newFunction,
                                          BasicBlock *newHeader,
-                                    ValueSet &inputs,
-                                    ValueSet &outputs);
+                                         ValueSet &inputs, ValueSet &outputs);
   };
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_CODEEXTRACTOR_H

llvm/lib/Transforms/Utils/CodeExtractor.cpp

if (Instruction *I = dyn_cast<Instruction>(Users[i]))
         I->replaceUsesOfWith(header, newHeader);
 
   return newFunction;
 }
 
 /// emitCallAndSwitchStatement - This method sets up the caller side by adding
 /// the call instruction, splitting any PHI nodes in the header block as
 /// necessary.
-void CodeExtractor::
-emitCallAndSwitchStatement(Function *newFunction, BasicBlock *codeReplacer,
-                           ValueSet &inputs, ValueSet &outputs) {
+CallInst *CodeExtractor::emitCallAndSwitchStatement(Function *newFunction,
+                                                    BasicBlock *codeReplacer,
+                                                    ValueSet &inputs,
+                                                    ValueSet &outputs) {
   // Emit a call to the new function, passing in: *pointer to struct (if
   // aggregating parameters), or plan inputs and allocated memory for outputs
   std::vector<Value *> params, StructValues, ReloadOutputs, Reloads;
 
   Module *M = newFunction->getParent();
   LLVMContext &Context = M->getContext();
   const DataLayout &DL = M->getDataLayout();
+  CallInst *call = nullptr;
 
   // Add inputs as params, or to be filled into the struct
   for (Value *input : inputs)
     if (AggregateArgs)
       StructValues.push_back(input);
     else
       params.push_back(input);
 
for (unsigned i = 0, e = inputs.size(); i != e; ++i) {
           StructArgTy, Struct, Idx, "gep_" + StructValues[i]->getName());
       codeReplacer->getInstList().push_back(GEP);
       StoreInst *SI = new StoreInst(StructValues[i], GEP);
       codeReplacer->getInstList().push_back(SI);
     }
   }
 
   // Emit the call to the function
-  CallInst *call = CallInst::Create(newFunction, params,
-                                    NumExitBlocks > 1 ? "targetBlock" : "");
+  call = CallInst::Create(newFunction, params,
   // Add debug location to the new call, if the original function has debug
   // info. In that case, the terminator of the entry block of the extracted
   // function contains the first debug location of the extracted function,
   // set in extractCodeRegion.
   if (codeReplacer->getParent()->getSubprogram()) {
     if (auto DL = newFunction->getEntryBlock().getTerminator()->getDebugLoc())
       call->setDebugLoc(DL);
   }
default:
     // Otherwise, make the default destination of the switch instruction be one
     // of the other successors.
     TheSwitch->setCondition(call);
     TheSwitch->setDefaultDest(TheSwitch->getSuccessor(NumExitBlocks));
     // Remove redundant case
     TheSwitch->removeCase(SwitchInst::CaseIt(TheSwitch, NumExitBlocks-1));
     break;
   }
+
+  return call;
 }
 
 void CodeExtractor::moveCodeToFunction(Function *newFunction) {
   Function *oldFunc = (*Blocks.begin())->getParent();
   Function::BasicBlockListType &oldBlocks = oldFunc->getBasicBlockList();
   Function::BasicBlockListType &newBlocks = newFunction->getBasicBlockList();
 
   for (BasicBlock *Block : Blocks) {
for (unsigned I = 0, E = BranchDist.Weights.size(); I < E; ++I) {
     BranchProbability BP(Weight.Amount, BranchDist.Total);
     BPI->setEdgeProbability(CodeReplacer, Weight.TargetNode.Index, BP);
   }
   TI->setMetadata(
       LLVMContext::MD_prof,
       MDBuilder(TI->getContext()).createBranchWeights(BranchWeights));
 }
 
+/// Scan the extraction region for lifetime markers which reference inputs.
+/// Erase these markers. Return the inputs which were referenced.
+///
+/// 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 SetVector<Value *>
+eraseLifetimeMarkersOnInputs(const SetVector<BasicBlock *> &Blocks,
+                             const SetVector<Value *> &SunkAllocas) {
+  SetVector<Value *> InputObjectsWithLifetime;
+  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;
+
+      InputObjectsWithLifetime.insert(Mem);
+      II->eraseFromParent();
+    }
+  }
+  return InputObjectsWithLifetime;
+}
+
+/// Insert lifetime start/end markers surrounding the call to the new function
+/// for objects defined in the caller.
+static void insertLifetimeMarkersSurroundingCall(
+    Module *M, const SetVector<Value *> &InputObjectsWithLifetime,
+    CallInst *TheCall) {
+  if (InputObjectsWithLifetime.empty())
+    return;
+
+  LLVMContext &Ctx = M->getContext();
+  auto Int8PtrTy = Type::getInt8PtrTy(Ctx);
+  auto NegativeOne = ConstantInt::getSigned(Type::getInt64Ty(Ctx), -1);
+  auto LifetimeStartFn = llvm::Intrinsic::getDeclaration(
+      M, llvm::Intrinsic::lifetime_start, Int8PtrTy);
+  auto LifetimeEndFn = llvm::Intrinsic::getDeclaration(
+      M, llvm::Intrinsic::lifetime_end, Int8PtrTy);
+  for (Value *Mem : InputObjectsWithLifetime) {
+    assert((!isa<Instruction>(Mem) ||
+            cast<Instruction>(Mem)->getFunction() == TheCall->getFunction()) &&
+           "Input memory not defined in original function");
+    Value *MemAsI8Ptr = nullptr;
+    if (Mem->getType() == Int8PtrTy)
+      MemAsI8Ptr = Mem;
+    else
+      MemAsI8Ptr =
+          CastInst::CreatePointerCast(Mem, Int8PtrTy, "lt.cast", TheCall);
+
+    auto StartMarker =
+        CallInst::Create(LifetimeStartFn, {NegativeOne, MemAsI8Ptr});
+    StartMarker->insertBefore(TheCall);
+    auto EndMarker = CallInst::Create(LifetimeEndFn, {NegativeOne, MemAsI8Ptr});
+    EndMarker->insertAfter(TheCall);
+  }
+}
+
 Function *CodeExtractor::extractCodeRegion() {
   if (!isEligible())
     return nullptr;
 
   // Assumption: this is a single-entry code region, and the header is the first
   // block in the region.
   BasicBlock *header = *Blocks.begin();
   Function *oldFunction = header->getParent();
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
+  // markers must be replicated in the calling function to prevent the stack
+  // coloring pass from merging slots which store input objects.
+  ValueSet InputObjectsWithLifetime =
+      eraseLifetimeMarkersOnInputs(Blocks, SinkingCands);
+
   // Construct new function based on inputs/outputs & add allocas for all defs.
-  Function *newFunction = constructFunction(inputs, outputs, header,
-                                            newFuncRoot,
-                                            codeReplacer, oldFunction,
-                                            oldFunction->getParent());
+  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(),
+                                       InputObjectsWithLifetime, 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
-; CHECK: call void @llvm.lifetime
 bb:
   %tmp = alloca  i8, align 4
   %tmp2 = load i32, i32* @g, align 4, !tbaa !2
   %tmp3 = add nsw i32 %tmp2, 1
   %tmp4 = icmp slt i32 %arg, 0
   br i1 %tmp4, label %bb6, label %bb5
 
 bb5:                                              ; preds = %bb

llvm/test/Transforms/CodeExtractor/PartialInlineAlloca5.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
 
 define i32 @callee_unknown_use2(i32 %arg) local_unnamed_addr #0 {
 ; CHECK-LABEL:define{{.*}}@callee_unknown_use2.{{[0-9]}}
 ; CHECK-NOT: alloca
-; CHECK: call void @llvm.lifetime
 bb:
   %tmp = alloca i32, align 4
   %tmp1 = bitcast i32* %tmp to i8*
   %tmp2 = load i32, i32* @g, align 4, !tbaa !2
   %tmp3 = add nsw i32 %tmp2, 1
   %tmp4 = icmp slt i32 %arg, 0
   br i1 %tmp4, label %bb6, label %bb5
 

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

+; RUN: opt -S -hotcoldsplit < %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 @use(i8*)
+
+declare void @cold_use2(i8*, i8*) cold
+
+; CHECK-LABEL: define {{.*}}@foo(
+define void @foo() {
+entry:
+  %local1 = alloca i256
+  %local2 = alloca i256
+  %local1_cast = bitcast i256* %local1 to i8*
+  %local2_cast = bitcast i256* %local2 to i8*
+  br i1 undef, label %normalPath, label %outlinedPath
+
+normalPath:
+  ; These two uses of stack slots are non-overlapping. Based on this alone,
+  ; the stack slots could be merged.
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %local1_cast)
+  call void @use(i8* %local1_cast)
+  call void @llvm.lifetime.end.p0i8(i64 1, i8* %local1_cast)
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %local2_cast)
+  call void @use(i8* %local2_cast)
+  call void @llvm.lifetime.end.p0i8(i64 1, i8* %local2_cast)
+  ret void
+
+; CHECK-LABEL: codeRepl:
+; CHECK: [[local1_cast:%.*]] = bitcast i256* %local1 to i8*
+; CHECK: call void @llvm.lifetime.start.p0i8(i64 -1, i8* [[local1_cast]])
+; CHECK: [[local2_cast:%.*]] = bitcast i256* %local2 to i8*
+; CHECK: call void @llvm.lifetime.start.p0i8(i64 -1, i8* [[local2_cast]])
+; CHECK: call i1 @foo.cold.1(i8* %local1_cast, i8* %local2_cast)
+; CHECK: call void @llvm.lifetime.end.p0i8(i64 -1, i8* [[local2_cast]])
+; CHECK: call void @llvm.lifetime.end.p0i8(i64 -1, i8* [[local1_cast]])
+; CHECK: 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)
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %local2_cast)
+  call void @cold_use2(i8* %local1_cast, i8* %local2_cast)
+  call void @llvm.lifetime.end.p0i8(i64 1, i8* %gep1)
+  call void @llvm.lifetime.end.p0i8(i64 1, i8* %local2_cast)
+  br i1 undef, label %outlinedPath2, label %outlinedPathExit
+
+outlinedPath2:
+  ; These extra lifetime markers are used to test that we emit only one
+  ; pair of guard markers in the caller per memory object.
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %local2_cast)
+  call void @use(i8* %local2_cast)
+  call void @llvm.lifetime.end.p0i8(i64 1, i8* %local2_cast)
+  ret void
+
+outlinedPathExit:
+  ret void
+}
+
+; CHECK-LABEL: define {{.*}}@foo.cold.1(
+; CHECK-NOT: @llvm.lifetime