StackColoring for SafeStack.

[eugenis]

Summary

This is a fix for PR27842.

An IR-level implementation of stack coloring tailored to work with
SafeStack. It is a bit weaker than the MI implementation in that it
does not the "lifetime start at first access" logic. This can be
improved in the future.

This patch also replaces the naive implementation of stack frame
layout with a greedy algorithm that can split existing stack slots
and even fit small objects inside the alignment padding of other
objects.

[hiraditya]

lib/CodeGen/SafeStack.cpp

532

DataLayout::getPrefTypeAlignment already returns unsigned.

lib/CodeGen/SafeStackColoring.h

62

It seems you can use DenseMap<const BasicBlock *, std::pair<unsigned, unsigned>>. Also it will be good if you could add some comments about what the numbers in 'std::pair<unsigned, unsigned>' represent; similar to what you have in line 77-78

[eugenis]

ping

[pcc]

lib/CodeGen/SafeStack.cpp

527

Since the GrowsDown parameter is always true, can you remove it?

lib/CodeGen/SafeStackColoring.cpp

230–234

You can remove all braces here.

lib/CodeGen/SafeStackColoring.h

41

Say which instructions.

42–50

Inconsistent capitalization of member functions.

59

Define "interesting"

73

Define "interesting"

lib/CodeGen/SafeStackLayout.cpp

70

This could be a range-for loop.

110

Same here

127

Same here

lib/CodeGen/SafeStackLayout.h

19

Is this the right name for this class? The name suggests that it handles layout for the safe stack.

[eugenis]

lib/CodeGen/SafeStackColoring.h

59

I've added a few words in the class-level comment.

lib/CodeGen/SafeStackLayout.h

19

Would you prefer UnsafeStackLayout?
There is nothing in this class that deals with the notion of safe or unsafe alloca - it's frame layout code for the "SafeStack" feature, and I like that it starts with "SafeStack".

[pcc]

LGTM

lib/CodeGen/SafeStackColoring.h

25

Nitpick: "as sets".

lib/CodeGen/SafeStackLayout.h

20

I'm partial to the pattern of moving pass implementation details into a namespace named after the pass, so maybe something like llvm::safestack::StackLayout?

But I'd be happy with whatever you decide to do here.

[eugenis]

lib/CodeGen/SafeStackLayout.h

20

I like this!
Done.
I've also renamed SafeStackColoring to safestack::StackColoring.

[pcc]

LGTM

[eugenis]

r274162
Thanks for the review!

lib/CodeGen/CMakeLists.txt

