[CodeGen] Speedup stack slot sharing during stack coloring (interval overlapping test).

[vpykhtin]

Summary

AMDGPU code with enabled address sanitizer generates tons of stack objects (> 200000 in my testcase) and
takes forever to compile due to the time spent on stack slot sharing.

While LiveRange::overlaps method has logarithmic complexity on the number of segments in the involved
liveranges the problem is that when a new interval is assigned to a used color it's tested against
overlapping every other assigned interval for that color.

Instead I decided to join all assigned intervals for a color into a single interval and this allows to
have logarithmic complexity on the number of segments for the joined interval.

This patch reduced time spent on stack slot coloring pass from 628 to 3 seconds on my testcase.

[arsenm]

llvm/lib/CodeGen/StackSlotColoring.cpp

99

Typo sigle

124

Don't understand why this is the case? Don't these need to be renumbered?

[vpykhtin]

llvm/lib/CodeGen/StackSlotColoring.cpp

124

Joined interval isn't used for anything else than range overlapping test, I don't need value numbers.

[arsenm]

llvm/lib/CodeGen/StackSlotColoring.cpp

124

You might not need it, but will it fail the verifier?

[vpykhtin]

llvm/lib/CodeGen/StackSlotColoring.cpp

124

It never reaches verifier I believe, it lives locally and dies after the overlapping test is done. It's not a fully valid object I agree.

[vpykhtin]

• fixed typo, debug build, added comment.

[MatzeB]

Hmm the main RAGreedy allocator is using the LiveRegMatrix and within that the LiveIntervalUnion classes to keep things performant. Maybe we can use those here too? In principle they should have even better runtime behavior than merging here...

[vpykhtin]

In D146057#4193602, @MatzeB wrote:

Hmm the main RAGreedy allocator is using the LiveRegMatrix and within that the LiveIntervalUnion classes to keep things performant. Maybe we can use those here too? In principle they should have even better runtime behavior than merging here...

Thank you for the clue, I'll take a look.

[vpykhtin]

Tried version with LiveIntervalUnion but it works slightly slower than first:

~4.2 vs ~3 seconds per pass on my testcase.

There're 48005704 overlap tests and 277607 interval joins in my testcase so it
looks like LiveRange::overlaps works faster than LiveIntervalUnion::Query::checkInterference.

Not sure which version to choose.

[vpykhtin]

• Rebased.

[arsenm]

llvm/lib/CodeGen/StackSlotColoring.cpp

101–103

I don't understand what this is doing, the parameter is just the allocator? Why can't you just normally construct a LiveIntervalUnion without placement new?

[vpykhtin]

llvm/lib/CodeGen/StackSlotColoring.cpp

101–103

ColorAssignmentInfo objects are placed in SmallVector so they need default constructor.

[arsenm]

LGTM

[vpykhtin]

Rebased before landing.

llvm/lib/CodeGen/StackSlotColoring.cpp

 //===- StackSlotColoring.cpp - Stack slot coloring pass. ------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the stack slot coloring pass.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/CodeGen/LiveInterval.h"
+#include "llvm/CodeGen/LiveIntervalUnion.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/LiveStacks.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstr.h"
class StackSlotColoring : public MachineFunctionPass {
     SmallVector<BitVector, 2> AllColors;
 
     // NextColor - Next "color" that's not yet used. This is per stack ID.
     SmallVector<int, 2> NextColors = { -1 };
 
     // UsedColors - "Colors" that have been assigned. This is per stack ID
     SmallVector<BitVector, 2> UsedColors;
 
+    // Join all intervals sharing one color into a single LiveIntervalUnion to
+    // speedup range overlap test.
+    class ColorAssignmentInfo {
+      // Single liverange (used to avoid creation of LiveIntervalUnion).
+      LiveInterval *SingleLI = nullptr;
+      // LiveIntervalUnion to perform overlap test.
+      LiveIntervalUnion *LIU = nullptr;
+      // LiveIntervalUnion has a parameter in its constructor so doing this
+      // dirty magic.
+      uint8_t LIUPad[sizeof(LiveIntervalUnion)];
+
+    public:
+      ~ColorAssignmentInfo() {
+        if (LIU)
+          LIU->~LiveIntervalUnion(); // Dirty magic again.
+      }
+
+      // Return true if LiveInterval overlaps with any
+      // intervals that have already been assigned to this color.
+      bool overlaps(LiveInterval *LI) const {
+        if (LIU)
+          return LiveIntervalUnion::Query(*LI, *LIU).checkInterference();
+        return SingleLI ? SingleLI->overlaps(*LI) : false;
+      }
+
+      // Add new LiveInterval to this color.
+      void add(LiveInterval *LI, LiveIntervalUnion::Allocator &Alloc) {
+        assert(!overlaps(LI));
+        if (LIU) {
+          LIU->unify(*LI, *LI);
+        } else if (SingleLI) {
+          LIU = new (LIUPad) LiveIntervalUnion(Alloc);
+          LIU->unify(*SingleLI, *SingleLI);
+          LIU->unify(*LI, *LI);
+          SingleLI = nullptr;
+        } else
+          SingleLI = LI;
+      }
+    };
+
+    LiveIntervalUnion::Allocator LIUAlloc;
+
     // Assignments - Color to intervals mapping.
-    SmallVector<SmallVector<LiveInterval*,4>, 16> Assignments;
+    SmallVector<ColorAssignmentInfo, 16> Assignments;
 
