StackColoring: smarter check for slot overlap

[arielb1]

Summary

The old check for slot overlap treated 2 slots S and T as
overlapping if there existed a CFG node in which both of the slots could
possibly be active. That is overly conservative and caused stack blowups
in Rust programs. Instead, check whether there is a single CFG node in
which both of the slots are possibly active *together*.

Fixes PR32488.

[thanm]

Hi Ariel,

I am thinking that maybe there is something wrong with the patch on phab-- I downloaded it and it would not compile:

StackColoring.cpp:1050:5: error: use of undeclared identifier 'UI'
    UI->getNextValue(Indexes->getZeroIndex(), VNInfoAllocator);

I fixed that up in the obvious way.

Aside from that, I'd say that so far I don't think I completely understand how your strategy works. Part of this may be due simple to terminology and the names you've given to the new things.

For example, I don't get why you've chosen the name "Use" for the new set in BlockInfo. The fact that gets recorded in this set is not really a "use" of the stack slot. Consider these two blocks:

BB1:
  lifetime_end slot 5
  branch

BB2:
  lifetime_end slot 5
  lifetime_start slot 5
  use of slot 5
  branch

The way the code works now, "Use" will be set for BB1 and not for BB2, even though BB2 has an actual use (reference to) the slot and BB2 does not. This is very confusing for people reading the code.