add_llvm_library(LLVMCodeGen
   RegisterCoalescer.cpp
   RegisterPressure.cpp
   RegisterScavenging.cpp
   RenameIndependentSubregs.cpp
   RegisterUsageInfo.cpp
   RegUsageInfoCollector.cpp
   RegUsageInfoPropagate.cpp
   SafeStack.cpp
+  SafeStackColoring.cpp
+  SafeStackLayout.cpp
   ScheduleDAG.cpp
   ScheduleDAGInstrs.cpp
   ScheduleDAGPrinter.cpp
   ScoreboardHazardRecognizer.cpp
   ShadowStackGCLowering.cpp
   ShrinkWrap.cpp
   SjLjEHPrepare.cpp
   SlotIndexes.cpp

lib/CodeGen/SafeStack.cpp

 // This pass splits the stack into the safe stack (kept as-is for LLVM backend)
 // and the unsafe stack (explicitly allocated and managed through the runtime
 // support library).
 //
 // http://clang.llvm.org/docs/SafeStack.html
 //
 //===----------------------------------------------------------------------===//
 
+#include "SafeStackColoring.h"
+#include "SafeStackLayout.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DIBuilder.h"

 #include "llvm/Support/raw_os_ostream.h"
 #include "llvm/Target/TargetLowering.h"
 #include "llvm/Target/TargetSubtargetInfo.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 
 using namespace llvm;
+using namespace llvm::safestack;
 
 #define DEBUG_TYPE "safestack"
 
 enum UnsafeStackPtrStorageVal { ThreadLocalUSP, SingleThreadUSP };
 
 static cl::opt<UnsafeStackPtrStorageVal> USPStorage("safe-stack-usp-storage",
     cl::Hidden, cl::init(ThreadLocalUSP),
     cl::desc("Type of storage for the unsafe stack pointer"),
Value *SafeStack::moveStaticAllocasToUnsafeStack(
     IRBuilder<> &IRB, Function &F, ArrayRef<AllocaInst *> StaticAllocas,
     ArrayRef<Argument *> ByValArguments, ArrayRef<ReturnInst *> Returns,
     Instruction *BasePointer, AllocaInst *StackGuardSlot) {
   if (StaticAllocas.empty() && ByValArguments.empty())
     return BasePointer;
 
   DIBuilder DIB(*F.getParent());
 
-  // Compute maximum alignment among static objects on the unsafe stack.
-  unsigned MaxAlignment = 0;
+  StackColoring SSC(F, StaticAllocas);
+  SSC.run();
+  SSC.removeAllMarkers();
+
+  // Unsafe stack always grows down.
+  StackLayout SSL(StackAlignment);
+  if (StackGuardSlot) {
+    Type *Ty = StackGuardSlot->getAllocatedType();
+    unsigned Align =
+        std::max(DL->getPrefTypeAlignment(Ty), StackGuardSlot->getAlignment());
+    SSL.addObject(StackGuardSlot, getStaticAllocaAllocationSize(StackGuardSlot),
+                  Align, SSC.getLiveRange(StackGuardSlot));
+  }
+
   for (Argument *Arg : ByValArguments) {
     Type *Ty = Arg->getType()->getPointerElementType();
+    uint64_t Size = DL->getTypeStoreSize(Ty);
+    if (Size == 0)
+      Size = 1; // Don't create zero-sized stack objects.
+
+    // Ensure the object is properly aligned.
     unsigned Align = std::max((unsigned)DL->getPrefTypeAlignment(Ty),
                               Arg->getParamAlignment());
-    if (Align > MaxAlignment)
-      MaxAlignment = Align;
+    SSL.addObject(Arg, Size, Align, SSC.getFullLiveRange());
   }
+
   for (AllocaInst *AI : StaticAllocas) {
     Type *Ty = AI->getAllocatedType();
+    uint64_t Size = getStaticAllocaAllocationSize(AI);
+    if (Size == 0)
+      Size = 1; // Don't create zero-sized stack objects.
+
+    // Ensure the object is properly aligned.
     unsigned Align =
         std::max((unsigned)DL->getPrefTypeAlignment(Ty), AI->getAlignment());
-    if (Align > MaxAlignment)
-      MaxAlignment = Align;
+
+    SSL.addObject(AI, Size, Align, SSC.getLiveRange(AI));
   }
 
-  if (MaxAlignment > StackAlignment) {
+  SSL.computeLayout();
+  unsigned FrameAlignment = SSL.getFrameAlignment();
+
+  // FIXME: tell SSL that we start at a less-then-MaxAlignment aligned location
+  // (AlignmentSkew).
+  if (FrameAlignment > StackAlignment) {
     // Re-align the base pointer according to the max requested alignment.
-    assert(isPowerOf2_32(MaxAlignment));
+    assert(isPowerOf2_32(FrameAlignment));
     IRB.SetInsertPoint(BasePointer->getNextNode());
     BasePointer = cast<Instruction>(IRB.CreateIntToPtr(
         IRB.CreateAnd(IRB.CreatePtrToInt(BasePointer, IntPtrTy),
-                      ConstantInt::get(IntPtrTy, ~uint64_t(MaxAlignment - 1))),
+                      ConstantInt::get(IntPtrTy, ~uint64_t(FrameAlignment - 1))),
         StackPtrTy));
   }
 
-  int64_t StaticOffset = 0; // Current stack top.
   IRB.SetInsertPoint(BasePointer->getNextNode());
 
   if (StackGuardSlot) {
-    StaticOffset += getStaticAllocaAllocationSize(StackGuardSlot);
+    unsigned Offset = SSL.getObjectOffset(StackGuardSlot);
     Value *Off = IRB.CreateGEP(BasePointer, // BasePointer is i8*
-                               ConstantInt::get(Int32Ty, -StaticOffset));
+                               ConstantInt::get(Int32Ty, -Offset));
     Value *NewAI =
         IRB.CreateBitCast(Off, StackGuardSlot->getType(), "StackGuardSlot");
 
     // Replace alloc with the new location.
     StackGuardSlot->replaceAllUsesWith(NewAI);
     StackGuardSlot->eraseFromParent();
   }
 
   for (Argument *Arg : ByValArguments) {
+    unsigned Offset = SSL.getObjectOffset(Arg);
     Type *Ty = Arg->getType()->getPointerElementType();
 
     uint64_t Size = DL->getTypeStoreSize(Ty);
     if (Size == 0)
       Size = 1; // Don't create zero-sized stack objects.
 
-    // Ensure the object is properly aligned.
-    unsigned Align = std::max((unsigned)DL->getPrefTypeAlignment(Ty),
-                              Arg->getParamAlignment());
-
-    // Add alignment.
-    // NOTE: we ensure that BasePointer itself is aligned to >= Align.
-    StaticOffset += Size;
-    StaticOffset = alignTo(StaticOffset, Align);
-
     Value *Off = IRB.CreateGEP(BasePointer, // BasePointer is i8*
-                               ConstantInt::get(Int32Ty, -StaticOffset));
+                               ConstantInt::get(Int32Ty, -Offset));
     Value *NewArg = IRB.CreateBitCast(Off, Arg->getType(),
                                      Arg->getName() + ".unsafe-byval");
 
     // Replace alloc with the new location.
     replaceDbgDeclare(Arg, BasePointer, BasePointer->getNextNode(), DIB,
-                      /*Deref=*/true, -StaticOffset);
+                      /*Deref=*/true, -Offset);
     Arg->replaceAllUsesWith(NewArg);
     IRB.SetInsertPoint(cast<Instruction>(NewArg)->getNextNode());
     IRB.CreateMemCpy(Off, Arg, Size, Arg->getParamAlignment());
   }
 
   // Allocate space for every unsafe static AllocaInst on the unsafe stack.
   for (AllocaInst *AI : StaticAllocas) {
     IRB.SetInsertPoint(AI);
+    unsigned Offset = SSL.getObjectOffset(AI);
 
-    Type *Ty = AI->getAllocatedType();
     uint64_t Size = getStaticAllocaAllocationSize(AI);
     if (Size == 0)
       Size = 1; // Don't create zero-sized stack objects.
 
-    // Ensure the object is properly aligned.
-    unsigned Align =
-        std::max((unsigned)DL->getPrefTypeAlignment(Ty), AI->getAlignment());
-
-    // Add alignment.
-    // NOTE: we ensure that BasePointer itself is aligned to >= Align.
-    StaticOffset += Size;
-    StaticOffset = alignTo(StaticOffset, Align);
-
-    replaceDbgDeclareForAlloca(AI, BasePointer, DIB, /*Deref=*/true, -StaticOffset);
-    replaceDbgValueForAlloca(AI, BasePointer, DIB, -StaticOffset);
+    replaceDbgDeclareForAlloca(AI, BasePointer, DIB, /*Deref=*/true, -Offset);
+    replaceDbgValueForAlloca(AI, BasePointer, DIB, -Offset);
 
     // Replace uses of the alloca with the new location.
     // Insert address calculation close to each use to work around PR27844.
     std::string Name = std::string(AI->getName()) + ".unsafe";
     while (!AI->use_empty()) {
       Use &U = *AI->use_begin();
       Instruction *User = cast<Instruction>(U.getUser());
 
       Instruction *InsertBefore;
       if (auto *PHI = dyn_cast<PHINode>(User))
         InsertBefore = PHI->getIncomingBlock(U)->getTerminator();
       else
         InsertBefore = User;
 
       IRBuilder<> IRBUser(InsertBefore);
       Value *Off = IRBUser.CreateGEP(BasePointer, // BasePointer is i8*
-                                     ConstantInt::get(Int32Ty, -StaticOffset));
+                                     ConstantInt::get(Int32Ty, -Offset));
       Value *Replacement = IRBUser.CreateBitCast(Off, AI->getType(), Name);
 
       if (auto *PHI = dyn_cast<PHINode>(User)) {
         // PHI nodes may have multiple incoming edges from the same BB (why??),
         // all must be updated at once with the same incoming value.
         auto *BB = PHI->getIncomingBlock(U);
         for (unsigned I = 0; I < PHI->getNumIncomingValues(); ++I)
           if (PHI->getIncomingBlock(I) == BB)
             PHI->setIncomingValue(I, Replacement);
       } else {
         U.set(Replacement);
       }
     }
 
     AI->eraseFromParent();
   }
 
   // Re-align BasePointer so that our callees would see it aligned as
   // expected.
   // FIXME: no need to update BasePointer in leaf functions.
-  StaticOffset = alignTo(StaticOffset, StackAlignment);
+  unsigned FrameSize = alignTo(SSL.getFrameSize(), StackAlignment);
 
   // Update shadow stack pointer in the function epilogue.
   IRB.SetInsertPoint(BasePointer->getNextNode());
 
   Value *StaticTop =