   public:
     static char ID; // Pass identification
 
     StackSlotColoring() : MachineFunctionPass(ID) {
       initializeStackSlotColoringPass(*PassRegistry::getPassRegistry());
     }
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
       AU.setPreservesCFG();
       AU.addRequired<SlotIndexes>();
       AU.addPreserved<SlotIndexes>();
       AU.addRequired<LiveStacks>();
       AU.addRequired<MachineBlockFrequencyInfo>();
       AU.addPreserved<MachineBlockFrequencyInfo>();
       AU.addPreservedID(MachineDominatorsID);
       MachineFunctionPass::getAnalysisUsage(AU);
     }
 
     bool runOnMachineFunction(MachineFunction &MF) override;
 
   private:
     void InitializeSlots();
     void ScanForSpillSlotRefs(MachineFunction &MF);
-    bool OverlapWithAssignments(LiveInterval *li, int Color) const;
     int ColorSlot(LiveInterval *li);
     bool ColorSlots(MachineFunction &MF);
     void RewriteInstruction(MachineInstr &MI, SmallVectorImpl<int> &SlotMapping,
                             MachineFunction &MF);
     bool RemoveDeadStores(MachineBasicBlock* MBB);
   };
 
 } // end anonymous namespace
void StackSlotColoring::InitializeSlots() {
 
   NextColors.resize(AllColors.size());
 
   // Get first "color".
   for (unsigned I = 0, E = AllColors.size(); I != E; ++I)
     NextColors[I] = AllColors[I].find_first();
 }
 
-/// OverlapWithAssignments - Return true if LiveInterval overlaps with any
-/// LiveIntervals that have already been assigned to the specified color.
-bool
-StackSlotColoring::OverlapWithAssignments(LiveInterval *li, int Color) const {
-  const SmallVectorImpl<LiveInterval *> &OtherLIs = Assignments[Color];
-  for (unsigned i = 0, e = OtherLIs.size(); i != e; ++i) {
-    LiveInterval *OtherLI = OtherLIs[i];
-    if (OtherLI->overlaps(*li))
-      return true;
-  }
-  return false;
-}
-
 /// ColorSlot - Assign a "color" (stack slot) to the specified stack slot.
 int StackSlotColoring::ColorSlot(LiveInterval *li) {
   int Color = -1;
   bool Share = false;
   int FI = Register::stackSlot2Index(li->reg());
   uint8_t StackID = MFI->getStackID(FI);
 
   if (!DisableSharing) {
 
     // Check if it's possible to reuse any of the used colors.
     Color = UsedColors[StackID].find_first();
     while (Color != -1) {
-      if (!OverlapWithAssignments(li, Color)) {
+      if (!Assignments[Color].overlaps(li)) {
         Share = true;
         ++NumEliminated;
         break;
       }
       Color = UsedColors[StackID].find_next(Color);
     }
   }
 
if (!Share) {
     Color = NextColors[StackID];
     UsedColors[StackID].set(Color);
     NextColors[StackID] = AllColors[StackID].find_next(NextColors[StackID]);
   }
 
   assert(MFI->getStackID(Color) == MFI->getStackID(FI));
 
   // Record the assignment.
-  Assignments[Color].push_back(li);
+  Assignments[Color].add(li, LIUAlloc);
   LLVM_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.
   Align Alignment = OrigAlignments[FI];
   if (!Share || Alignment > MFI->getObjectAlign(Color))
     MFI->setObjectAlignment(Color, Alignment);
bool StackSlotColoring::runOnMachineFunction(MachineFunction &MF) {
   SSIntervals.clear();
   for (unsigned i = 0, e = SSRefs.size(); i != e; ++i)
     SSRefs[i].clear();
   SSRefs.clear();
   OrigAlignments.clear();
   OrigSizes.clear();
   AllColors.clear();
   UsedColors.clear();
-  for (unsigned i = 0, e = Assignments.size(); i != e; ++i)
-    Assignments[i].clear();
   Assignments.clear();
 
   return Changed;
 }