[StackColoring] Change the StackColoring logic + enables it to handle spills

llvm/include/llvm/CodeGen/LiveInterval.h

@@ -237,9 +237,9 @@ namespace llvm {
     }
 
     /// Constructs a new LiveRange object.
-    LiveRange(bool UseSegmentSet = false)
-        : segmentSet(UseSegmentSet ? std::make_unique<SegmentSet>()
-                                   : nullptr) {}
+    explicit LiveRange(bool UseSegmentSet = false)
+        : segmentSet(UseSegmentSet ? std::make_unique<SegmentSet>() : nullptr) {
+    }
 
     /// Constructs a new LiveRange object by copying segments and valnos from
     /// another LiveRange.

llvm/include/llvm/CodeGen/LiveStacks.h

@@ -40,49 +40,43 @@ class LiveStacks {
   ///
   VNInfo::Allocator VNInfoAllocator;
 
-  /// S2IMap - Stack slot indices to live interval mapping.
-  using SS2IntervalMap = std::unordered_map<int, LiveInterval>;
-  SS2IntervalMap S2IMap;
-
-  /// S2RCMap - Stack slot indices to register class mapping.
-  std::map<int, const TargetRegisterClass *> S2RCMap;
+  int StartIdx = -1;
+  SmallVector<LiveInterval *> S2LI;
+  SmallVector<const TargetRegisterClass *> S2RC;
 
 public:
-  using iterator = SS2IntervalMap::iterator;
-  using const_iterator = SS2IntervalMap::const_iterator;
+  using iterator = SmallVector<LiveInterval *>::iterator;
+  using const_iterator = SmallVector<LiveInterval *>::const_iterator;
 
-  const_iterator begin() const { return S2IMap.begin(); }
-  const_iterator end() const { return S2IMap.end(); }
-  iterator begin() { return S2IMap.begin(); }
-  iterator end() { return S2IMap.end(); }
+  const_iterator begin() const { return S2LI.begin(); }
+  const_iterator end() const { return S2LI.end(); }
+  iterator begin() { return S2LI.begin(); }
+  iterator end() { return S2LI.end(); }
 
-  unsigned getNumIntervals() const { return (unsigned)S2IMap.size(); }
+  unsigned getStartIdx() const { return StartIdx; }
+  unsigned getNumIntervals() const { return (unsigned)S2LI.size(); }
 
   LiveInterval &getOrCreateInterval(int Slot, const TargetRegisterClass *RC);
 
   LiveInterval &getInterval(int Slot) {
     assert(Slot >= 0 && "Spill slot indice must be >= 0");
-    SS2IntervalMap::iterator I = S2IMap.find(Slot);
-    assert(I != S2IMap.end() && "Interval does not exist for stack slot");
-    return I->second;
+    return *S2LI[Slot - StartIdx];
   }
 
   const LiveInterval &getInterval(int Slot) const {
     assert(Slot >= 0 && "Spill slot indice must be >= 0");
-    SS2IntervalMap::const_iterator I = S2IMap.find(Slot);
-    assert(I != S2IMap.end() && "Interval does not exist for stack slot");
-    return I->second;
+    return *S2LI[Slot - StartIdx];
   }
 
-  bool hasInterval(int Slot) const { return S2IMap.count(Slot); }
+  bool hasInterval(int Slot) const {
+    if (Slot < StartIdx || StartIdx == -1)
+      return false;
+    return !getInterval(Slot).empty();
+  }
 
   const TargetRegisterClass *getIntervalRegClass(int Slot) const {
     assert(Slot >= 0 && "Spill slot indice must be >= 0");
-    std::map<int, const TargetRegisterClass *>::const_iterator I =
-        S2RCMap.find(Slot);
-    assert(I != S2RCMap.end() &&
-           "Register class info does not exist for stack slot");
-    return I->second;
+    return S2RC[Slot - StartIdx];
   }
 
   VNInfo::Allocator &getVNInfoAllocator() { return VNInfoAllocator; }

llvm/include/llvm/CodeGen/MachineFrameInfo.h

@@ -120,12 +120,18 @@ class MachineFrameInfo {
                       ///< triggered protection.  3rd closest to the protector.
   };
 
+  static constexpr int NoUnderlyingSlot = std::numeric_limits<int>::min();
+  static constexpr int IsUnderlyingSlot = std::numeric_limits<int>::min() + 1;
+
 private:
   // Represent a single object allocated on the stack.
   struct StackObject {
     // The offset of this object from the stack pointer on entry to
     // the function.  This field has no meaning for a variable sized element.
-    int64_t SPOffset;
+    // After getting placed this is relative to SP
+    // If UnderlyingSlot is not NoUnderlyingSlot, this is relative to the start
+    // of the UnderlyingSlot
+    int64_t Offset;
 
     // The size of this object on the stack. 0 means a variable sized object,
     // ~0ULL means a dead object.
@@ -134,6 +140,10 @@ class MachineFrameInfo {
     // The required alignment of this stack slot.
     Align Alignment;
 
+    // If not NoUnderlyingSlot, it Indicate that this slot should be placed
+    // at Offset, into the slot UnderlyingSlot
+    int UnderlyingSlot = NoUnderlyingSlot;
+
     // If true, the value of the stack object is set before
     // entering the function and is not modified inside the function. By
     // default, fixed objects are immutable unless marked otherwise.
@@ -183,10 +193,10 @@ class MachineFrameInfo {
 
     uint8_t SSPLayout = SSPLK_None;
 
-    StackObject(uint64_t Size, Align Alignment, int64_t SPOffset,
+    StackObject(uint64_t Size, Align Alignment, int64_t Offset,
                 bool IsImmutable, bool IsSpillSlot, const AllocaInst *Alloca,
                 bool IsAliased, uint8_t StackID = 0)
-        : SPOffset(SPOffset), Size(Size), Alignment(Alignment),
+        : Offset(Offset), Size(Size), Alignment(Alignment),
           isImmutable(IsImmutable), isSpillSlot(IsSpillSlot), StackID(StackID),
           Alloca(Alloca), isAliased(IsAliased) {}
   };
@@ -532,7 +542,7 @@ class MachineFrameInfo {
            "Invalid Object Idx!");
     assert(!isDeadObjectIndex(ObjectIdx) &&
            "Getting frame offset for a dead object?");
-    return Objects[ObjectIdx+NumFixedObjects].SPOffset;
+    return Objects[ObjectIdx + NumFixedObjects].Offset;
   }
 
   bool isObjectZExt(int ObjectIdx) const {
@@ -561,12 +571,12 @@ class MachineFrameInfo {
 
   /// Set the stack frame offset of the specified object. The
   /// offset is relative to the stack pointer on entry to the function.
-  void setObjectOffset(int ObjectIdx, int64_t SPOffset) {
+  void setObjectOffset(int ObjectIdx, int64_t Offset) {
     assert(unsigned(ObjectIdx+NumFixedObjects) < Objects.size() &&
            "Invalid Object Idx!");
     assert(!isDeadObjectIndex(ObjectIdx) &&
            "Setting frame offset for a dead object?");
-    Objects[ObjectIdx+NumFixedObjects].SPOffset = SPOffset;
+    Objects[ObjectIdx + NumFixedObjects].Offset = Offset;
   }
 
   SSPLayoutKind getObjectSSPLayout(int ObjectIdx) const {
@@ -762,6 +772,18 @@ class MachineFrameInfo {
     // If ID == 0, MaxAlignment will need to be updated separately.
   }
 
+  int getUnderlyingSlot(int ObjectIdx) const {
+    assert(unsigned(ObjectIdx + NumFixedObjects) < Objects.size() &&
+           "Invalid Object Idx!");
+    return Objects[ObjectIdx + NumFixedObjects].UnderlyingSlot;
+  }
+
+  void setUnderlyingSlot(int ObjectIdx, int Underlying) {
+    assert(unsigned(ObjectIdx + NumFixedObjects) < Objects.size() &&
+           "Invalid Object Idx!");
+    Objects[ObjectIdx + NumFixedObjects].UnderlyingSlot = Underlying;
+  }
+
   /// Returns true if the specified index corresponds to a dead object.
   bool isDeadObjectIndex(int ObjectIdx) const {
     assert(unsigned(ObjectIdx+NumFixedObjects) < Objects.size() &&

llvm/include/llvm/CodeGen/MachineInstr.h

@@ -1340,7 +1340,9 @@ class MachineInstr
   }
 
   // True if the instruction represents a position in the function.
-  bool isPosition() const { return isLabel() || isCFIInstruction(); }
+  bool isPosition() const {
+    return isLifetimeMarker() || isLabel() || isCFIInstruction();
+  }
 
   bool isNonListDebugValue() const {
     return getOpcode() == TargetOpcode::DBG_VALUE;

llvm/include/llvm/DebugInfo/PDB/Native/FormatUtil.h

@@ -62,49 +62,6 @@ LLVM_ABI std::string formatChunkKind(codeview::DebugSubsectionKind Kind,
 LLVM_ABI std::string formatSymbolKind(codeview::SymbolKind K);
 LLVM_ABI std::string formatTypeLeafKind(codeview::TypeLeafKind K);
 
-/// Returns the number of digits in the given integer.
-inline int NumDigits(uint64_t N) {
-  if (N < 10ULL)
-    return 1;
-  if (N < 100ULL)
-    return 2;
-  if (N < 1000ULL)
-    return 3;
-  if (N < 10000ULL)
-    return 4;
-  if (N < 100000ULL)
-    return 5;
-  if (N < 1000000ULL)
-    return 6;
-  if (N < 10000000ULL)
-    return 7;
-  if (N < 100000000ULL)
-    return 8;
-  if (N < 1000000000ULL)
-    return 9;
-  if (N < 10000000000ULL)
-    return 10;
-  if (N < 100000000000ULL)
-    return 11;
-  if (N < 1000000000000ULL)
-    return 12;
-  if (N < 10000000000000ULL)
-    return 13;
-  if (N < 100000000000000ULL)
-    return 14;
-  if (N < 1000000000000000ULL)
-    return 15;
-  if (N < 10000000000000000ULL)
-    return 16;
-  if (N < 100000000000000000ULL)
-    return 17;
-  if (N < 1000000000000000000ULL)
-    return 18;
-  if (N < 10000000000000000000ULL)
-    return 19;
-  return 20;
-}
-
 namespace detail {
 template <typename T>
 struct EndianAdapter final

llvm/include/llvm/Support/MathExtras.h

@@ -795,6 +795,49 @@ using stack_float_t = volatile float;
 using stack_float_t = float;
 #endif
 
+/// Returns the number of digits in the given integer.
+inline int NumDigits(uint64_t N) {
+  if (N < 10ULL)
+    return 1;
+  if (N < 100ULL)
+    return 2;
+  if (N < 1000ULL)
+    return 3;
+  if (N < 10000ULL)
+    return 4;
+  if (N < 100000ULL)
+    return 5;
+  if (N < 1000000ULL)
+    return 6;
+  if (N < 10000000ULL)
+    return 7;
+  if (N < 100000000ULL)
+    return 8;
+  if (N < 1000000000ULL)
+    return 9;
+  if (N < 10000000000ULL)
+    return 10;
+  if (N < 100000000000ULL)
+    return 11;
+  if (N < 1000000000000ULL)
+    return 12;
+  if (N < 10000000000000ULL)
+    return 13;
+  if (N < 100000000000000ULL)
+    return 14;
+  if (N < 1000000000000000ULL)
+    return 15;
+  if (N < 10000000000000000ULL)
+    return 16;
+  if (N < 100000000000000000ULL)
+    return 17;
+  if (N < 1000000000000000000ULL)
+    return 18;
+  if (N < 10000000000000000000ULL)
+    return 19;
+  return 20;
+}
+
 } // namespace llvm
 
 #endif

llvm/lib/CodeGen/LiveStacks.cpp

@@ -37,10 +37,12 @@ void LiveStacksWrapperLegacy::getAnalysisUsage(AnalysisUsage &AU) const {
 }
 
 void LiveStacks::releaseMemory() {
+  for (int Idx = 0; Idx < (int)S2LI.size(); ++Idx)
+    S2LI[Idx]->~LiveInterval();
   // Release VNInfo memory regions, VNInfo objects don't need to be dtor'd.
   VNInfoAllocator.Reset();
-  S2IMap.clear();
-  S2RCMap.clear();
+  S2LI.clear();
+  S2RC.clear();
 }
 
 void LiveStacks::init(MachineFunction &MF) {
@@ -52,20 +54,22 @@ void LiveStacks::init(MachineFunction &MF) {
 LiveInterval &
 LiveStacks::getOrCreateInterval(int Slot, const TargetRegisterClass *RC) {
   assert(Slot >= 0 && "Spill slot indice must be >= 0");
-  SS2IntervalMap::iterator I = S2IMap.find(Slot);
-  if (I == S2IMap.end()) {
-    I = S2IMap
-            .emplace(
-                std::piecewise_construct, std::forward_as_tuple(Slot),
-                std::forward_as_tuple(Register::index2StackSlot(Slot), 0.0F))
-            .first;
-    S2RCMap.insert(std::make_pair(Slot, RC));
+  if (StartIdx == -1)
+    StartIdx = Slot;
+
+  int Idx = Slot - StartIdx;
+  assert(Idx >= 0 && "Slot not in order ?");
+  if (Idx < (int)S2LI.size()) {
+    S2RC[Idx] = TRI->getCommonSubClass(S2RC[Idx], RC);
   } else {
-    // Use the largest common subclass register class.
-    const TargetRegisterClass *&OldRC = S2RCMap[Slot];
-    OldRC = TRI->getCommonSubClass(OldRC, RC);
+    S2RC.resize(Idx + 1);
+    S2LI.resize(Idx + 1);
+    S2LI[Idx] = this->VNInfoAllocator.Allocate<LiveInterval>();
+    new (S2LI[Idx]) LiveInterval(Register::index2StackSlot(Slot), 0.0F);
+    S2RC[Idx] = RC;
   }
-  return I->second;
+  assert(S2RC.size() == S2LI.size());
+  return *S2LI[Idx];
 }
 
 AnalysisKey LiveStacksAnalysis::Key;
@@ -96,13 +100,12 @@ void LiveStacksWrapperLegacy::print(raw_ostream &OS, const Module *) const {
 }
 
 /// print - Implement the dump method.
-void LiveStacks::print(raw_ostream &OS, const Module*) const {
+void LiveStacks::print(raw_ostream &OS, const Module *) const {
 
   OS << "********** INTERVALS **********\n";
-  for (const_iterator I = begin(), E = end(); I != E; ++I) {
-    I->second.print(OS);
-    int Slot = I->first;
-    const TargetRegisterClass *RC = getIntervalRegClass(Slot);
+  for (int Idx = 0; Idx < (int)S2LI.size(); ++Idx) {
+    S2LI[Idx]->print(OS);
+    const TargetRegisterClass *RC = S2RC[Idx];
     if (RC)
       OS << " [" << TRI->getRegClassName(RC) << "]\n";
     else