Similar confusion over your use of the term "active" in the comment ("active" is also used at various other places in the file comments, and I don't think your definition of "active" is the same as that used in the other places). Ditto for "IntervalStarts" -- why is it named this way?

Another point of confusion is over the use of the term "interference graph" in the comments-- the data structure being used to capture interferences is not really a graph, there are no explicit edges to speak of. A real interference graph would be an improvement over what's there (could be more precise in a number of cases), but that's not what exists in the code today.

A final question is what sort of testing you've done so far to make sure this works for all the odd corner cases. LNT and bootstrap build are two things that come to mind.

With that said, please keep working on this, I think it's worth pursuing. I am OOO all this week and possibly next, so I won't have a lot of time to look at it, but I'll get there eventually.

Regards,

Than

[arielb1]

Typofix UI -> SI

[arielb1]

In D31583#718444, @thanm wrote:

Hi Ariel,

I am thinking that maybe there is something wrong with the patch on phab-- I downloaded it and it would not compile:

StackColoring.cpp:1050:5: error: use of undeclared identifier 'UI'
    UI->getNextValue(Indexes->getZeroIndex(), VNInfoAllocator);

I fixed that up in the obvious way.

Fixed that.

Aside from that, I'd say that so far I don't think I completely understand how your strategy works. Part of this may be due simple to terminology and the names you've given to the new things.

For example, I don't get why you've chosen the name "Use" for the new set in BlockInfo. The fact that gets recorded in this set is not really a "use" of the stack slot. Consider these two blocks:

BB1:
  lifetime_end slot 5
  branch

BB2:
  lifetime_end slot 5
  lifetime_start slot 5
  use of slot 5
  branch

The way the code works now, "Use" will be set for BB1 and not for BB2, even though BB2 has an actual use (reference to) the slot and BB2 does not. This is very confusing for people reading the code.

BLI.Use refes, as it says, to a slot that begins and ends within the same block. Maybe I should have called it "BeginAndEnd" or something? Maybe I should have Use also include Begin (actually, let's do that)?

Similar confusion over your use of the term "active" in the comment ("active" is also used at various other places in the file comments, and I don't think your definition of "active" is the same as that used in the other places). Ditto for "IntervalStarts" -- why is it named this way?

I thought "active" was a new fresh term. I informally referred to it as "live", but unfortunately "live" has another meaning. Is there a better term than "active"?

Another point of confusion is over the use of the term "interference graph" in the comments-- the data structure being used to capture interferences is not really a graph, there are no explicit edges to speak of. A real interference graph would be an improvement over what's there (could be more precise in a number of cases), but that's not what exists in the code today.

I was thinking about things theoretically.

A final question is what sort of testing you've done so far to make sure this works for all the odd corner cases. LNT and bootstrap build are two things that come to mind.

I bootstrapped rustc and passed all of its tests. I don't know how to bootstrap LLVM, so I didn't do that.

With that said, please keep working on this, I think it's worth pursuing. I am OOO all this week and possibly next, so I won't have a lot of time to look at it, but I'll get there eventually.

Thanks.

Regards,

Than

[• dberlin]

To try to echo a bit what than is saying:, i think it would be helpful if you could write some paragraphs of what you think the high level design of this new part is.

IE something like:
we create a dataflow problem that tells us x
we do this by computing these attributes on a per-block basis, and then finding a maximal fixpoint.
We then use x to prune y.

or whatever.

I can think of a lot of different approaches than stack coloring currently takes (IE you could do real liveness analysis, as than mentions, you could build interference graphs and simplify them, etc), but they are more fundamental changes than i see here.

[arielb1]

Change names and comment to not mention interference graph.

[arielb1]

In D31583#719207, @dberlin wrote:

To try to echo a bit what than is saying:, i think it would be helpful if you could write some paragraphs of what you think the high level design of this new part is.

IE something like:
we create a dataflow problem that tells us x
we do this by computing these attributes on a per-block basis, and then finding a maximal fixpoint.
We then use x to prune y.

or whatever.

I can think of a lot of different approaches than stack coloring currently takes (IE you could do real liveness analysis, as than mentions, you could build interference graphs and simplify them, etc), but they are more fundamental changes than i see here.

I have a big comment in the source code that is supposed to be that explanation. Is there anything missing?

[• dberlin]

lib/CodeGen/StackColoring.cpp

1270

Or you could just augment the standard dataflow problem with more info?
You also just kind of assert it doesn't give very good results, but it's pretty widely used with good results :)

1308

Errr, what?
It is very possible to do this in N time with N dataflow facts.

[arielb1]

lib/CodeGen/StackColoring.cpp

1270

What I said is that what the old code was doing - propagating the N S active facts across the CFG and that
saying that 2 nodes interfere if they are both potentially active at a point - is overconservative at merge points.
Merging this PR caused x16 stack usage reductions in real world Rust compiler stack usage.

My algorithm also propagates the N S active facts, but uses them along with S active starts non-propagated facts to compute S active & T active more precisely.

Liveness can make it smarter, but aliasing makes it annoying to compute (and if there are opaque function calls in the merge BB, impossible to compute).

[dotdash]

test/CodeGen/X86/StackColoring.ll

600

I think you want to use %bar here.

[arielb1]

Clean up code to LLVM standards - thanks @doener. Also add support for multiple live ranges.

[dotdash]

Since I had a hand in changing the logic for the live interval calculation, I
feel like I should take some time to explain the approach that was taken here.

Originally, the stack coloring pass collected the lifetime markers and assumed
that there are at most two such markers per MBB. Either in the order start ->
end, creating a single segment in the live interval. Or in the order end ->
start, creating two segments in the live interval, the slot being dead in the
middle of the MBB.

Later the code was adjusted to handle more than two markers per block. Because
the collected markers were in a random order, instead of forming multiple
segments only a single segment was created and extended as needed.
Unfortunately this actually broke the logic for handling two markers in the
order end -> start. This is because the LiveIn and LiveOut bits would then also
be set, and the segment gets extended to cover the whole MBB.

By now, we're iterating over the instructions in the MBB in order anyway, so we
can actually handle multiple start/end markers properly and create multiple
segments in the live interval even within a single block.

The semantics are as follows:

A slot for which there are no lifetime markers is always live.

For slots that have lifetime markers:

In the entry MBB, the slot starts out as dead. In other MBBs, the slot starts
out as live if the dataflow analysis determined it to be live coming into this
block, otherwise it starts as dead.

Iterating over the instructions in the MBB in sequential order:

A dead slot becomes live when it encounters a START (which could be a use), and
stays live until it encounters an END. Any START encountered while a slot is
live has no effect. This is necessary, because the START could be a plain use
rather than a lifetime_start and may thus not shorten the live interval.

A live slot becomes dead when it encounters an END. At this point, the slots
live interval gets a new segment that starts at the index where the slot became
live, and ends at the index the END was encountered. Any END encountered while
a slot is dead has no effect.

Building upon that, the improvement this patch was initially motivated by is
described below.

Because a lifetime end on a dead slot has no effect, a frontend may choose to
combine certain code paths to produce fewer BBs. Consider this pseudocode:

A = alloca TYPE
B = alloca TYPE
C = alloca TYPE_WITH_DTOR

main:
  LT_START(C)

  if COND {
    LT_START(A)
    INVOKE func UNWIND cleanup_A
    LT_END(A)
  } else {
    LT_START(B)
    INVOKE func UNWIND cleanup_B
    LT_END(B)
  }

  DTOR(C)
  LT_END(C)

  RETURN

cleanup_A:
  LP
  LT_END(A)
  br cleanup;

cleanup_B:
  LP
  LT_END(B)
  br cleanup;

cleanup:
  DTOR(C)
  LT_END(C)
  RESUME

Here we need two distinct cleanup paths and landing pads just to ensure that
the old stack coloring code can merge slots A and B. But assuming that a
lifetime end on a dead slot is a no-op, we could also write:

A = alloca TYPE
B = alloca TYPE
C = alloca TYPE_WITH_DTOR

main:
  LT_START(C)

  if COND {
    LT_START(A)
    INVOKE func UNWIND cleanup_A
    LT_END(A)
  } else {
    LT_START(B)
    INVOKE func UNWIND cleanup_B
    LT_END(B)
  }

  DTOR(C)
  LT_END(C)

  RETURN

cleanup:
  LP
  LT_END(A)
  LT_END(B)
  DTOR(C)
  LT_END(C)
  RESUME

Now the problem is that both A and B are only "possibly live" coming into
"cleanup". "Possibly live" meaning that a slot is live on one but not
necessarily all incoming edges. Using a plain overlap check on the live
intervals created using this information causes false positives, stopping the
slots from being merged.

To solve this, we can use an alternative approach to check whether two slots
are live at the same time. For two slots to be live at the same time, one of
them needs to become live when the other is live as well. We can check this by
keeping track of the points at which a slot becomes live. As these points tell
us that a slot is "definitely" live, we get more accurate results.

I hope this explains the approach we've taken well enough.

[arielb1]

Ready and just waiting for review.

[arielb1]

rnk can you review this? This is stuck in the queue for almost a month.

[thanm]

In D31583#754945, @arielb1 wrote:

Ready and just waiting for review.

Thanks for the ping/reminder.

I read through Björn's April 17th comment. The explanation makes sense, but this is something that should be in the code as opposed to only in a Phab review. Please add something along these lines to the "Implementation Notes" command section -- that is where other material exists relating to how the data flow analysis works.

My existing comments regarding naming still stand (overloading of the terms "active" and naming of the BB set "Use").

Ditto on my previous comment regarding testing. I recommend running a clang/llvm bootstrap build to make sure you haven't broken anything. You can find instructions on how to do bootstrap builds at http://llvm.org/docs/AdvancedBuilds.html. This is purely an advisory suggestion, you are free to ignore it (I mention it mainly since when I made changes to stack coloring previously it was helpful in flushing out problems that I hadn't anticipated).

[arielb1]

Tests run locally.

[arielb1]

Updated comment.

[arielb1]

Works locally.

[arielb1]

Ping. I think the new comment is good enough to go.

[thanm]

I pulled your patch and did a bootstrap build with it, no issues. Overall this seems like an improvement over the initial patch. Please see inline comments.

lib/CodeGen/StackColoring.cpp

186

This is kind of a nit, but if I am reading this correctly, your use of "non-conservative" here refers to a slot that is not degenerate, which is not talked about until down on line 342. Maybe it would make sense to put in a forward reference?

833

I see "if (!IsStart && ..." -- seems like the "!IsStart" is redundant at this point, no?

1188

A couple of observations:

First, relative to your first patch I think I can make more more sense out of what's going on. I can also see where you're getting the improvement, since doing the conflict check based on starts is inherently more precise.

Second, there is a concern that in a large function with many packets overlapped together you could get into quadratic compile time behavior (since in the worst case the interval has O(N) items and the LiveStarts vector has O(N) items). If you haven't seen any issues in practice, though, I suppose it's probably not worrying too much about.

[arielb1]

Changed else/if nesting in that part.

I'm not too worried about any O(n^2) part - both the old and the new code do an O(# slots) * O(# slots) overlap check, and I don't remember it being a particularly expensive pass.

[arielb1]

Use "degenerate" in all comments.

[thanm]

LGTM.

FWIW I tried a self-build of clang with and without the change and looked at the -stats output to see what sort of improvement there is. Looks pretty modest (<1%) but it is positive overall.

lib/CodeGen/StackColoring.cpp

LifetimeStartOnFirstUse("stackcoloring-lifetime-start-on-first-use",
         cl::desc("Treat stack lifetimes as starting on first use, not on START marker."));
 
 
 STATISTIC(NumMarkerSeen,  "Number of lifetime markers found.");
 STATISTIC(StackSpaceSaved, "Number of bytes saved due to merging slots.");
 STATISTIC(StackSlotMerged, "Number of stack slot merged.");
 STATISTIC(EscapedAllocas, "Number of allocas that escaped the lifetime region");
 
+//===----------------------------------------------------------------------===//
+//                           StackColoring Pass
+//===----------------------------------------------------------------------===//
+//
+// Stack Coloring reduces stack usage by merging stack slots when they
+// can't be used together. For example, consider the following C program:
+//
+//     void bar(char *, int);
+//     void foo(bool var) {
+//         A: {
+//             char z[4096];
+//             bar(z, 0);
+//         }
+//
+//         char *p;
+//         char x[4096];
+//         char y[4096];
+//         if (var) {
+//             p = x;
+//         } else {
+//             bar(y, 1);
+//             p = y + 1024;
+//         }
+//     B:
+//         bar(p, 2);
+//     }
+//
+// Naively-compiled, this program would use 12k of stack space. However, the
+// stack slot corresponding to `z` is always destroyed before either of the
+// stack slots for `x` or `y` are used, and then `x` is only used if `var`
+// is true, while `y` is only used if `var` is false. So in no time are 2
+// of the stack slots used together, and therefore we can merge them,
+// compiling the function using only a single 4k alloca:
+//
+//     void foo(bool var) { // equivalent
+//         char x[4096];
+//         char *p;
+//         bar(x, 0);
+//         if (var) {
+//             p = x;
+//         } else {
+//             bar(x, 1);
+//             p = x + 1024;
+//         }
+//         bar(p, 2);
+//     }
+//
+// This is an important optimization if we want stack space to be under
+// control in large functions, both open-coded ones and ones created by
+// inlining.
 //
 // Implementation Notes:
 // ---------------------
 //
+// An important part of the above reasoning is that `z` can't be accessed
+// while the latter 2 calls to `bar` are running. This is justified because
+// `z`'s lifetime is over after we exit from block `A:`, so any further
+// accesses to it would be UB. The way we represent this information
+// in LLVM is by having frontends delimit blocks with `lifetime.start`
+// and `lifetime.end` intrinsics.
+//
+// The effect of these intrinsics seems to be as follows (maybe I should
+// specify this in the reference?):
+//
+//   L1) at start, each stack-slot is marked as *out-of-scope*, unless no
+//   lifetime intrinsic refers to that stack slot, in which case
+//   it is marked as *in-scope*.
+//   L2) on a `lifetime.start`, a stack slot is marked as *in-scope* and
+//   the stack slot is overwritten with `undef`.
+//   L3) on a `lifetime.end`, a stack slot is marked as *out-of-scope*.
+//   L4) on function exit, all stack slots are marked as *out-of-scope*.
+//   L5) `lifetime.end` is a no-op when called on a slot that is already
+//   *out-of-scope*.
+//   L6) memory accesses to *out-of-scope* stack slots are UB.
+//   L7) when a stack-slot is marked as *out-of-scope*, all pointers to it
+//   are invalidated unless it looks like they might be used (?). This
+//   is used to justify not marking slots as live until the pointer
+//   to them is used, but I think this should be clarified better.
+//
+// Now, let's ground stack coloring on these rules. We'll define a slot
+// as *in-use* at a (dynamic) point in execution if it either can be
+// written to at that point, or if it has a live and non-undef content
+// at that point.
+//
+// Obviously, slots that are never *in-use* together can be merged, and
+// in our example `foo`, the slots for `x`, `y` and `z` are never
+// in-use together (of course, sometimes slots that *are* in-use together
+// might still be mergable, but we don't care about that here).
+//
+// In this implementation, we successively merge pairs of slots that are
+// not *in-use* together. We could be smarter - for example, we could merge
+// a single large slot with 2 small slots, or we could construct the
+// interference graph and run a "smart" graph coloring algorithm, but with
+// that aside, how do we find out whether a pair of slots might be *in-use*
+// together?
+//
+// From our rules, we see that *out-of-scope* slots are never *in-use*,
+// and from (L7) we see that "non-conservative" slots remain non-*in-use*
+// until their address is taken. Therefore, we can approximate slot activity
+// using dataflow.
+//
+// A subtle point: naively, we might try to figure out which pairs of
+// stack-slots interfere by propagating `S in-use` through the CFG for every
+// stack-slot `S`, and having `S` and `T` interfere if there is a CFG point in
+// which they are both *in-use*.
+//
+// That is sound, but overly conservative in some cases: in our (artificial)
+// example `foo`, either `x` or `y` might be in use at the label `B:`, but
+// as `x` is only in use if we came in from the `var` edge and `y` only
+// if we came from the `!var` edge, they still can't be in use together.
+// See PR32488 for an important real-life case.
+//
+// If we wanted to find all points of interference precisely, we could
+// propagate `S in-use` and `S&T in-use` predicates through the CFG. That
+// would be precise, but requires propagating `O(n^2)` dataflow facts.
+//
+// However, we aren't interested in the *set* of points of interference
+// between 2 stack slots, only *whether* there *is* such a point. So we
+// can rely on a little trick: for `S` and `T` to be in-use together,
+// one of them needs to become in-use while the other is in-use (or
+// they might both become in use simultaneously). We can check this
+// by also keeping track of the points at which a stack slot might *start*
+// being in-use.
+//
+// Comments on exact first use:
+// ----------------------------
+//
 // Consider the following motivating example:
 //
 //     int foo() {
 //       char b1[1024], b2[1024];
 //       if (...) {
 //         char b3[1024];
 //         <uses of b1, b3>;
 //         return x;

 //         abc(&b);
 //    }
 //
 // If in RPO ordering chosen to walk the CFG  we happen to visit the b[k]
 // before visiting the memcpy block (which will contain the lifetime start
 // for "b" then it will appear that 'b' has a degenerate lifetime.
 //
 
-//===----------------------------------------------------------------------===//
-//                           StackColoring Pass
-//===----------------------------------------------------------------------===//
-
 namespace {
 /// StackColoring - A machine pass for merging disjoint stack allocations,
 /// marked by the LIFETIME_START and LIFETIME_END pseudo instructions.
 class StackColoring : public MachineFunctionPass {
   MachineFrameInfo *MFI;
   MachineFunction *MF;
 
   /// A class representing liveness information for a single basic block.
class StackColoring : public MachineFunctionPass {
   typedef DenseMap<const MachineBasicBlock*, BlockLifetimeInfo> LivenessMap;
   LivenessMap BlockLiveness;
 
   /// Maps serial numbers to basic blocks.
   DenseMap<const MachineBasicBlock*, int> BasicBlocks;
   /// Maps basic blocks to a serial number.
   SmallVector<const MachineBasicBlock*, 8> BasicBlockNumbering;
 
-  /// Maps liveness intervals for each slot.
+  /// Maps slots to their use interval. Outside of this interval, slots
+  /// values are either dead or `undef` and they will not be written to.
   SmallVector<std::unique_ptr<LiveInterval>, 16> Intervals;
+  /// Maps slots to the points where they can become in-use.
+  SmallVector<SmallVector<SlotIndex, 4>, 16> LiveStarts;
   /// VNInfo is used for the construction of LiveIntervals.
   VNInfo::Allocator VNInfoAllocator;
   /// SlotIndex analysis object.
   SlotIndexes *Indexes;
   /// The stack protector object.
   StackProtector *SP;
 
   /// The list of lifetime markers found. These markers are to be removed
while (changed) {
     }
   }// while changed.
 
   NumIterations = NumIters;
 }
 
 void StackColoring::calculateLiveIntervals(unsigned NumSlots) {
   SmallVector<SlotIndex, 16> Starts;
-  SmallVector<SlotIndex, 16> Finishes;
+  SmallVector<bool, 16> DefinitelyInUse;
 
   // For each block, find which slots are active within this block
   // and update the live intervals.
   for (const MachineBasicBlock &MBB : *MF) {
     Starts.clear();
     Starts.resize(NumSlots);
-    Finishes.clear();
-    Finishes.resize(NumSlots);
+    DefinitelyInUse.clear();
+    DefinitelyInUse.resize(NumSlots);
+
+    // Start the interval of the slots that we previously found to be 'in-use'.
+    BlockLifetimeInfo &MBBLiveness = BlockLiveness[&MBB];
+    for (int pos = MBBLiveness.LiveIn.find_first(); pos != -1;
+         pos = MBBLiveness.LiveIn.find_next(pos)) {
+      Starts[pos] = Indexes->getMBBStartIdx(&MBB);
+    }
 
     // Create the interval for the basic blocks containing lifetime begin/end.
     for (const MachineInstr &MI : MBB) {
 
       SmallVector<int, 4> slots;
       bool IsStart = false;
       if (!isLifetimeStartOrEnd(MI, slots, IsStart))
         continue;
       SlotIndex ThisIndex = Indexes->getInstructionIndex(MI);
       for (auto Slot : slots) {
         if (IsStart) {
-          if (!Starts[Slot].isValid() || Starts[Slot] > ThisIndex)
+          // If a slot is already definitely in use, we don't have to emit
+          // a new start marker because there is already a pre-existing
+          // one.
+          if (!DefinitelyInUse[Slot]) {
+            LiveStarts[Slot].push_back(ThisIndex);
+            DefinitelyInUse[Slot] = true;
+          }
+          if (!Starts[Slot].isValid())
             Starts[Slot] = ThisIndex;
-        } else {
-          if (!Finishes[Slot].isValid() || Finishes[Slot] < ThisIndex)
-            Finishes[Slot] = ThisIndex;
+        } else if (!IsStart && Starts[Slot].isValid()) {
+          VNInfo *VNI = Intervals[Slot]->getValNumInfo(0);
+          Intervals[Slot]->addSegment(
+              LiveInterval::Segment(Starts[Slot], ThisIndex, VNI));
+          Starts[Slot] = SlotIndex(); // Invalidate the start index
+          DefinitelyInUse[Slot] = false;
         }
       }
     }
 
-    // Create the interval of the blocks that we previously found to be 'alive'.
-    BlockLifetimeInfo &MBBLiveness = BlockLiveness[&MBB];
-    for (unsigned pos : MBBLiveness.LiveIn.set_bits()) {
-      Starts[pos] = Indexes->getMBBStartIdx(&MBB);
-    }
-    for (unsigned pos : MBBLiveness.LiveOut.set_bits()) {
-      Finishes[pos] = Indexes->getMBBEndIdx(&MBB);
-    }
-
+    // Finish up started segments
     for (unsigned i = 0; i < NumSlots; ++i) {
-      //
-      // When LifetimeStartOnFirstUse is turned on, data flow analysis
-      // is forward (from starts to ends), not bidirectional. A
-      // consequence of this is that we can wind up in situations
-      // where Starts[i] is invalid but Finishes[i] is valid and vice
-      // versa. Example:
-      //
-      //     LIFETIME_START x
-      //     if (...) {
-      //       <use of x>
-      //       throw ...;
-      //     }
-      //     LIFETIME_END x
-      //     return 2;
-      //
-      //
-      // Here the slot for "x" will not be live into the block
-      // containing the "return 2" (since lifetimes start with first
-      // use, not at the dominating LIFETIME_START marker).
-      //
-      if (Starts[i].isValid() && !Finishes[i].isValid()) {
-        Finishes[i] = Indexes->getMBBEndIdx(&MBB);
-      }
       if (!Starts[i].isValid())
         continue;
 
-      assert(Starts[i] && Finishes[i] && "Invalid interval");
-      VNInfo *ValNum = Intervals[i]->getValNumInfo(0);
-      SlotIndex S = Starts[i];
-      SlotIndex F = Finishes[i];
-      if (S < F) {
-        // We have a single consecutive region.
-        Intervals[i]->addSegment(LiveInterval::Segment(S, F, ValNum));
-      } else {
-        // We have two non-consecutive regions. This happens when
-        // LIFETIME_START appears after the LIFETIME_END marker.
-        SlotIndex NewStart = Indexes->getMBBStartIdx(&MBB);
-        SlotIndex NewFin = Indexes->getMBBEndIdx(&MBB);
-        Intervals[i]->addSegment(LiveInterval::Segment(NewStart, F, ValNum));
-        Intervals[i]->addSegment(LiveInterval::Segment(S, NewFin, ValNum));
-      }
+      SlotIndex EndIdx = Indexes->getMBBEndIdx(&MBB);
+      VNInfo *VNI = Intervals[i]->getValNumInfo(0);
+      Intervals[i]->addSegment(LiveInterval::Segment(Starts[i], EndIdx, VNI));
     }
   }
 }
 
 bool StackColoring::removeAllMarkers() {
   unsigned Count = 0;
   for (MachineInstr *MI : Markers) {
     MI->eraseFromParent();
bool StackColoring::runOnMachineFunction(MachineFunction &Func) {
   MFI = &MF->getFrameInfo();
   Indexes = &getAnalysis<SlotIndexes>();
   SP = &getAnalysis<StackProtector>();
   BlockLiveness.clear();
   BasicBlocks.clear();
   BasicBlockNumbering.clear();
   Markers.clear();
   Intervals.clear();
+  LiveStarts.clear();
   VNInfoAllocator.Reset();
 
   unsigned NumSlots = MFI->getObjectIndexEnd();
 
   // If there are no stack slots then there are no markers to remove.
   if (!NumSlots)
     return false;
 
   SmallVector<int, 8> SortedSlots;
   SortedSlots.reserve(NumSlots);
   Intervals.reserve(NumSlots);
+  LiveStarts.resize(NumSlots);
 
   unsigned NumMarkers = collectMarkers(NumSlots);
 
   unsigned TotalSize = 0;
   DEBUG(dbgs()<<"Found "<<NumMarkers<<" markers and "<<NumSlots<<" slots\n");
   DEBUG(dbgs()<<"Slot structure:\n");
 
   for (int i=0; i < MFI->getObjectIndexEnd(); ++i) {
std::stable_sort(SortedSlots.begin(), SortedSlots.end(),
                    [this](int LHS, int RHS) {
     // We use -1 to denote a uninteresting slot. Place these slots at the end.
     if (LHS == -1) return false;
     if (RHS == -1) return true;
     // Sort according to size.
     return MFI->getObjectSize(LHS) > MFI->getObjectSize(RHS);
   });
 
+  for (auto &s : LiveStarts)
+    std::sort(s.begin(), s.end());
+
   bool Changed = true;
   while (Changed) {
     Changed = false;
     for (unsigned I = 0; I < NumSlots; ++I) {
       if (SortedSlots[I] == -1)
         continue;
 
       for (unsigned J=I+1; J < NumSlots; ++J) {
         if (SortedSlots[J] == -1)
           continue;
 
         int FirstSlot = SortedSlots[I];
         int SecondSlot = SortedSlots[J];
         LiveInterval *First = &*Intervals[FirstSlot];
         LiveInterval *Second = &*Intervals[SecondSlot];
+        auto &FirstS = LiveStarts[FirstSlot];
+        auto &SecondS = LiveStarts[SecondSlot];
         assert (!First->empty() && !Second->empty() && "Found an empty range");
 
-        // Merge disjoint slots.
-        if (!First->overlaps(*Second)) {
+        // Merge disjoint slots. This is a little bit tricky - see the
+        // Implementation Notes section for an explanation.
+        if (!First->isLiveAtIndexes(SecondS) &&
+            !Second->isLiveAtIndexes(FirstS)) {
           Changed = true;
           First->MergeSegmentsInAsValue(*Second, First->getValNumInfo(0));
+
+          int OldSize = FirstS.size();
+          FirstS.append(SecondS.begin(), SecondS.end());
+          auto Mid = FirstS.begin() + OldSize;
+          std::inplace_merge(FirstS.begin(), Mid, FirstS.end());
+
           SlotRemap[SecondSlot] = FirstSlot;
           SortedSlots[J] = -1;
           DEBUG(dbgs()<<"Merging #"<<FirstSlot<<" and slots #"<<
                 SecondSlot<<" together.\n");
           unsigned MaxAlignment = std::max(MFI->getObjectAlignment(FirstSlot),
                                            MFI->getObjectAlignment(SecondSlot));
 
           assert(MFI->getObjectSize(FirstSlot) >=
DEBUG(dbgs()<<"Merge "<<RemovedSlots<<" slots. Saved "<<
         ReducedSize<<" bytes\n");
 
   // Scan the entire function and update all machine operands that use frame
   // indices to use the remapped frame index.
   expungeSlotMap(SlotRemap, NumSlots);
   remapInstructions(SlotRemap);
 
   return removeAllMarkers();
 }

test/CodeGen/X86/StackColoring.ll

if.then:                                          ; preds = %entry
   call void @llvm.lifetime.end.p0i8(i64 48, i8* %tmp1)
   br label %if.end
 
 if.end:                                           ; preds = %if.then, %entry
   %x.addr.0 = phi i32 [ %add.i9, %if.then ], [ %add.i, %entry ]
   ret i32 %x.addr.0
 }
 
+;CHECK-LABEL: multi_segment:
+;YESCOLOR: subq  $256, %rsp
+;NOFIRSTUSE: subq  $256, %rsp
+;NOCOLOR: subq  $512, %rsp
+define i1 @multi_segment(i1, i1)
+{
+entry-block:
+  %foo = alloca [32 x i64]
+  %bar = alloca [32 x i64]
+  %foo_i8 = bitcast [32 x i64]* %foo to i8*
+  %bar_i8 = bitcast [32 x i64]* %bar to i8*
+  call void @llvm.lifetime.start.p0i8(i64 256, i8* %bar_i8)
+  call void @baz([32 x i64]* %bar, i32 1)
+  call void @llvm.lifetime.end.p0i8(i64 256, i8* %bar_i8)
+  call void @llvm.lifetime.start.p0i8(i64 256, i8* %foo_i8)
+  call void @baz([32 x i64]* %foo, i32 1)
+  call void @llvm.lifetime.end.p0i8(i64 256, i8* %foo_i8)
+  call void @llvm.lifetime.start.p0i8(i64 256, i8* %bar_i8)
+  call void @baz([32 x i64]* %bar, i32 1)
+  call void @llvm.lifetime.end.p0i8(i64 256, i8* %bar_i8)
+  ret i1 true
+}
+
+;CHECK-LABEL: pr32488:
+;YESCOLOR: subq  $256, %rsp
+;NOFIRSTUSE: subq  $256, %rsp
+;NOCOLOR: subq  $512, %rsp
+define i1 @pr32488(i1, i1)
+{
+entry-block:
+  %foo = alloca [32 x i64]
+  %bar = alloca [32 x i64]
+  %foo_i8 = bitcast [32 x i64]* %foo to i8*
+  %bar_i8 = bitcast [32 x i64]* %bar to i8*
+  br i1 %0, label %if_false, label %if_true
+if_false:
+  call void @llvm.lifetime.start.p0i8(i64 256, i8* %bar_i8)
+  call void @baz([32 x i64]* %bar, i32 0)
+  br i1 %1, label %if_false.1, label %onerr
+if_false.1:
+  call void @llvm.lifetime.end.p0i8(i64 256, i8* %bar_i8)
+  br label %merge
+if_true:
+  call void @llvm.lifetime.start.p0i8(i64 256, i8* %foo_i8)
+  call void @baz([32 x i64]* %foo, i32 1)
+  br i1 %1, label %if_true.1, label %onerr
+if_true.1:
+  call void @llvm.lifetime.end.p0i8(i64 256, i8* %foo_i8)
+  br label %merge
+merge:
+  ret i1 false
+onerr:
+  call void @llvm.lifetime.end.p0i8(i64 256, i8* %foo_i8)
+  call void @llvm.lifetime.end.p0i8(i64 256, i8* %bar_i8)
+  call void @destructor()
+  ret i1 true
+}
+
+%Data = type { [32 x i64] }
+
+declare void @destructor()
+
 declare void @inita(i32*)
 
 declare void @initb(i32*,i32*,i32*)
 
 declare void @bar([100 x i32]* , [100 x i32]*) nounwind
 
+declare void @baz([32 x i64]*, i32)
+
 declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) nounwind
 
 declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) nounwind
 
 declare i32 @foo(i32, i8*)