[PHIElimination] Handle subranges in LiveInterval updates (#69429)

Add subrange tracking and handling for LiveIntervals during PHI
elimination.
This requires extending MachineBasicBlock::SplitCriticalEdge to also
update subrange intervals.

Author: perlfu

llvm/include/llvm/CodeGen/LiveInterval.h

@@ -520,8 +520,9 @@ namespace llvm {
         endIndex() < End.getBoundaryIndex();
     }
 
-    /// Remove the specified segment from this range.  Note that the segment
-    /// must be a single Segment in its entirety.
+    /// Remove the specified interval from this live range.
+    /// Does nothing if interval is not part of this live range.
+    /// Note that the interval must be within a single Segment in its entirety.
     void removeSegment(SlotIndex Start, SlotIndex End,
                        bool RemoveDeadValNo = false);
 

llvm/include/llvm/CodeGen/LiveIntervals.h

@@ -139,6 +139,13 @@ class VirtRegMap;
       return LI;
     }
 
+    /// Return an existing interval for \p Reg.
+    /// If \p Reg has no interval then this creates a new empty one instead.
+    /// Note: does not trigger interval computation.
+    LiveInterval &getOrCreateEmptyInterval(Register Reg) {
+      return hasInterval(Reg) ? getInterval(Reg) : createEmptyInterval(Reg);
+    }
+
     /// Interval removal.
     void removeInterval(Register Reg) {
       delete VirtRegIntervals[Reg];

llvm/lib/CodeGen/LiveInterval.cpp

@@ -563,13 +563,15 @@ VNInfo *LiveRange::extendInBlock(SlotIndex StartIdx, SlotIndex Kill) {
   return CalcLiveRangeUtilVector(this).extendInBlock(StartIdx, Kill);
 }
 
-/// Remove the specified segment from this range.  Note that the segment must
-/// be in a single Segment in its entirety.
 void LiveRange::removeSegment(SlotIndex Start, SlotIndex End,
                               bool RemoveDeadValNo) {
   // Find the Segment containing this span.
   iterator I = find(Start);
-  assert(I != end() && "Segment is not in range!");
+
+  // No Segment found, so nothing to do.
+  if (I == end())
+    return;
+
   assert(I->containsInterval(Start, End)
          && "Segment is not entirely in range!");
 

llvm/lib/CodeGen/LiveIntervals.cpp

@@ -862,7 +862,7 @@ float LiveIntervals::getSpillWeight(bool isDef, bool isUse,
 
 LiveRange::Segment
 LiveIntervals::addSegmentToEndOfBlock(Register Reg, MachineInstr &startInst) {
-  LiveInterval &Interval = createEmptyInterval(Reg);
+  LiveInterval &Interval = getOrCreateEmptyInterval(Reg);
   VNInfo *VN = Interval.getNextValue(
       SlotIndex(getInstructionIndex(startInst).getRegSlot()),
       getVNInfoAllocator());

llvm/lib/CodeGen/MachineBasicBlock.cpp

@@ -1283,6 +1283,8 @@ MachineBasicBlock *MachineBasicBlock::SplitCriticalEdge(
           assert(VNI &&
                  "PHI sources should be live out of their predecessors.");
           LI.addSegment(LiveInterval::Segment(StartIndex, EndIndex, VNI));
+          for (auto &SR : LI.subranges())
+            SR.addSegment(LiveInterval::Segment(StartIndex, EndIndex, VNI));
         }
       }
     }
@@ -1302,8 +1304,16 @@ MachineBasicBlock *MachineBasicBlock::SplitCriticalEdge(
         VNInfo *VNI = LI.getVNInfoAt(PrevIndex);
         assert(VNI && "LiveInterval should have VNInfo where it is live.");
         LI.addSegment(LiveInterval::Segment(StartIndex, EndIndex, VNI));
+        // Update subranges with live values
+        for (auto &SR : LI.subranges()) {
+          VNInfo *VNI = SR.getVNInfoAt(PrevIndex);
+          if (VNI)
+            SR.addSegment(LiveInterval::Segment(StartIndex, EndIndex, VNI));
+        }
       } else if (!isLiveOut && !isLastMBB) {
         LI.removeSegment(StartIndex, EndIndex);
+        for (auto &SR : LI.subranges())
+          SR.removeSegment(StartIndex, EndIndex);
       }
     }
 

llvm/lib/CodeGen/PHIElimination.cpp

@@ -389,7 +389,7 @@ void PHIElimination::LowerPHINode(MachineBasicBlock &MBB,
     if (IncomingReg) {
       // Add the region from the beginning of MBB to the copy instruction to
       // IncomingReg's live interval.
-      LiveInterval &IncomingLI = LIS->createEmptyInterval(IncomingReg);
+      LiveInterval &IncomingLI = LIS->getOrCreateEmptyInterval(IncomingReg);
       VNInfo *IncomingVNI = IncomingLI.getVNInfoAt(MBBStartIndex);
       if (!IncomingVNI)
         IncomingVNI = IncomingLI.getNextValue(MBBStartIndex,
@@ -400,24 +400,47 @@ void PHIElimination::LowerPHINode(MachineBasicBlock &MBB,
     }
 
     LiveInterval &DestLI = LIS->getInterval(DestReg);
-    assert(!DestLI.empty() && "PHIs should have nonempty LiveIntervals.");
-    if (DestLI.endIndex().isDead()) {
-      // A dead PHI's live range begins and ends at the start of the MBB, but
-      // the lowered copy, which will still be dead, needs to begin and end at
-      // the copy instruction.
-      VNInfo *OrigDestVNI = DestLI.getVNInfoAt(MBBStartIndex);
-      assert(OrigDestVNI && "PHI destination should be live at block entry.");
-      DestLI.removeSegment(MBBStartIndex, MBBStartIndex.getDeadSlot());
-      DestLI.createDeadDef(DestCopyIndex.getRegSlot(),
-                           LIS->getVNInfoAllocator());
-      DestLI.removeValNo(OrigDestVNI);
-    } else {
-      // Otherwise, remove the region from the beginning of MBB to the copy
-      // instruction from DestReg's live interval.
-      DestLI.removeSegment(MBBStartIndex, DestCopyIndex.getRegSlot());
-      VNInfo *DestVNI = DestLI.getVNInfoAt(DestCopyIndex.getRegSlot());
+    assert(!DestLI.empty() && "PHIs should have non-empty LiveIntervals.");
+
+    SlotIndex NewStart = DestCopyIndex.getRegSlot();
+
+    SmallVector<LiveRange *> ToUpdate({&DestLI});
+    for (auto &SR : DestLI.subranges())
+      ToUpdate.push_back(&SR);
+
+    for (auto LR : ToUpdate) {
+      auto DestSegment = LR->find(MBBStartIndex);
+      assert(DestSegment != LR->end() &&
+             "PHI destination must be live in block");
+
+      if (LR->endIndex().isDead()) {
+        // A dead PHI's live range begins and ends at the start of the MBB, but
+        // the lowered copy, which will still be dead, needs to begin and end at
+        // the copy instruction.
+        VNInfo *OrigDestVNI = LR->getVNInfoAt(DestSegment->start);
+        assert(OrigDestVNI && "PHI destination should be live at block entry.");
+        LR->removeSegment(DestSegment->start, DestSegment->start.getDeadSlot());
+        LR->createDeadDef(NewStart, LIS->getVNInfoAllocator());
+        LR->removeValNo(OrigDestVNI);
+        continue;
+      }
+
+      // Destination copies are not inserted in the same order as the PHI nodes
+      // they replace. Hence the start of the live range may need to be adjusted
+      // to match the actual slot index of the copy.
+      if (DestSegment->start > NewStart) {
+        VNInfo *VNI = LR->getVNInfoAt(DestSegment->start);
+        assert(VNI && "value should be defined for known segment");
+        LR->addSegment(
+            LiveInterval::Segment(NewStart, DestSegment->start, VNI));
+      } else if (DestSegment->start < NewStart) {
+        assert(DestSegment->start >= MBBStartIndex);
+        assert(DestSegment->end >= DestCopyIndex.getRegSlot());
+        LR->removeSegment(DestSegment->start, NewStart);
+      }
+      VNInfo *DestVNI = LR->getVNInfoAt(NewStart);
       assert(DestVNI && "PHI destination should be live at its definition.");
-      DestVNI->def = DestCopyIndex.getRegSlot();
+      DestVNI->def = NewStart;
     }
   }
 
@@ -612,6 +635,10 @@ void PHIElimination::LowerPHINode(MachineBasicBlock &MBB,
           SlotIndex LastUseIndex = LIS->getInstructionIndex(*KillInst);
           SrcLI.removeSegment(LastUseIndex.getRegSlot(),
                               LIS->getMBBEndIdx(&opBlock));
+          for (auto &SR : SrcLI.subranges()) {
+            SR.removeSegment(LastUseIndex.getRegSlot(),
+                             LIS->getMBBEndIdx(&opBlock));
+          }
         }
       }
     }

llvm/test/CodeGen/AMDGPU/split-mbb-lis-subrange.mir

@@ -1,15 +1,15 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 2
-# RUN: llc -march=amdgcn -mcpu=gfx1100 -verify-machineinstrs -run-pass liveintervals -o - %s | FileCheck -check-prefixes=GCN %s
+# RUN: llc -march=amdgcn -mcpu=gfx1100 -verify-machineinstrs -run-pass liveintervals,phi-node-elimination -o - %s | FileCheck -check-prefixes=GCN %s
 
-# This test simply checks that liveintervals pass verification.
+# This checks liveintervals pass verification and phi-node-elimination correctly preserves them.
 
 ---
 name: split_critical_edge_subranges
 tracksRegLiveness: true
 body:             |
   ; GCN-LABEL: name: split_critical_edge_subranges
   ; GCN: bb.0:
-  ; GCN-NEXT:   successors: %bb.3(0x40000000), %bb.1(0x40000000)
+  ; GCN-NEXT:   successors: %bb.5(0x40000000), %bb.1(0x40000000)
   ; GCN-NEXT: {{  $}}
   ; GCN-NEXT:   %coord:vreg_64 = IMPLICIT_DEF
   ; GCN-NEXT:   %desc:sgpr_256 = IMPLICIT_DEF
@@ -20,14 +20,22 @@ body:             |
   ; GCN-NEXT:   %s0a:vgpr_32 = COPY %load.sub0
   ; GCN-NEXT:   %s0b:vgpr_32 = COPY %load.sub1
   ; GCN-NEXT:   S_CMP_EQ_U32 %c0, %c1, implicit-def $scc
-  ; GCN-NEXT:   S_CBRANCH_SCC1 %bb.3, implicit $scc
-  ; GCN-NEXT:   S_BRANCH %bb.1
+  ; GCN-NEXT:   S_CBRANCH_SCC0 %bb.1, implicit $scc
+  ; GCN-NEXT: {{  $}}
+  ; GCN-NEXT: bb.5:
+  ; GCN-NEXT:   successors: %bb.3(0x80000000)
+  ; GCN-NEXT: {{  $}}
+  ; GCN-NEXT:   [[COPY:%[0-9]+]]:vgpr_32 = COPY %s0a
+  ; GCN-NEXT:   [[COPY1:%[0-9]+]]:vgpr_32 = COPY %s0b
+  ; GCN-NEXT:   S_BRANCH %bb.3
   ; GCN-NEXT: {{  $}}
   ; GCN-NEXT: bb.1:
   ; GCN-NEXT:   successors: %bb.3(0x80000000)
   ; GCN-NEXT: {{  $}}
   ; GCN-NEXT:   %s0c:vgpr_32 = V_ADD_F32_e64 0, %s0a, 0, %const, 0, 0, implicit $mode, implicit $exec
   ; GCN-NEXT:   %s0d:vgpr_32 = V_ADD_F32_e64 0, %s0b, 0, %const, 0, 0, implicit $mode, implicit $exec
+  ; GCN-NEXT:   [[COPY2:%[0-9]+]]:vgpr_32 = COPY %s0c
+  ; GCN-NEXT:   [[COPY3:%[0-9]+]]:vgpr_32 = COPY %s0d
   ; GCN-NEXT:   S_BRANCH %bb.3
   ; GCN-NEXT: {{  $}}
   ; GCN-NEXT: bb.2:
@@ -37,8 +45,8 @@ body:             |
   ; GCN-NEXT: bb.3:
   ; GCN-NEXT:   successors: %bb.4(0x80000000)
   ; GCN-NEXT: {{  $}}
-  ; GCN-NEXT:   %phi0:vgpr_32 = PHI %s0a, %bb.0, %s0c, %bb.1
-  ; GCN-NEXT:   %phi1:vgpr_32 = PHI %s0b, %bb.0, %s0d, %bb.1
+  ; GCN-NEXT:   %phi1:vgpr_32 = COPY [[COPY3]]
+  ; GCN-NEXT:   %phi0:vgpr_32 = COPY [[COPY2]]
   ; GCN-NEXT:   S_BRANCH %bb.4
   ; GCN-NEXT: {{  $}}
   ; GCN-NEXT: bb.4: