[PHIElimination] Handle subranges in LiveInterval updates

[critson]

Summary

Add handling for subrange updates in LiveInterval preservation.
This requires extending MachineBasicBlock::SplitCriticalEdge
to also update subrange intervals.

[arsenm]

llvm/lib/CodeGen/LiveIntervals.cpp

865–866

Probably should move to getOrCreateEmptyInterval helper

[critson]

• Add getOrCreateEmptyInterval helper to LiveIntervals

[vitalybuka]

There is a memory leak likely from this patch https://lab.llvm.org/buildbot/#/builders/5/builds/36572
Unfortunately build was broken by unrelated patch, not it's fixed, and we see the leak.

Please reply if you have a quick fix, or I will try to revert it.

[critson]

@vitalybuka - thanks for reverting, I will investigate.

[critson]

• Add additional fixes to SplitCriticalEdge

[critson]

llvm/lib/CodeGen/MachineBasicBlock.cpp

1169

@arsenm / @foad - can you let me know if you are happy with this fix to obviously wrong use after free code that existed previously?

Basically the old code takes pointers to terminators that will be removed by TII->removeBranch() then dereferences them later to update the maps.

Other possible ways I can see to this:

1. Provide a way to remove pointers from maps directly, without dereferencing (this is fairly trivial to implement).
2. Add a class variable which stores a pointer to SlotIndexes to update, and make erase() use this to call removeMachineInstrFromMaps(). Manipulate the new variable during SplitCriticalEdge to catch any erase() calls.

The problems are:
(1) implies SlotIndexes should be able to retain pointers to deleted data, which is a bad idea.
(2) adds (minor) overhead to every erase() call, just to support this edge case.

[foad]

llvm/lib/CodeGen/MachineBasicBlock.cpp

1169

I am not familiar with this code.

Could you pass an optional SlotIndexes * into updateTerminator, and have it do the updates itself?

[arsenm]

llvm/lib/CodeGen/MachineBasicBlock.cpp

1169

I'd expect these updates to happen at the time the actual mutation happens, so have updateTerminator doing t

[critson]

llvm/lib/CodeGen/MachineBasicBlock.cpp

1169

Simply passing SlotIndexes to updateTerminator is insufficient.
The actual instruction manipulation is performed by TargetInstrInfo::removeBranch so the instructions removed are target dependent.

To truly pass SlotIndexes all the way down is an all targets change.
I don't think it is unreasonable to pursue such a change, but it'll need to be a separate patch with many more reviewers.

[foad]

llvm/lib/CodeGen/MachineBasicBlock.cpp

1169

Hmm, I think it's probably still the right way to do it though.

[critson]

llvm/lib/CodeGen/MachineBasicBlock.cpp

1169

I have put in a PR for fixing the SlotIndexes issue: https://github.com/llvm/llvm-project/pull/66188

llvm/include/llvm/CodeGen/LiveIntervals.h

public:
     }
 
     LiveInterval &createAndComputeVirtRegInterval(Register Reg) {
       LiveInterval &LI = createEmptyInterval(Reg);
       computeVirtRegInterval(LI);
       return LI;
     }
 
+    LiveInterval &getOrCreateEmptyInterval(Register Reg) {
+      return hasInterval(Reg) ? getInterval(Reg) : createEmptyInterval(Reg);
+    }
+
     /// Interval removal.
     void removeInterval(Register Reg) {
       delete VirtRegIntervals[Reg];
       VirtRegIntervals[Reg] = nullptr;
     }
 
     /// Given a register and an instruction, adds a live segment from that
     /// instruction to the end of its MBB.

llvm/lib/CodeGen/LiveIntervals.cpp

 float LiveIntervals::getSpillWeight(bool isDef, bool isUse,
                                     const MachineBlockFrequencyInfo *MBFI,
                                     const MachineBasicBlock *MBB) {
   return (isDef + isUse) * MBFI->getBlockFreqRelativeToEntryBlock(MBB);
 }
 
 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());
   LiveRange::Segment S(SlotIndex(getInstructionIndex(startInst).getRegSlot()),
                        getMBBEndIdx(startInst.getParent()), VN);
   Interval.addSegment(S);
 
   return S;
 }

