[AMDGPU] Improve liveness copying in si-optimize-exec-masking-pre-ra

[critson]

Summary

Further improve liveness copying for CC register post optimization
by mirroring live internal splits.
The fixes a bug in register allocation when CC register liveness
is extended across a branches instead of split.

[critson]

I am still working on suitably simple test case.

[arsenm]

llvm/lib/Target/AMDGPU/SIOptimizeExecMaskingPreRA.cpp

194–195

I'm increasingly worried about just ignoring physical registers. Is there a reason why we bother to try handling them here?

[critson]

llvm/lib/Target/AMDGPU/SIOptimizeExecMaskingPreRA.cpp

194–195

Thanks for pointing this out.
We don't need to handle physical register for SelReg anymore as do not optimize these anymore.
I've simplified the code.

[critson]

• Simplify code based on not handling physical registers

[critson]

• Add test, unfortunately this has to be a bug specific MIR test from an llvm-reduce

[arsenm]

llvm/lib/Target/AMDGPU/SIOptimizeExecMaskingPreRA.cpp

192

Don't need the &

llvm/test/CodeGen/AMDGPU/optimize-exec-mask-pre-ra-alloc-failure.mir

1 ↗

(On Diff #444138)

Don't need -mattr. Also I would prefer specifying the specific greedy run, and run to the end of virtregrewriter

12 ↗

(On Diff #444138)

Can drop register section (and then compact the numbers with -run-pass=none)

[critson]

llvm/test/CodeGen/AMDGPU/optimize-exec-mask-pre-ra-alloc-failure.mir

12 ↗

(On Diff #444138)

I have dropped all the register definitions that did not change the test result, 3 with preferred register values of $vcc remain.

[critson]

• Address reviewer comments

[cdevadas]

llvm/test/CodeGen/AMDGPU/optimize-exec-mask-pre-ra-alloc-failure.mir

1 ↗

(On Diff #444503)

How do the two '-stop-after' fields work here?

[arsenm]

llvm/test/CodeGen/AMDGPU/optimize-exec-mask-pre-ra-alloc-failure.mir

1 ↗

(On Diff #444138)

Should only have one stop-after. I think the second one wins

[arsenm]

llvm/lib/Target/AMDGPU/SIOptimizeExecMaskingPreRA.cpp

227

Typo defintion. The assert seems a bit redundant with the FindSegmentContaining below

[arsenm]

llvm/test/CodeGen/AMDGPU/optimize-exec-mask-pre-ra-alloc-failure.mir

26 ↗

(On Diff #444503)

Should drop the register types

[critson]

• Address reviewer comments

[critson]

llvm/lib/Target/AMDGPU/SIOptimizeExecMaskingPreRA.cpp

227

Fixed the typo.
I would prefer to the keep the assertion because while the code won't fail without segment, it suggest something is wrong which will impact code gen.

llvm/test/CodeGen/AMDGPU/optimize-exec-mask-pre-ra-alloc-failure.mir

26 ↗

(On Diff #444503)

I assume you mean the variable specification (p3) or (p4), I have removed these.
The register type information is required on most instructions for code gen.

[arsenm]

llvm/lib/Target/AMDGPU/SIOptimizeExecMaskingPreRA.cpp

227

I mean you could assert FindSegmentContaining isn't the end

llvm/lib/Target/AMDGPU/SIOptimizeExecMaskingPreRA.cpp

bool SIOptimizeExecMaskingPreRA::optimizeVcndVcmpPair(MachineBasicBlock &MBB) {
 
   Register CCReg = CC->getReg();
 
   // If there was a def between the select and the and, we would need to move it
   // to fold this.
   if (isDefBetween(*TRI, LIS, CCReg, *Sel, *And))
     return false;
 
+  // Cannot safely mirror live intervals with PHI nodes, so check for these
+  // before optimization.
+  SlotIndex SelIdx = LIS->getInstructionIndex(*Sel);
+  LiveInterval *SelLI = &LIS->getInterval(SelReg);
+  if (llvm::any_of(SelLI->vnis(),
+                    [&](const VNInfo *VNI) {
+                      return VNI->isPHIDef();
+                    }))
+    return false;
+
   // TODO: Guard against implicit def operands?
   LLVM_DEBUG(dbgs() << "Folding sequence:\n\t" << *Sel << '\t' << *Cmp << '\t'
                     << *And);
 
   MachineInstr *Andn2 =
       BuildMI(MBB, *And, And->getDebugLoc(), TII->get(Andn2Opc),
               And->getOperand(0).getReg())
           .addReg(ExecReg)
           .addReg(CCReg, getUndefRegState(CC->isUndef()), CC->getSubReg());
   MachineOperand &AndSCC = And->getOperand(3);
   assert(AndSCC.getReg() == AMDGPU::SCC);
   MachineOperand &Andn2SCC = Andn2->getOperand(3);
   assert(Andn2SCC.getReg() == AMDGPU::SCC);
   Andn2SCC.setIsDead(AndSCC.isDead());
 
   SlotIndex AndIdx = LIS->ReplaceMachineInstrInMaps(*And, *Andn2);
   And->eraseFromParent();
 
   LLVM_DEBUG(dbgs() << "=>\n\t" << *Andn2 << '\n');
 
-  SlotIndex CmpIdx = LIS->getInstructionIndex(*Cmp);
-  SlotIndex SelIdx = LIS->getInstructionIndex(*Sel);
-
-  LiveInterval *CmpLI =
-      CmpReg.isVirtual() ? &LIS->getInterval(CmpReg) : nullptr;
-  LiveInterval *SelLI =
-      SelReg.isVirtual() ? &LIS->getInterval(SelReg) : nullptr;
-
   // Update live intervals for CCReg before potentially removing CmpReg/SelReg,
   // and their associated liveness information.
+  SlotIndex CmpIdx = LIS->getInstructionIndex(*Cmp);
   if (CCReg.isVirtual()) {
-    // Note: this ignores that SelLI might have multiple internal values
-    // or splits and simply extends the live range to cover all cases
-    // where the result of the v_cndmask_b32 was live (e.g. loops).
-    // This could yield worse register allocation in rare edge cases.
-    SlotIndex EndIdx = AndIdx.getRegSlot();
-    if (SelLI && SelLI->endIndex() > EndIdx && SelLI->endIndex().isBlock())
-      EndIdx = SelLI->endIndex();
+    // Apply live ranges from SelLI to CCReg potentially matching splits
+    // and extending to loop boundaries.
+
+    auto applyLiveRanges = [&](LiveRange &Dst, VNInfo *VNI) {
+      // Copy live ranges from SelLI, adjusting start and end as required
+      assert(SelLI->getSegmentContaining(SelIdx.getRegSlot()) &&
+             "No live interval segment covering defintion?");
+      for (auto I = SelLI->FindSegmentContaining(SelIdx.getRegSlot());
+                I != SelLI->end(); ++I) {
+        SlotIndex Start = I->start < SelIdx.getRegSlot() ?
+                          SelIdx.getRegSlot() : I->start;
+        SlotIndex End = I->end < AndIdx.getRegSlot() || I->end.isBlock() ?
+                        I->end : AndIdx.getRegSlot();
+        Dst.addSegment(LiveRange::Segment(Start, End, VNI));
+      }
+      // If SelLI does not cover AndIdx (because Cmp killed Sel) then extend.
+      if (!SelLI->getSegmentContaining(AndIdx.getRegSlot()))
+        Dst.addSegment(LiveRange::Segment(CmpIdx.getRegSlot(), AndIdx.getRegSlot(), VNI));
+    };
 
     LiveInterval &CCLI = LIS->getInterval(CCReg);
     auto CCQ = CCLI.Query(SelIdx.getRegSlot());
-    if (CCQ.valueIn()) {
-      CCLI.addSegment(LiveRange::Segment(SelIdx.getRegSlot(),
-                                         EndIdx, CCQ.valueIn()));
-    }
+    if (CCQ.valueIn())
+      applyLiveRanges(CCLI, CCQ.valueIn());
 
     if (CC->getSubReg()) {
       LaneBitmask Mask = TRI->getSubRegIndexLaneMask(CC->getSubReg());
       BumpPtrAllocator &Allocator = LIS->getVNInfoAllocator();
       CCLI.refineSubRanges(
           Allocator, Mask,
           [=](LiveInterval::SubRange &SR) {
             auto CCQS = SR.Query(SelIdx.getRegSlot());
-            if (CCQS.valueIn()) {
-              SR.addSegment(LiveRange::Segment(
-                  SelIdx.getRegSlot(), EndIdx, CCQS.valueIn()));
-            }
+            if (CCQS.valueIn())
+              applyLiveRanges(SR, CCQS.valueIn());
           },
           *LIS->getSlotIndexes(), *TRI);
       CCLI.removeEmptySubRanges();
 
       SmallVector<LiveInterval *> SplitLIs;
       LIS->splitSeparateComponents(CCLI, SplitLIs);
     }
   } else
     LIS->removeAllRegUnitsForPhysReg(CCReg);
 
   // Try to remove compare. Cmp value should not used in between of cmp
   // and s_and_b64 if VCC or just unused if any other register.
-  if ((CmpReg.isVirtual() && CmpLI && CmpLI->Query(AndIdx.getRegSlot()).isKill()) ||
+  LiveInterval *CmpLI = CmpReg.isVirtual() ? &LIS->getInterval(CmpReg) : nullptr;
+  if ((CmpLI && CmpLI->Query(AndIdx.getRegSlot()).isKill()) ||
       (CmpReg == Register(CondReg) &&
        std::none_of(std::next(Cmp->getIterator()), Andn2->getIterator(),
                     [&](const MachineInstr &MI) {
                       return MI.readsRegister(CondReg, TRI);
                     }))) {
     LLVM_DEBUG(dbgs() << "Erasing: " << *Cmp << '\n');
     if (CmpLI)
       LIS->removeVRegDefAt(*CmpLI, CmpIdx.getRegSlot());
     LIS->RemoveMachineInstrFromMaps(*Cmp);
     Cmp->eraseFromParent();
 
     // Try to remove v_cndmask_b32.
-    if (SelLI) {
-      // Kill status must be checked before shrinking the live range.
-      bool IsKill = SelLI->Query(CmpIdx.getRegSlot()).isKill();
-      LIS->shrinkToUses(SelLI);
-      bool IsDead = SelLI->Query(SelIdx.getRegSlot()).isDeadDef();
-      if (MRI->use_nodbg_empty(SelReg) && (IsKill || IsDead)) {
-        LLVM_DEBUG(dbgs() << "Erasing: " << *Sel << '\n');
 
-        LIS->removeVRegDefAt(*SelLI, SelIdx.getRegSlot());
-        LIS->RemoveMachineInstrFromMaps(*Sel);
-        Sel->eraseFromParent();
-      }
-    }
-  }
 
   return true;
 }
 
 // Optimize sequence
 //    %dst = S_OR_SAVEEXEC %src
 //    ... instructions not modifying exec ...
 //    %tmp = S_AND $exec, %dst