-      IRB.CreateGEP(BasePointer, ConstantInt::get(Int32Ty, -StaticOffset),
+      IRB.CreateGEP(BasePointer, ConstantInt::get(Int32Ty, -FrameSize),
                     "unsafe_stack_static_top");
   IRB.CreateStore(StaticTop, UnsafeStackPtr);
   return StaticTop;
 }
 
 void SafeStack::moveDynamicAllocasToUnsafeStack(
     Function &F, Value *UnsafeStackPtr, AllocaInst *DynamicTop,
     ArrayRef<AllocaInst *> DynamicAllocas) {

lib/CodeGen/SafeStackColoring.h

+//===-- SafeStackColoring.h - SafeStack frame coloring ---------*- C++ -*--===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_CODEGEN_SAFESTACKCOLORING_H
+#define LLVM_LIB_CODEGEN_SAFESTACKCOLORING_H
+
+#include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/IR/Function.h"
+#include "llvm/Support/raw_os_ostream.h"
+
+namespace llvm {
+class AllocaInst;
+
+namespace safestack {
+/// Compute live ranges of allocas.
+/// Live ranges are represented as sets of "interesting" instructions, which are
+/// defined as instructions that may start or end an alloca's lifetime. These
+/// are:
+/// * lifetime.start and lifetime.end intrinsics
+/// * first instruction of any basic block
+/// Interesting instructions are numbered in the depth-first walk of the CFG,
+/// and in the program order inside each basic block.
+class StackColoring {
+  /// A class representing liveness information for a single basic block.
+  /// Each bit in the BitVector represents the liveness property
+  /// for a different stack slot.
+  struct BlockLifetimeInfo {
+    /// Which slots BEGINs in each basic block.
+    BitVector Begin;
+    /// Which slots ENDs in each basic block.
+    BitVector End;
+    /// Which slots are marked as LIVE_IN, coming into each basic block.
+    BitVector LiveIn;
+    /// Which slots are marked as LIVE_OUT, coming out of each basic block.
+    BitVector LiveOut;
+  };
+
+public:
+  /// This class represents a set of interesting instructions where an alloca is
+  /// live.
+  struct LiveRange {
+    BitVector bv;
+    void SetMaximum(int size) { bv.resize(size); }
+    void AddRange(unsigned start, unsigned end) { bv.set(start, end); }
+    bool Overlaps(const LiveRange &Other) const {
+      return bv.anyCommon(Other.bv);
+    }
+    void Join(const LiveRange &Other) { bv |= Other.bv; }
+  };
+
+private:
+  Function &F;
+
+  /// Maps active slots (per bit) for each basic block.
+  typedef DenseMap<BasicBlock *, BlockLifetimeInfo> LivenessMap;
+  LivenessMap BlockLiveness;
+
+  /// Number of interesting instructions.
+  int NumInst;
+  /// Numeric ids for interesting instructions.
+  DenseMap<Instruction *, unsigned> InstructionNumbering;
+  /// A range [Start, End) of instruction ids for each basic block.
+  /// Instructions inside each BB have monotonic and consecutive ids.
+  DenseMap<const BasicBlock *, std::pair<unsigned, unsigned>> BlockInstRange;
+
+  ArrayRef<AllocaInst *> Allocas;
+  unsigned NumAllocas;
+  DenseMap<AllocaInst *, unsigned> AllocaNumbering;
+  /// LiveRange for allocas.
+  SmallVector<LiveRange, 8> LiveRanges;
+
+  /// The set of allocas that have at least one lifetime.start. All other
+  /// allocas get LiveRange that corresponds to the entire function.
+  BitVector InterestingAllocas;
+  SmallVector<Instruction *, 8> Markers;
+
+  struct Marker {
+    unsigned AllocaNo;
+    bool IsStart;
+  };
+
+  /// List of {InstNo, {AllocaNo, IsStart}} for each BB, ordered by InstNo.
+  DenseMap<BasicBlock *, SmallVector<std::pair<unsigned, Marker>, 4>> BBMarkers;
+
+  void dumpAllocas();
+  void dumpBlockLiveness();
+  void dumpLiveRanges();
+
+  bool readMarker(Instruction *I, bool *IsStart);
+  void collectMarkers();
+  void calculateLocalLiveness();
+  void calculateLiveIntervals();
+
+public:
+  StackColoring(Function &F, ArrayRef<AllocaInst *> Allocas)
+      : F(F), NumInst(-1), Allocas(Allocas), NumAllocas(Allocas.size()) {}
+
+  void run();
+  void removeAllMarkers();
+
+  /// Returns a set of "interesting" instructions where the given alloca is
+  /// live. Not all instructions in a function are interesting: we pick a set
+  /// that is large enough for LiveRange::Overlaps to be correct.
+  const LiveRange &getLiveRange(AllocaInst *AI);
+
+  /// Returns a live range that represents an alloca that is live throughout the
+  /// entire function.
+  LiveRange getFullLiveRange() {
+    assert(NumInst >= 0);
+    LiveRange R;
+    R.SetMaximum(NumInst);
+    R.AddRange(0, NumInst);
+    return R;
+  }
+};
+
+static inline raw_ostream &operator<<(raw_ostream &OS, const BitVector &V) {
+  OS << "{";
+  int idx = V.find_first();
+  bool first = true;
+  while (idx >= 0) {
+    if (!first) {
+      OS << ", ";
+    }
+    first = false;
+    OS << idx;
+    idx = V.find_next(idx);
+  }
+  OS << "}";
+  return OS;
+}
+
+static inline raw_ostream &operator<<(raw_ostream &OS,
+                                      const StackColoring::LiveRange &R) {
+  return OS << R.bv;
+}
+
+} // namespace safestack
+} // namespace llvm
+
+#endif // LLVM_LIB_CODEGEN_SAFESTACKCOLORING_H

lib/CodeGen/SafeStackColoring.cpp

+//===-- SafeStackColoring.cpp - SafeStack frame coloring -------*- C++ -*--===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "SafeStackColoring.h"
+
+#include "llvm/ADT/DepthFirstIterator.h"
+#include "llvm/IR/CFG.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/Support/Debug.h"
+
+using namespace llvm;
+using namespace llvm::safestack;
+
+#define DEBUG_TYPE "safestackcoloring"
+
+static cl::opt<bool> ClColoring("safe-stack-coloring",
+                                cl::desc("enable safe stack coloring"),
+                                cl::Hidden, cl::init(true));
+
+const StackColoring::LiveRange &StackColoring::getLiveRange(AllocaInst *AI) {
+  return LiveRanges[AllocaNumbering[AI]];
+}
+
+bool StackColoring::readMarker(Instruction *I, bool *IsStart) {
+  auto *II = dyn_cast<IntrinsicInst>(I);
+  if (!II || (II->getIntrinsicID() != Intrinsic::lifetime_start &&
+              II->getIntrinsicID() != Intrinsic::lifetime_end))
+    return false;
+
+  *IsStart = II->getIntrinsicID() == Intrinsic::lifetime_start;
+  return true;
+}
+
+void StackColoring::removeAllMarkers() {
+  for (auto *I : Markers) {
+    auto *Op = dyn_cast<Instruction>(I->getOperand(1));
+    I->eraseFromParent();
+    // Remove the operand bitcast, too, if it has no more uses left.
+    if (Op && Op->use_empty())
+      Op->eraseFromParent();
+  }
+}
+
+void StackColoring::collectMarkers() {
+  InterestingAllocas.resize(NumAllocas);
+  DenseMap<BasicBlock *, SmallDenseMap<Instruction *, Marker>> BBMarkerSet;
+
+  // Compute the set of start/end markers per basic block.
+  for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo) {
+    AllocaInst *AI = Allocas[AllocaNo];
+    SmallVector<Instruction *, 8> WorkList;
+    WorkList.push_back(AI);
+    while (!WorkList.empty()) {
+      Instruction *I = WorkList.pop_back_val();
+      for (User *U : I->users()) {
+        if (auto *BI = dyn_cast<BitCastInst>(U)) {
+          WorkList.push_back(BI);
+          continue;
+        }
+        auto *UI = dyn_cast<Instruction>(U);
+        if (!UI)
+          continue;
+        bool IsStart;
+        if (!readMarker(UI, &IsStart))
+          continue;
+        if (IsStart)
+          InterestingAllocas.set(AllocaNo);
+        BBMarkerSet[UI->getParent()][UI] = {AllocaNo, IsStart};
+        Markers.push_back(UI);
+      }
+    }
+  }
+
+  // Compute instruction numbering. Only the following instructions are
+  // considered:
+  // * Basic block entries
+  // * Lifetime markers
+  // For each basic block, compute
+  // * the list of markers in the instruction order
+  // * the sets of allocas whose lifetime starts or ends in this BB
+  DEBUG(dbgs() << "Instructions:\n");
+  unsigned InstNo = 0;
+  for (BasicBlock *BB : depth_first(&F)) {
+    DEBUG(dbgs() << "  " << InstNo << ": BB " << BB->getName() << "\n");
+    unsigned BBStart = InstNo++;
+
+    BlockLifetimeInfo &BlockInfo = BlockLiveness[BB];
+    BlockInfo.Begin.resize(NumAllocas);
+    BlockInfo.End.resize(NumAllocas);
+    BlockInfo.LiveIn.resize(NumAllocas);
+    BlockInfo.LiveOut.resize(NumAllocas);
+
+    auto &BlockMarkerSet = BBMarkerSet[BB];
+    if (BlockMarkerSet.empty()) {
+      unsigned BBEnd = InstNo;
+      BlockInstRange[BB] = std::make_pair(BBStart, BBEnd);
+      continue;
+    }
+
+    auto ProcessMarker = [&](Instruction *I, const Marker &M) {
+      DEBUG(dbgs() << "  " << InstNo << ":  "
+                   << (M.IsStart ? "start " : "end   ") << M.AllocaNo << ", "
+                   << *I << "\n");
+
+      BBMarkers[BB].push_back({InstNo, M});
+
+      InstructionNumbering[I] = InstNo++;
+
+      if (M.IsStart) {
+        if (BlockInfo.End.test(M.AllocaNo))
+          BlockInfo.End.reset(M.AllocaNo);
+        BlockInfo.Begin.set(M.AllocaNo);
+      } else {
+        if (BlockInfo.Begin.test(M.AllocaNo))
+          BlockInfo.Begin.reset(M.AllocaNo);
+        BlockInfo.End.set(M.AllocaNo);
+      }
+    };
+
+    if (BlockMarkerSet.size() == 1) {
+      ProcessMarker(BlockMarkerSet.begin()->getFirst(),
+                    BlockMarkerSet.begin()->getSecond());
+    } else {
+      // Scan the BB to determine the marker order.
+      for (Instruction &I : *BB) {
+        auto It = BlockMarkerSet.find(&I);
+        if (It == BlockMarkerSet.end())
+          continue;
+        ProcessMarker(&I, It->getSecond());
+      }
+    }
+
+    unsigned BBEnd = InstNo;
+    BlockInstRange[BB] = std::make_pair(BBStart, BBEnd);
+  }
+  NumInst = InstNo;
+}
+
+void StackColoring::calculateLocalLiveness() {
+  bool changed = true;
+  while (changed) {
+    changed = false;
+
+    for (BasicBlock *BB : depth_first(&F)) {
+      BlockLifetimeInfo &BlockInfo = BlockLiveness[BB];
+
+      // Compute LiveIn by unioning together the LiveOut sets of all preds.
+      BitVector LocalLiveIn;
+      for (auto *PredBB : predecessors(BB)) {
+        LivenessMap::const_iterator I = BlockLiveness.find(PredBB);
+        assert(I != BlockLiveness.end() && "Predecessor not found");
+        LocalLiveIn |= I->second.LiveOut;
+      }
+
+      // Compute LiveOut by subtracting out lifetimes that end in this
+      // block, then adding in lifetimes that begin in this block.  If
+      // we have both BEGIN and END markers in the same basic block
+      // then we know that the BEGIN marker comes after the END,
+      // because we already handle the case where the BEGIN comes
+      // before the END when collecting the markers (and building the
+      // BEGIN/END vectors).
+      BitVector LocalLiveOut = LocalLiveIn;
+      LocalLiveOut.reset(BlockInfo.End);
+      LocalLiveOut |= BlockInfo.Begin;
+
+      // Update block LiveIn set, noting whether it has changed.
+      if (LocalLiveIn.test(BlockInfo.LiveIn)) {
+        changed = true;
+        BlockInfo.LiveIn |= LocalLiveIn;
+      }
+
+      // Update block LiveOut set, noting whether it has changed.
+      if (LocalLiveOut.test(BlockInfo.LiveOut)) {
+        changed = true;
+        BlockInfo.LiveOut |= LocalLiveOut;
+      }
+    }
+  } // while changed.
+}
+
+void StackColoring::calculateLiveIntervals() {
+  for (auto IT : BlockLiveness) {
+    BasicBlock *BB = IT.getFirst();
+    BlockLifetimeInfo &BlockInfo = IT.getSecond();
+    unsigned BBStart, BBEnd;
+    std::tie(BBStart, BBEnd) = BlockInstRange[BB];
+
+    BitVector Started, Ended;
+    Started.resize(NumAllocas);
+    Ended.resize(NumAllocas);
+    SmallVector<unsigned, 8> Start;
+    Start.resize(NumAllocas);
+
+    // LiveIn ranges start at the first instruction.
+    for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo) {
+      if (BlockInfo.LiveIn.test(AllocaNo)) {
+        Started.set(AllocaNo);
+        Start[AllocaNo] = BBStart;
+      }
+    }
+
+    for (auto &It : BBMarkers[BB]) {
+      unsigned InstNo = It.first;
+      bool IsStart = It.second.IsStart;
+      unsigned AllocaNo = It.second.AllocaNo;
+
+      if (IsStart) {
+        assert(!Started.test(AllocaNo));
+        Started.set(AllocaNo);
+        Ended.reset(AllocaNo);
+        Start[AllocaNo] = InstNo;
+      } else {
+        assert(!Ended.test(AllocaNo));
+        if (Started.test(AllocaNo)) {
+          LiveRanges[AllocaNo].AddRange(Start[AllocaNo], InstNo);
+          Started.reset(AllocaNo);
+        }
+        Ended.set(AllocaNo);
+      }
+    }
+
+    for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo)
+      if (Started.test(AllocaNo))
+        LiveRanges[AllocaNo].AddRange(Start[AllocaNo], BBEnd);
+  }
+}
+
+LLVM_DUMP_METHOD void StackColoring::dumpAllocas() {
+  dbgs() << "Allocas:\n";
+  for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo)
+    dbgs() << "  " << AllocaNo << ": " << *Allocas[AllocaNo] << "\n";
+}
+
+LLVM_DUMP_METHOD void StackColoring::dumpBlockLiveness() {
+  dbgs() << "Block liveness:\n";
+  for (auto IT : BlockLiveness) {
+    BasicBlock *BB = IT.getFirst();
+    BlockLifetimeInfo &BlockInfo = BlockLiveness[BB];
+    auto BlockRange = BlockInstRange[BB];
+    dbgs() << "  BB [" << BlockRange.first << ", " << BlockRange.second
+           << "): begin " << BlockInfo.Begin << ", end " << BlockInfo.End
+           << ", livein " << BlockInfo.LiveIn << ", liveout "
+           << BlockInfo.LiveOut << "\n";
+  }
+}
+
+LLVM_DUMP_METHOD void StackColoring::dumpLiveRanges() {
+  dbgs() << "Alloca liveness:\n";
+  for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo) {
+    LiveRange &Range = LiveRanges[AllocaNo];
+    dbgs() << "  " << AllocaNo << ": " << Range << "\n";
+  }
+}
+
+void StackColoring::run() {
+  DEBUG(dumpAllocas());
+
+  for (unsigned I = 0; I < NumAllocas; ++I)
+    AllocaNumbering[Allocas[I]] = I;
+  LiveRanges.resize(NumAllocas);
+
+  collectMarkers();
+
+  if (!ClColoring) {
+    for (auto &R : LiveRanges) {
+      R.SetMaximum(1);
+      R.AddRange(0, 1);
+    }
+    return;
+  }
+
+  for (auto &R : LiveRanges)
+    R.SetMaximum(NumInst);
+  for (unsigned I = 0; I < NumAllocas; ++I)
+    if (!InterestingAllocas.test(I))
+      LiveRanges[I] = getFullLiveRange();
+
+  calculateLocalLiveness();
+  DEBUG(dumpBlockLiveness());
+  calculateLiveIntervals();
+  DEBUG(dumpLiveRanges());
+}