llvm/lib/CodeGen/MachineBasicBlock.cpp

for (MachineInstr &MI :
         if (!is_contained(UsedRegs, Reg))
           UsedRegs.push_back(Reg);
       }
     }
   }
 
   ReplaceUsesOfBlockWith(Succ, NMBB);
 
-  // If updateTerminator() removes instructions, we need to remove them from
-  // SlotIndexes.
+  if (!ChangedIndirectJump) {
+    // If updateTerminator() removes instructions they need to be removed from
+    // SlotIndexes, but we do not know which will be removed, so we have to
+    // remove all pre-emptively as their pointers will no longer be valid.
     SmallVector<MachineInstr*, 4> Terminators;
     if (Indexes) {
       for (MachineInstr &MI :
            llvm::make_range(getFirstInstrTerminator(), instr_end()))
-      Terminators.push_back(&MI);
+        Indexes->removeMachineInstrFromMaps(MI);
     }
 
     // Since we replaced all uses of Succ with NMBB, that should also be treated
     // as the fallthrough successor
     if (Succ == PrevFallthrough)
       PrevFallthrough = NMBB;
 
-  if (!ChangedIndirectJump)
     updateTerminator(PrevFallthrough);
 
+    // Reinsert terminators into maps.
     if (Indexes) {
-    SmallVector<MachineInstr*, 4> NewTerminators;
       for (MachineInstr &MI :
            llvm::make_range(getFirstInstrTerminator(), instr_end()))
-      NewTerminators.push_back(&MI);
-
-    for (MachineInstr *Terminator : Terminators) {
-      if (!is_contained(NewTerminators, Terminator))
-        Indexes->removeMachineInstrFromMaps(*Terminator);
+        Indexes->insertMachineInstrInMaps(MI);
     }
   }
 
   // Insert unconditional "jump Succ" instruction in NMBB if necessary.
   NMBB->addSuccessor(Succ);
   if (!NMBB->isLayoutSuccessor(Succ)) {
     SmallVector<MachineOperand, 4> Cond;
     const TargetInstrInfo *TII = getParent()->getSubtarget().getInstrInfo();
for (MachineBasicBlock::instr_iterator
           if (MO.isUndef())
             continue;
 
           LiveInterval &LI = LIS->getInterval(Reg);
           VNInfo *VNI = LI.getVNInfoAt(PrevIndex);
           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));
         }
       }
     }
 
     MachineRegisterInfo *MRI = &getParent()->getRegInfo();
     for (unsigned i = 0, e = MRI->getNumVirtRegs(); i != e; ++i) {
       Register Reg = Register::index2VirtReg(i);
       if (PHISrcRegs.count(Reg) || !LIS->hasInterval(Reg))
         continue;
 
       LiveInterval &LI = LIS->getInterval(Reg);
       if (!LI.liveAt(PrevIndex))
         continue;
 
       bool isLiveOut = LI.liveAt(LIS->getMBBStartIdx(Succ));
       if (isLiveOut && isLastMBB) {
         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()) {
+          if (SR.overlaps(StartIndex, EndIndex))
+            SR.removeSegment(StartIndex, EndIndex);
+        }
       }
     }
 
     // Update all intervals for registers whose uses may have been modified by
     // updateTerminator().
     LIS->repairIntervalsInRange(this, getFirstTerminator(), end(), UsedRegs);
   }
 

llvm/lib/CodeGen/PHIElimination.cpp

INITIALIZE_PASS_BEGIN(PHIElimination, DEBUG_TYPE,
                       "Eliminate PHI nodes for register allocation",
                       false, false)
 INITIALIZE_PASS_DEPENDENCY(LiveVariables)
 INITIALIZE_PASS_END(PHIElimination, DEBUG_TYPE,
                     "Eliminate PHI nodes for register allocation", false, false)
 
 void PHIElimination::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addUsedIfAvailable<LiveVariables>();
