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

[Valery Pykhtin via llvm-commits]

Author: Valery Pykhtin
Date: 2023-04-06T07:23:45+02:00
New Revision: e09b33feec1991876a3093d86348818df8b8edae

URL: https://github.com/llvm/llvm-project/commit/e09b33feec1991876a3093d86348818df8b8edae
DIFF: https://github.com/llvm/llvm-project/commit/e09b33feec1991876a3093d86348818df8b8edae.diff

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

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.

Reviewed By: arsenm

Differential Revision: https://reviews.llvm.org/D146057

Added: 
    

Modified: 
    llvm/lib/CodeGen/StackSlotColoring.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/StackSlotColoring.cpp b/llvm/lib/CodeGen/StackSlotColoring.cpp
index b8c750688914f..d6c401c3ce33f 100644
--- a/llvm/lib/CodeGen/StackSlotColoring.cpp
+++ b/llvm/lib/CodeGen/StackSlotColoring.cpp
@@ -14,6 +14,7 @@
 #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"
@@ -90,8 +91,50 @@ namespace {
     // 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
@@ -116,7 +159,6 @@ namespace {
   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,
@@ -247,19 +289,6 @@ void StackSlotColoring::InitializeSlots() {
     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;
@@ -272,7 +301,7 @@ int StackSlotColoring::ColorSlot(LiveInterval *li) {
     // 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;
@@ -298,7 +327,7 @@ int StackSlotColoring::ColorSlot(LiveInterval *li) {
   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
@@ -515,8 +544,6 @@ bool StackSlotColoring::runOnMachineFunction(MachineFunction &MF) {
   OrigSizes.clear();
   AllColors.clear();
   UsedColors.clear();
-  for (unsigned i = 0, e = Assignments.size(); i != e; ++i)
-    Assignments[i].clear();
   Assignments.clear();
 
   return Changed;