lib/CodeGen/SafeStackLayout.h

+//===-- SafeStackLayout.h - SafeStack frame layout -------------*- C++ -*--===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_CODEGEN_SAFESTACKLAYOUT_H
+#define LLVM_LIB_CODEGEN_SAFESTACKLAYOUT_H
+
+#include "SafeStackColoring.h"
+
+namespace llvm {
+namespace safestack {
+
+/// Compute the layout of an unsafe stack frame.
+class StackLayout {
+  unsigned MaxAlignment;
+
+  struct StackRegion {
+    unsigned Start;
+    unsigned End;
+    StackColoring::LiveRange Range;
+    StackRegion(unsigned Start, unsigned End,
+                const StackColoring::LiveRange &Range)
+        : Start(Start), End(End), Range(Range) {}
+  };
+  /// The list of current stack regions, sorted by StackRegion::Start.
+  SmallVector<StackRegion, 16> Regions;
+
+  struct StackObject {
+    const Value *Handle;
+    unsigned Size, Alignment;
+    StackColoring::LiveRange Range;
+  };
+  SmallVector<StackObject, 8> StackObjects;
+
+  DenseMap<const Value *, unsigned> ObjectOffsets;
+
+  void layoutObject(StackObject &Obj);
+
+public:
+  StackLayout(unsigned StackAlignment) : MaxAlignment(StackAlignment) {}
+  /// Add an object to the stack frame. Value pointer is opaque and used as a
+  /// handle to retrieve the object's offset in the frame later.
+  void addObject(const Value *V, unsigned Size, unsigned Alignment,
+                 const StackColoring::LiveRange &Range);
+
+  /// Run the layout computation for all previously added objects.
+  void computeLayout();
+
+  /// Returns the offset to the object start in the stack frame.
+  unsigned getObjectOffset(const Value *V) { return ObjectOffsets[V]; }
+
+  /// Returns the size of the entire frame.
+  unsigned getFrameSize() { return Regions.empty() ? 0 : Regions.back().End; }
+
+  /// Returns the alignment of the frame.
+  unsigned getFrameAlignment() { return MaxAlignment; }
+  void print(raw_ostream &OS);
+};
+
+} // namespace safestack
+} // namespace llvm
+
+#endif // LLVM_LIB_CODEGEN_SAFESTACKLAYOUT_H

lib/CodeGen/SafeStackLayout.cpp

+//===-- SafeStackLayout.cpp - SafeStack frame layout -----------*- C++ -*--===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "SafeStackLayout.h"
+
+#include "llvm/IR/Instructions.h"
+#include "llvm/Support/Debug.h"
+
+using namespace llvm;
+using namespace llvm::safestack;
+
+#define DEBUG_TYPE "safestacklayout"
+
+static cl::opt<bool> ClLayout("safe-stack-layout",
+                              cl::desc("enable safe stack layout"), cl::Hidden,
+                              cl::init(true));
+
+LLVM_DUMP_METHOD void StackLayout::print(raw_ostream &OS) {
+  OS << "Stack regions:\n";
+  for (unsigned i = 0; i < Regions.size(); ++i) {
+    OS << "  " << i << ": [" << Regions[i].Start << ", " << Regions[i].End
+       << "), range " << Regions[i].Range << "\n";
+  }
+  OS << "Stack objects:\n";
+  for (auto &IT : ObjectOffsets) {
+    OS << "  at " << IT.getSecond() << ": " << *IT.getFirst() << "\n";
+  }
+}
+
+void StackLayout::addObject(const Value *V, unsigned Size, unsigned Alignment,
+                            const StackColoring::LiveRange &Range) {
+  StackObjects.push_back({V, Size, Alignment, Range});
+  MaxAlignment = std::max(MaxAlignment, Alignment);
+}
+
+static unsigned AdjustStackOffset(unsigned Offset, unsigned Size,
+                                  unsigned Alignment) {
+  return alignTo(Offset + Size, Alignment) - Size;
+}
+
+void StackLayout::layoutObject(StackObject &Obj) {
+  if (!ClLayout) {
+    // If layout is disabled, just grab the next aligned address.
+    // This effectively disables stack coloring as well.
+    unsigned LastRegionEnd = Regions.empty() ? 0 : Regions.back().End;
+    unsigned Start = AdjustStackOffset(LastRegionEnd, Obj.Size, Obj.Alignment);
+    unsigned End = Start + Obj.Size;
+    Regions.emplace_back(Start, End, Obj.Range);
+    ObjectOffsets[Obj.Handle] = End;
+    return;
+  }
+
+  DEBUG(dbgs() << "Layout: size " << Obj.Size << ", align " << Obj.Alignment
+               << ", range " << Obj.Range << "\n");
+  assert(Obj.Alignment <= MaxAlignment);
+  unsigned Start = AdjustStackOffset(0, Obj.Size, Obj.Alignment);
+  unsigned End = Start + Obj.Size;
+  DEBUG(dbgs() << "  First candidate: " << Start << " .. " << End << "\n");
+  for (const StackRegion &R : Regions) {
+    DEBUG(dbgs() << "  Examining region: " << R.Start << " .. " << R.End
+                 << ", range " << R.Range << "\n");
+    assert(End >= R.Start);
+    if (Start >= R.End) {
+      DEBUG(dbgs() << "  Does not intersect, skip.\n");
+      continue;
+    }
+    if (Obj.Range.Overlaps(R.Range)) {
+      // Find the next appropriate location.
+      Start = AdjustStackOffset(R.End, Obj.Size, Obj.Alignment);
+      End = Start + Obj.Size;
+      DEBUG(dbgs() << "  Overlaps. Next candidate: " << Start << " .. " << End
+                   << "\n");
+      continue;
+    }
+    if (End <= R.End) {
+      DEBUG(dbgs() << "  Reusing region(s).\n");
+      break;
+    }
+  }
+
+  unsigned LastRegionEnd = Regions.empty() ? 0 : Regions.back().End;
+  if (End > LastRegionEnd) {
+    // Insert a new region at the end. Maybe two.
+    if (Start > LastRegionEnd) {
+      DEBUG(dbgs() << "  Creating gap region: " << LastRegionEnd << " .. "
+                   << Start << "\n");
+      Regions.emplace_back(LastRegionEnd, Start, StackColoring::LiveRange());
+      LastRegionEnd = Start;
+    }
+    DEBUG(dbgs() << "  Creating new region: " << LastRegionEnd << " .. " << End
+                 << ", range " << Obj.Range << "\n");
+    Regions.emplace_back(LastRegionEnd, End, Obj.Range);
+    LastRegionEnd = End;
+  }
+
+  // Split starting and ending regions if necessary.
+  for (StackRegion &R : Regions) {
+    if (Start > R.Start && Start < R.End) {
+      StackRegion R0 = R;
+      R.Start = R0.End = Start;
+      Regions.insert(&R, R0);
+      continue;
+    }
+    if (End > R.Start && End < R.End) {
+      StackRegion R0 = R;
+      R0.End = R.Start = End;
+      Regions.insert(&R, R0);
+      break;
+    }
+  }
+
+  // Update live ranges for all affected regions.
+  for (StackRegion &R : Regions) {
+    if (Start < R.End && End > R.Start)
+      R.Range.Join(Obj.Range);
+    if (End <= R.End)
+      break;
+  }
+
+  ObjectOffsets[Obj.Handle] = End;
+}
+
+void StackLayout::computeLayout() {
+  // Simple greedy algorithm.
+  // If this is replaced with something smarter, it must preserve the property
+  // that the first object is always at the offset 0 in the stack frame (for
+  // StackProtectorSlot), or handle stack protector in some other way.
+  for (auto &Obj : StackObjects)
+    layoutObject(Obj);
+
+  DEBUG(print(dbgs()));
+}

