[StackSlotColoring] Fixed handling of StackID

[tpr]

Summary

StackID is a property of a stack slot that allows a target to have
multiple spaces for spilling into. AMDGPU uses it to allow SGPRs to be
spilled into a VGPR, instead of using a real stack slot.

This commit fixes StackSlotColoring to ensure that two stack objects
with different StackID really don't share a slot, even if they are
colored out of order, and that the StackID gets set up in the colored
frame ids for later use.

Change-Id: I6a4d88ad7739d5d0550162b907e394b57fb3e437

[tpr]

Sorry, no test. My reproducer is enormous. I tried to modify test/CodeGen/AMDGPU/stack-slot-color-sgpr-vgpr-spills.mir to color the two slots out of order by adding extra references to one to change the weight, but it still did it in order. So I couldn't get it to fail without this change.

[arsenm]

Duplicate D45968

[dfukalov]

I have a quite small .ll that causes assertion in SIFrameLowering.cpp, if you need

reduced2.ll3 KBDownload

[tpr]

Thanks Daniil.

[tpr]

Use D45968 instead.

lib/CodeGen/StackSlotColoring.cpp

class StackSlotColoring : public MachineFunctionPass {
     SmallVector<SmallVector<MachineMemOperand *, 8>, 16> SSRefs;
 
     // OrigAlignments - Alignments of stack objects before coloring.
     SmallVector<unsigned, 16> OrigAlignments;
 
     // OrigSizes - Sizess of stack objects before coloring.
     SmallVector<unsigned, 16> OrigSizes;
 
+    // OrigStackIDs - Identifier for stack memory type, before coloring.
+    SmallVector<uint8_t, 16> OrigStackIDs;
+
     // AllColors - If index is set, it's a spill slot, i.e. color.
     // FIXME: This assumes PEI locate spill slot with smaller indices
     // closest to stack pointer / frame pointer. Therefore, smaller
     // index == better color.
     BitVector AllColors;
 
     // NextColor - Next "color" that's not yet used.
     int NextColor = -1;

 }
 
 /// InitializeSlots - Process all spill stack slot liveintervals and add them
 /// to a sorted (by weight) list.
 void StackSlotColoring::InitializeSlots() {
   int LastFI = MFI->getObjectIndexEnd();
   OrigAlignments.resize(LastFI);
   OrigSizes.resize(LastFI);
+  OrigStackIDs.resize(LastFI);
   AllColors.resize(LastFI);
   UsedColors.resize(LastFI);
   Assignments.resize(LastFI);
 
   using Pair = std::iterator_traits<LiveStacks::iterator>::value_type;
 
   SmallVector<Pair *, 16> Intervals;
 
for (auto *I : Intervals) {
     LiveInterval &li = I->second;
     DEBUG(li.dump());
     int FI = TargetRegisterInfo::stackSlot2Index(li.reg);
     if (MFI->isDeadObjectIndex(FI))
       continue;
     SSIntervals.push_back(&li);
     OrigAlignments[FI] = MFI->getObjectAlignment(FI);
     OrigSizes[FI]      = MFI->getObjectSize(FI);
+    OrigStackIDs[FI]   = MFI->getStackID(FI);
     AllColors.set(FI);
   }
   DEBUG(dbgs() << '\n');
 
   // Sort them by weight.
   std::stable_sort(SSIntervals.begin(), SSIntervals.end(), IntervalSorter());
 
   // Get first "color".
while (Color != -1) {
         Share = true;
         ++NumEliminated;
         break;
       }
       Color = UsedColors.find_next(Color);
     }
   }
 
-  if (Color != -1 && MFI->getStackID(Color) != MFI->getStackID(FI)) {
+  if (Color != -1 && MFI->getStackID(Color) != OrigStackIDs[FI]) {
     DEBUG(dbgs() << "cannot share FIs with different stack IDs\n");
     Share = false;
   }
 
   // Assign it to the first available color (assumed to be the best) if it's
   // not possible to share a used color with other objects.
   if (!Share) {
     assert(NextColor != -1 && "No more spill slots?");
     Color = NextColor;
     UsedColors.set(Color);
     NextColor = AllColors.find_next(NextColor);
   }
 
   // Record the assignment.
   Assignments[Color].push_back(li);
   DEBUG(dbgs() << "Assigning fi#" << FI << " to fi#" << Color << "\n");
 
   // Change size and alignment of the allocated slot. If there are multiple
   // objects sharing the same slot, then make sure the size and alignment
-  // are large enough for all.
+  // are large enough for all. Also copy across the StackID.
   unsigned Align = OrigAlignments[FI];
   if (!Share || Align > MFI->getObjectAlignment(Color))
     MFI->setObjectAlignment(Color, Align);
   int64_t Size = OrigSizes[FI];
   if (!Share || Size > MFI->getObjectSize(Color))
     MFI->setObjectSize(Color, Size);
+  MFI->setStackID(Color, OrigStackIDs[FI]);
   return Color;
 }
 
 /// Colorslots - Color all spill stack slots and rewrite all frameindex machine
 /// operands in the function.
 bool StackSlotColoring::ColorSlots(MachineFunction &MF) {
   unsigned NumObjs = MFI->getObjectIndexEnd();
   SmallVector<int, 16> SlotMapping(NumObjs, -1);
bool StackSlotColoring::runOnMachineFunction(MachineFunction &MF) {
 
   NextColor = -1;
   SSIntervals.clear();
   for (unsigned i = 0, e = SSRefs.size(); i != e; ++i)
     SSRefs[i].clear();
   SSRefs.clear();
   OrigAlignments.clear();
   OrigSizes.clear();
+  OrigStackIDs.clear();
   AllColors.clear();
   UsedColors.clear();
   for (unsigned i = 0, e = Assignments.size(); i != e; ++i)
     Assignments[i].clear();
   Assignments.clear();
 
   return Changed;
 }