+  AU.addUsedIfAvailable<LiveIntervals>();
   AU.addPreserved<LiveVariables>();
   AU.addPreserved<SlotIndexes>();
   AU.addPreserved<LiveIntervals>();
   AU.addPreserved<MachineDominatorTree>();
   AU.addPreserved<MachineLoopInfo>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
void PHIElimination::LowerPHINode(MachineBasicBlock &MBB,
   // Update LiveIntervals for the new copy or implicit def.
   if (LIS) {
     SlotIndex DestCopyIndex = LIS->InsertMachineInstrInMaps(*PHICopy);
 
     SlotIndex MBBStartIndex = LIS->getMBBStartIdx(&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,
                                               LIS->getVNInfoAllocator());
       IncomingLI.addSegment(LiveInterval::Segment(MBBStartIndex,
                                                   DestCopyIndex.getRegSlot(),
                                                   IncomingVNI));
     }
 
     LiveInterval &DestLI = LIS->getInterval(DestReg);
-    assert(!DestLI.empty() && "PHIs should have nonempty LiveIntervals.");
-    if (DestLI.endIndex().isDead()) {
+    assert(!DestLI.empty() && "PHIs should have non-empty LiveIntervals.");
+
+    SlotIndex NewStart = DestCopyIndex.getRegSlot();
+
+    SmallVector<LiveRange*> ToUpdate;
+    ToUpdate.push_back(&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 = DestLI.getVNInfoAt(MBBStartIndex);
+        VNInfo *OrigDestVNI = LR->getVNInfoAt(DestSegment->start);
         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 {
+        LR->removeSegment(DestSegment->start, DestSegment->start.getDeadSlot());
+        LR->createDeadDef(NewStart, LIS->getVNInfoAllocator());
+        LR->removeValNo(OrigDestVNI);
+        continue;
+      }
+
+      if (DestSegment->start > NewStart) {
+        // With a single PHI removed from block the index of the copy may be
+        // lower than the original PHI. Extend live range backward to cover
+        // the copy.
+        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) {
         // 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(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;
     }
   }
 
   // Adjust the VRegPHIUseCount map to account for the removal of this PHI node.
   for (unsigned i = 1; i != MPhi->getNumOperands(); i += 2) {
     if (!MPhi->getOperand(i).isUndef()) {
       --VRegPHIUseCount[BBVRegPair(
           MPhi->getOperand(i + 1).getMBB()->getNumber(),
if (LIS) {
             }
           }
           assert(KillInst->readsRegister(SrcReg) &&
                  "Cannot find kill instruction");
 
           SlotIndex LastUseIndex = LIS->getInstructionIndex(*KillInst);
           SrcLI.removeSegment(LastUseIndex.getRegSlot(),
                               LIS->getMBBEndIdx(&opBlock));
+          for (auto &SR : SrcLI.subranges()) {
+            SR.removeSegment(LastUseIndex.getRegSlot(),
+                                LIS->getMBBEndIdx(&opBlock));
+          }
         }
       }
     }
   }
 
   // Really delete the PHI instruction now, if it is not in the LoweredPHIs map.
   if (reusedIncoming || !IncomingReg) {
     if (LIS)

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

 # 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
   ; GCN-NEXT:   %c0:sreg_32 = IMPLICIT_DEF
   ; GCN-NEXT:   %c1:sreg_32 = IMPLICIT_DEF
   ; GCN-NEXT:   %const:vgpr_32 = IMPLICIT_DEF
   ; GCN-NEXT:   %load:vreg_64 = IMAGE_LOAD_V2_V2_gfx11 %coord, %desc, 3, 1, -1, 0, 0, 0, 0, 0, 0, implicit $exec :: (dereferenceable load (s64), align 16, addrspace 4)
   ; 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:
   ; GCN-NEXT:   S_NOP 0
   ; GCN-NEXT:   S_ENDPGM 0
   ; GCN-NEXT: {{  $}}
   ; 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:
   ; GCN-NEXT:   S_ENDPGM 0, implicit %phi0, implicit %phi1
   bb.0:
     %coord:vreg_64 = IMPLICIT_DEF
     %desc:sgpr_256 = IMPLICIT_DEF
     %c0:sreg_32 = IMPLICIT_DEF