test/Transforms/SafeStack/coloring.ll

+; RUN: opt -safe-stack -S -mtriple=i386-pc-linux-gnu < %s -o - | FileCheck %s
+; RUN: opt -safe-stack -S -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck %s
+
+define void @f() safestack {
+entry:
+; CHECK:  %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+; CHECK:  %[[USST:.*]] = getelementptr i8, i8* %[[USP]], i32 -16
+
+  %x = alloca i32, align 4
+  %x1 = alloca i32, align 4
+  %x2 = alloca i32, align 4
+  %0 = bitcast i32* %x to i8*
+  call void @llvm.lifetime.start(i64 4, i8* %0)
+
+; CHECK:  %[[A1:.*]] = getelementptr i8, i8* %[[USP]], i32 -4
+; CHECK:  %[[A2:.*]] = bitcast i8* %[[A1]] to i32*
+; CHECK:  call void @capture(i32* nonnull %[[A2]])
+
+  call void @capture(i32* nonnull %x)
+  call void @llvm.lifetime.end(i64 4, i8* %0)
+  %1 = bitcast i32* %x1 to i8*
+  call void @llvm.lifetime.start(i64 4, i8* %1)
+
+; CHECK:  %[[B1:.*]] = getelementptr i8, i8* %[[USP]], i32 -4
+; CHECK:  %[[B2:.*]] = bitcast i8* %[[B1]] to i32*
+; CHECK:  call void @capture(i32* nonnull %[[B2]])
+
+  call void @capture(i32* nonnull %x1)
+  call void @llvm.lifetime.end(i64 4, i8* %1)
+  %2 = bitcast i32* %x2 to i8*
+  call void @llvm.lifetime.start(i64 4, i8* %2)
+
+; CHECK:  %[[C1:.*]] = getelementptr i8, i8* %[[USP]], i32 -4
+; CHECK:  %[[C2:.*]] = bitcast i8* %[[C1]] to i32*
+; CHECK:  call void @capture(i32* nonnull %[[C2]])
+
+  call void @capture(i32* nonnull %x2)
+  call void @llvm.lifetime.end(i64 4, i8* %2)
+  ret void
+}
+
+declare void @llvm.lifetime.start(i64, i8* nocapture)
+declare void @llvm.lifetime.end(i64, i8* nocapture)
+declare void @capture(i32*)

test/Transforms/SafeStack/coloring2.ll

+; RUN: opt -safe-stack -S -mtriple=i386-pc-linux-gnu < %s -o - | FileCheck %s
+; RUN: opt -safe-stack -S -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck %s
+
+; x and y share the stack slot.
+define void @f() safestack {
+; CHECK-LABEL: define void @f
+entry:
+; CHECK:  %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -16
+
+  %x = alloca i32, align 4
+  %y = alloca i32, align 4
+  %z = alloca i32, align 4
+  %x0 = bitcast i32* %x to i8*
+  %y0 = bitcast i32* %y to i8*
+  %z0 = bitcast i32* %z to i8*
+
+  call void @llvm.lifetime.start(i64 -1, i8* %z0)
+  call void @llvm.lifetime.start(i64 -1, i8* %x0)
+
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -4
+  call void @capture32(i32* %x)
+  call void @llvm.lifetime.end(i64 -1, i8* %x0)
+  call void @llvm.lifetime.start(i64 -1, i8* %y0)
+
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -4
+  call void @capture32(i32* %y)
+  call void @llvm.lifetime.end(i64 -1, i8* %y0)
+
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -8
+  call void @capture32(i32* %z)
+  call void @llvm.lifetime.end(i64 -1, i8* %z0)
+
+  ret void
+}
+
+define void @no_markers() safestack {
+; CHECK-LABEL: define void @no_markers(
+entry:
+; CHECK:  %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -16
+
+  %x = alloca i32, align 4
+  %y = alloca i32, align 4
+  %x0 = bitcast i32* %x to i8*
+
+  call void @llvm.lifetime.start(i64 -1, i8* %x0)
+
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -4
+  call void @capture32(i32* %x)
+  call void @llvm.lifetime.end(i64 -1, i8* %x0)
+
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -8
+  call void @capture32(i32* %y)
+
+  ret void
+}
+
+; x and y can't share memory, but they can split z's storage.
+define void @g() safestack {
+; CHECK-LABEL: define void @g
+entry:
+; CHECK:  %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -16
+
+  %x = alloca i32, align 4
+  %y = alloca i32, align 4
+  %z = alloca i64, align 4
+  %x0 = bitcast i32* %x to i8*
+  %y0 = bitcast i32* %y to i8*
+  %z0 = bitcast i64* %z to i8*
+
+  call void @llvm.lifetime.start(i64 -1, i8* %x0)
+  call void @llvm.lifetime.start(i64 -1, i8* %y0)
+
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -4
+  call void @capture32(i32* %x)
+  call void @llvm.lifetime.end(i64 -1, i8* %x0)
+
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -8
+  call void @capture32(i32* %y)
+  call void @llvm.lifetime.end(i64 -1, i8* %y0)
+  call void @llvm.lifetime.start(i64 -1, i8* %z0)
+
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -8
+  call void @capture64(i64* %z)
+  call void @llvm.lifetime.end(i64 -1, i8* %z0)
+
+  ret void
+}
+
+; Both y and z fit in x's alignment gap.
+define void @h() safestack {
+; CHECK-LABEL: define void @h
+entry:
+; CHECK:  %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -16
+
+  %x = alloca i32, align 16
+  %z = alloca i64, align 4
+  %y = alloca i32, align 4
+  %x0 = bitcast i32* %x to i8*
+  %y0 = bitcast i32* %y to i8*
+  %z0 = bitcast i64* %z to i8*
+
+  call void @llvm.lifetime.start(i64 -1, i8* %x0)
+  call void @llvm.lifetime.start(i64 -1, i8* %y0)
+  call void @llvm.lifetime.start(i64 -1, i8* %z0)
+
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -16
+  call void @capture32(i32* %x)
+
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -12
+  call void @capture32(i32* %y)
+
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -8
+  call void @capture64(i64* %z)
+
+  call void @llvm.lifetime.end(i64 -1, i8* %x0)
+  call void @llvm.lifetime.end(i64 -1, i8* %y0)
+  call void @llvm.lifetime.end(i64 -1, i8* %z0)
+
+  ret void
+}
+
+; void f(bool a, bool b) {
+;   long x1, x2; capture64(&x1); capture64(&x2);
+;   if (a) {
+;     long y; capture64(&y);
+;     if (b) {
+;       long y1; capture64(&y1);
+;     } else {
+;       long y2; capture64(&y2);
+;     }
+;   } else {
+;     long z; capture64(&z);
+;     if (b) {
+;       long z1; capture64(&z1);
+;     } else {
+;       long z2; capture64(&z2);
+;     }
+;   }
+; }
+; Everything fits in 4 x 64-bit slots.
+define void @i(i1 zeroext %a, i1 zeroext %b) safestack {
+; CHECK-LABEL: define void @i
+entry:
+; CHECK:        %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+; CHECK-NEXT:   getelementptr i8, i8* %[[USP]], i32 -32
+  %x1 = alloca i64, align 8
+  %x2 = alloca i64, align 8
+  %y = alloca i64, align 8
+  %y1 = alloca i64, align 8
+  %y2 = alloca i64, align 8
+  %z = alloca i64, align 8
+  %z1 = alloca i64, align 8
+  %z2 = alloca i64, align 8
+  %0 = bitcast i64* %x1 to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %0)
+  %1 = bitcast i64* %x2 to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %1)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -8
+; CHECK:   call void @capture64(
+  call void @capture64(i64* nonnull %x1)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -16
+; CHECK:   call void @capture64(
+  call void @capture64(i64* nonnull %x2)
+  br i1 %a, label %if.then, label %if.else4
+
+if.then:                                          ; preds = %entry
+  %2 = bitcast i64* %y to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %2)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -24
+; CHECK:   call void @capture64(
+  call void @capture64(i64* nonnull %y)
+  br i1 %b, label %if.then3, label %if.else
+
+if.then3:                                         ; preds = %if.then
+  %3 = bitcast i64* %y1 to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %3)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -32
+; CHECK:   call void @capture64(
+  call void @capture64(i64* nonnull %y1)
+  call void @llvm.lifetime.end(i64 -1, i8* %3)
+  br label %if.end
+
+if.else:                                          ; preds = %if.then
+  %4 = bitcast i64* %y2 to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %4)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -32
+; CHECK:   call void @capture64(
+  call void @capture64(i64* nonnull %y2)
+  call void @llvm.lifetime.end(i64 -1, i8* %4)
+  br label %if.end
+
+if.end:                                           ; preds = %if.else, %if.then3
+  call void @llvm.lifetime.end(i64 -1, i8* %2)
+  br label %if.end9
+
+if.else4:                                         ; preds = %entry
+  %5 = bitcast i64* %z to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %5)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -24
+; CHECK:   call void @capture64(
+  call void @capture64(i64* nonnull %z)
+  br i1 %b, label %if.then6, label %if.else7
+
+if.then6:                                         ; preds = %if.else4
+  %6 = bitcast i64* %z1 to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %6)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -32
+; CHECK:   call void @capture64(
+  call void @capture64(i64* nonnull %z1)
+  call void @llvm.lifetime.end(i64 -1, i8* %6)
+  br label %if.end8
+
+if.else7:                                         ; preds = %if.else4
+  %7 = bitcast i64* %z2 to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %7)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -32
+; CHECK:   call void @capture64(
+  call void @capture64(i64* nonnull %z2)
+  call void @llvm.lifetime.end(i64 -1, i8* %7)
+  br label %if.end8
+
+if.end8:                                          ; preds = %if.else7, %if.then6
+  call void @llvm.lifetime.end(i64 -1, i8* %5)
+  br label %if.end9
+
+if.end9:                                          ; preds = %if.end8, %if.end
+  call void @llvm.lifetime.end(i64 -1, i8* %1)
+  call void @llvm.lifetime.end(i64 -1, i8* %0)
+  ret void
+}
+
+; lifetime for x ends in 2 different BBs
+define void @no_merge1(i1 %d) safestack {
+; CHECK-LABEL: define void @no_merge1(
+entry:
+; CHECK:        %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+; CHECK-NEXT:   getelementptr i8, i8* %[[USP]], i32 -16
+  %x = alloca i32, align 4
+  %y = alloca i32, align 4
+  %x0 = bitcast i32* %x to i8*
+  %y0 = bitcast i32* %y to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %x0)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -4
+; CHECK:   call void @capture32(
+  call void @capture32(i32* %x)
+  br i1 %d, label %bb2, label %bb3
+bb2:
+  call void @llvm.lifetime.start(i64 -1, i8* %y0)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -8
+; CHECK:   call void @capture32(
+  call void @capture32(i32* %y)
+  call void @llvm.lifetime.end(i64 -1, i8* %y0)
+  call void @llvm.lifetime.end(i64 -1, i8* %x0)
+  ret void
+bb3:
+  call void @llvm.lifetime.end(i64 -1, i8* %x0)
+  ret void
+}
+
+define void @merge1(i1 %d) safestack {
+; CHECK-LABEL: define void @merge1(
+entry:
+; CHECK:        %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+; CHECK-NEXT:   getelementptr i8, i8* %[[USP]], i32 -16
+  %x = alloca i32, align 4
+  %y = alloca i32, align 4
+  %x0 = bitcast i32* %x to i8*
+  %y0 = bitcast i32* %y to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %x0)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -4
+; CHECK:   call void @capture32(
+  call void @capture32(i32* %x)
+  call void @llvm.lifetime.end(i64 -1, i8* %x0)
+  br i1 %d, label %bb2, label %bb3
+bb2:
+  call void @llvm.lifetime.start(i64 -1, i8* %y0)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -4
+; CHECK:   call void @capture32(
+  call void @capture32(i32* %y)
+  call void @llvm.lifetime.end(i64 -1, i8* %y0)
+  ret void
+bb3:
+  ret void
+}
+
+; Missing lifetime.end
+define void @merge2_noend(i1 %d) safestack {
+; CHECK-LABEL: define void @merge2_noend(
+entry:
+; CHECK:        %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+; CHECK-NEXT:   getelementptr i8, i8* %[[USP]], i32 -16
+  %x = alloca i32, align 4
+  %y = alloca i32, align 4
+  %x0 = bitcast i32* %x to i8*
+  %y0 = bitcast i32* %y to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %x0)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -4
+; CHECK:   call void @capture32(
+  call void @capture32(i32* %x)
+  call void @llvm.lifetime.end(i64 -1, i8* %x0)
+  br i1 %d, label %bb2, label %bb3
+bb2:
+  call void @llvm.lifetime.start(i64 -1, i8* %y0)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -4
+; CHECK:   call void @capture32(
+  call void @capture32(i32* %y)
+  ret void
+bb3:
+  ret void
+}
+
+; Missing lifetime.end
+define void @merge3_noend(i1 %d) safestack {
+; CHECK-LABEL: define void @merge3_noend(
+entry:
+; CHECK:        %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+; CHECK-NEXT:   getelementptr i8, i8* %[[USP]], i32 -16
+  %x = alloca i32, align 4
+  %y = alloca i32, align 4
+  %x0 = bitcast i32* %x to i8*
+  %y0 = bitcast i32* %y to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %x0)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -4
+; CHECK:   call void @capture32(
+  call void @capture32(i32* %x)
+  br i1 %d, label %bb2, label %bb3
+bb2:
+  call void @llvm.lifetime.end(i64 -1, i8* %x0)
+  call void @llvm.lifetime.start(i64 -1, i8* %y0)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -4
+; CHECK:   call void @capture32(
+  call void @capture32(i32* %y)
+  ret void
+bb3:
+  ret void
+}
+
+; Missing lifetime.start
+define void @nomerge4_nostart(i1 %d) safestack {
+; CHECK-LABEL: define void @nomerge4_nostart(
+entry:
+; CHECK:        %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+; CHECK-NEXT:   getelementptr i8, i8* %[[USP]], i32 -16
+  %x = alloca i32, align 4
+  %y = alloca i32, align 4
+  %x0 = bitcast i32* %x to i8*
+  %y0 = bitcast i32* %y to i8*
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -4
+; CHECK:   call void @capture32(
+  call void @capture32(i32* %x)
+  call void @llvm.lifetime.end(i64 -1, i8* %x0)
+  br i1 %d, label %bb2, label %bb3
+bb2:
+  call void @llvm.lifetime.start(i64 -1, i8* %y0)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -8
+; CHECK:   call void @capture32(
+  call void @capture32(i32* %y)
+  ret void
+bb3:
+  ret void
+}
+
+define void @array_merge() safestack {
+; CHECK-LABEL: define void @array_merge(
+entry:
+; CHECK:        %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+; CHECK-NEXT:   getelementptr i8, i8* %[[USP]], i32 -800
+  %A.i1 = alloca [100 x i32], align 4
+  %B.i2 = alloca [100 x i32], align 4
+  %A.i = alloca [100 x i32], align 4
+  %B.i = alloca [100 x i32], align 4
+  %0 = bitcast [100 x i32]* %A.i to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %0)
+  %1 = bitcast [100 x i32]* %B.i to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %1)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -400
+; CHECK:   call void @capture100x32(
+  call void @capture100x32([100 x i32]* %A.i)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -800
+; CHECK:   call void @capture100x32(
+  call void @capture100x32([100 x i32]* %B.i)
+  call void @llvm.lifetime.end(i64 -1, i8* %0)
+  call void @llvm.lifetime.end(i64 -1, i8* %1)
+  %2 = bitcast [100 x i32]* %A.i1 to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %2)
+  %3 = bitcast [100 x i32]* %B.i2 to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %3)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -400
+; CHECK:   call void @capture100x32(
+  call void @capture100x32([100 x i32]* %A.i1)
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -800
+; CHECK:   call void @capture100x32(
+  call void @capture100x32([100 x i32]* %B.i2)
+  call void @llvm.lifetime.end(i64 -1, i8* %2)
+  call void @llvm.lifetime.end(i64 -1, i8* %3)
+  ret void
+}
+
+define void @myCall_pr15707() safestack {
+; CHECK-LABEL: define void @myCall_pr15707(
+entry:
+; CHECK:        %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+; CHECK-NEXT:   getelementptr i8, i8* %[[USP]], i32 -200000
+  %buf1 = alloca i8, i32 100000, align 16
+  %buf2 = alloca i8, i32 100000, align 16
+
+  call void @llvm.lifetime.start(i64 -1, i8* %buf1)
+  call void @llvm.lifetime.end(i64 -1, i8* %buf1)
+
+  call void @llvm.lifetime.start(i64 -1, i8* %buf1)
+  call void @llvm.lifetime.start(i64 -1, i8* %buf2)
+  call void @capture8(i8* %buf1)
+  call void @capture8(i8* %buf2)
+  ret void
+}
+
+; Check that we don't assert and crash even when there are allocas
+; outside the declared lifetime regions.
+define void @bad_range() safestack {
+; CHECK-LABEL: define void @bad_range(
+entry:
+; CHECK:        %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+; A.i and B.i unsafe, not merged
+; CHECK-NEXT:   getelementptr i8, i8* %[[USP]], i32 -800
+; A.i1 and B.i2 safe
+; CHECK: = alloca [100 x i32], align 4
+; CHECK: = alloca [100 x i32], align 4
+
+  %A.i1 = alloca [100 x i32], align 4
+  %B.i2 = alloca [100 x i32], align 4
+  %A.i = alloca [100 x i32], align 4
+  %B.i = alloca [100 x i32], align 4
+  %0 = bitcast [100 x i32]* %A.i to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %0) nounwind
+  %1 = bitcast [100 x i32]* %B.i to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %1) nounwind
+  call void @capture100x32([100 x i32]* %A.i)
+  call void @capture100x32([100 x i32]* %B.i)
+  call void @llvm.lifetime.end(i64 -1, i8* %0) nounwind
+  call void @llvm.lifetime.end(i64 -1, i8* %1) nounwind
+  br label %block2
+
+block2:
+  ; I am used outside the marked lifetime.
+  call void @capture100x32([100 x i32]* %A.i)
+  call void @capture100x32([100 x i32]* %B.i)
+  ret void
+}
+
+%struct.Klass = type { i32, i32 }
+
+define i32 @shady_range(i32 %argc, i8** nocapture %argv) safestack {
+; CHECK-LABEL: define i32 @shady_range(
+entry:
+; CHECK:        %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+; CHECK-NEXT:   getelementptr i8, i8* %[[USP]], i32 -64
+  %a.i = alloca [4 x %struct.Klass], align 16
+  %b.i = alloca [4 x %struct.Klass], align 16
+  %a8 = bitcast [4 x %struct.Klass]* %a.i to i8*
+  %b8 = bitcast [4 x %struct.Klass]* %b.i to i8*
+  ; I am used outside the lifetime zone below:
+  %z2 = getelementptr inbounds [4 x %struct.Klass], [4 x %struct.Klass]* %a.i, i64 0, i64 0, i32 0
+  call void @llvm.lifetime.start(i64 -1, i8* %a8)
+  call void @llvm.lifetime.start(i64 -1, i8* %b8)
+  call void @capture8(i8* %a8)
+  call void @capture8(i8* %b8)
+  %z3 = load i32, i32* %z2, align 16
+  call void @llvm.lifetime.end(i64 -1, i8* %a8)
+  call void @llvm.lifetime.end(i64 -1, i8* %b8)
+  ret i32 %z3
+}
+
+declare void @llvm.lifetime.start(i64, i8* nocapture)
+declare void @llvm.lifetime.end(i64, i8* nocapture)
+declare void @capture8(i8*)
+declare void @capture32(i32*)
+declare void @capture64(i64*)
+declare void @capture100x32([100 x i32]*)