Check subrange liveness at rematerialization

[rampitec]

Summary

LiveRangeEdit::allUsesAvailableAt checks that VNI at use is the same
as at the original use slot. However, the VNI can be the same while
a specific subrange needed for use can be dead at the new index.

This patch adds subrange liveness check if there is a subreg use.

Fixes: SWDEV-312810

[rampitec]

Added test with partially live subreg.

[arsenm]

llvm/lib/CodeGen/LiveRangeEdit.cpp

146–148

This isn't required?

llvm/test/CodeGen/AMDGPU/remat-dead-subreg.mir

4

Should include a brief description

[rampitec]

Added test description.

[rampitec]

llvm/lib/CodeGen/LiveRangeEdit.cpp

146–148

This is early exit. No need to continue iterating over the rest of the subranges if we have already checked all lanes we need. In case of wide register tuples this shall be faster.

[arsenm]

LGTM

llvm/lib/CodeGen/LiveRangeEdit.cpp

146–148

Probably should comment this is just for early exit (I'm also skeptical this helps anything)

[rampitec]

Added comment.

[rampitec]

llvm/lib/CodeGen/LiveRangeEdit.cpp

146–148

We have pretty wide tuples so the loop may have to process up to 32 subranges.

llvm/lib/CodeGen/LiveRangeEdit.cpp

for (const MachineOperand &MO : OrigMI->operands()) {
     // Don't allow rematerialization immediately after the original def.
     // It would be incorrect if OrigMI redefines the register.
     // See PR14098.
     if (SlotIndex::isSameInstr(OrigIdx, UseIdx))
       return false;
 
     if (OVNI != li.getVNInfoAt(UseIdx))
       return false;
+
+    // Check that subrange is live at UseIdx.
+    if (MO.getSubReg()) {
+      const TargetRegisterInfo *TRI = MRI.getTargetRegisterInfo();
+      LaneBitmask LM = TRI->getSubRegIndexLaneMask(MO.getSubReg());
+      for (LiveInterval::SubRange &SR : li.subranges()) {
+        if ((SR.LaneMask & LM).none())
+          continue;
+        if (!SR.liveAt(UseIdx))
+          return false;
+        // Early exit if all used lanes are checked. No need to continue.
+        LM &= ~SR.LaneMask;
+        if (LM.none())
+          break;
+      }
+    }
   }
   return true;
 }
 
 bool LiveRangeEdit::canRematerializeAt(Remat &RM, VNInfo *OrigVNI,
                                        SlotIndex UseIdx, bool cheapAsAMove) {
   assert(ScannedRemattable && "Call anyRematerializable first");
 

llvm/test/CodeGen/AMDGPU/remat-dead-subreg.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -march=amdgcn -start-before=greedy -stop-after=virtregrewriter -stress-regalloc=3 -verify-machineinstrs -o - %s | FileCheck -check-prefix=GCN %s
+
+# Check that subreg use is live at the point of materialization, not just the main range.
+# Do not rematerialize if used subreg is dead at a new index.
+
+---
+name:            dead_subreg
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    ; GCN-LABEL: name: dead_subreg
+    ; GCN: $m0 = IMPLICIT_DEF
+    ; GCN-NEXT: renamable $sgpr0_sgpr1 = S_MOV_B64 1, implicit $m0
+    ; GCN-NEXT: renamable $sgpr1 = S_MUL_I32 renamable $sgpr1, 3
+    ; GCN-NEXT: SI_SPILL_S32_SAVE killed renamable $sgpr1, %stack.0, implicit $exec, implicit $sp_reg :: (store (s32) into %stack.0, addrspace 5)
+    ; GCN-NEXT: renamable $sgpr2 = S_MOV_B32 2, implicit $m0
+    ; GCN-NEXT: renamable $sgpr1 = S_MOV_B32 3, implicit $m0
+    ; GCN-NEXT: dead %4:vgpr_32 = V_MOV_B32_e32 $sgpr0, implicit $exec, implicit killed $sgpr2, implicit killed $sgpr1
+    ; GCN-NEXT: renamable $sgpr1 = SI_SPILL_S32_RESTORE %stack.0, implicit $exec, implicit $sp_reg :: (load (s32) from %stack.0, addrspace 5)
+    ; GCN-NEXT: dead %5:vgpr_32 = V_MOV_B32_e32 killed $sgpr1, implicit $exec
+    ; GCN-NEXT: S_NOP 0, implicit killed renamable $sgpr0
+    $m0 = IMPLICIT_DEF
+    %0:sreg_64_xexec = S_MOV_B64 1, implicit $m0
+    %1:sreg_32 = S_MUL_I32 %0.sub1:sreg_64_xexec, 3
+    %2:sreg_32 = S_MOV_B32 2, implicit $m0
+    %3:sreg_32 = S_MOV_B32 3, implicit $m0
+    %4:vgpr_32 = V_MOV_B32_e32 %0.sub0:sreg_64_xexec, implicit $exec, implicit %2, implicit %3
+    %5:vgpr_32 = V_MOV_B32_e32 %1:sreg_32, implicit $exec
+    S_NOP 0, implicit %0.sub0
+...
+---
+name:            live_subreg
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    ; GCN-LABEL: name: live_subreg
+    ; GCN: $m0 = IMPLICIT_DEF
+    ; GCN-NEXT: renamable $sgpr4_sgpr5 = S_MOV_B64 1, implicit $m0
+    ; GCN-NEXT: renamable $sgpr2 = S_MOV_B32 2, implicit $m0
+    ; GCN-NEXT: SI_SPILL_S32_SAVE killed renamable $sgpr2, %stack.0, implicit $exec, implicit $sp_reg :: (store (s32) into %stack.0, addrspace 5)
+    ; GCN-NEXT: renamable $sgpr2 = S_MOV_B32 3, implicit $m0
+    ; GCN-NEXT: renamable $sgpr0 = SI_SPILL_S32_RESTORE %stack.0, implicit $exec, implicit $sp_reg :: (load (s32) from %stack.0, addrspace 5)
+    ; GCN-NEXT: dead %4:vgpr_32 = V_MOV_B32_e32 $sgpr4, implicit $exec, implicit killed $sgpr0, implicit killed $sgpr2
+    ; GCN-NEXT: renamable $sgpr0 = S_MUL_I32 renamable $sgpr5, 3
+    ; GCN-NEXT: dead %5:vgpr_32 = V_MOV_B32_e32 killed $sgpr0, implicit $exec
+    ; GCN-NEXT: S_NOP 0, implicit killed renamable $sgpr5
+    $m0 = IMPLICIT_DEF
+    %0:sreg_64_xexec = S_MOV_B64 1, implicit $m0
+    %1:sreg_32 = S_MUL_I32 %0.sub1:sreg_64_xexec, 3
+    %2:sreg_32 = S_MOV_B32 2, implicit $m0
+    %3:sreg_32 = S_MOV_B32 3, implicit $m0
+    %4:vgpr_32 = V_MOV_B32_e32 %0.sub0:sreg_64_xexec, implicit $exec, implicit %2, implicit %3
+    %5:vgpr_32 = V_MOV_B32_e32 %1:sreg_32, implicit $exec
+    S_NOP 0, implicit %0.sub1
+...
+---
+name:            partially_live_wide_subreg
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    ; GCN-LABEL: name: partially_live_wide_subreg
+    ; GCN: renamable $sgpr0_sgpr1 = IMPLICIT_DEF
+    ; GCN-NEXT: renamable $sgpr8_sgpr9_sgpr10_sgpr11 = S_LOAD_DWORDX4_IMM killed renamable $sgpr0_sgpr1, 1, 0
+    ; GCN-NEXT: renamable $sgpr4_sgpr5 = S_MOV_B64 renamable $sgpr10_sgpr11
+    ; GCN-NEXT: SI_SPILL_S64_SAVE killed renamable $sgpr4_sgpr5, %stack.0, implicit $exec, implicit $sp_reg :: (store (s64) into %stack.0, align 4, addrspace 5)
+    ; GCN-NEXT: renamable $sgpr4_sgpr5 = S_MOV_B64 3, implicit $m0
+    ; GCN-NEXT: renamable $sgpr0_sgpr1 = S_MOV_B64 2, implicit $m0
+    ; GCN-NEXT: dead %5:vgpr_32 = V_MOV_B32_e32 $sgpr8, implicit $exec, implicit killed $sgpr0_sgpr1, implicit killed $sgpr4_sgpr5
+    ; GCN-NEXT: renamable $sgpr4_sgpr5 = SI_SPILL_S64_RESTORE %stack.0, implicit $exec, implicit $sp_reg :: (load (s64) from %stack.0, align 4, addrspace 5)
+    ; GCN-NEXT: dead %6:vreg_64 = V_MOV_B64_PSEUDO killed $sgpr4_sgpr5, implicit $exec
+    ; GCN-NEXT: S_NOP 0, implicit killed renamable $sgpr8, implicit renamable $sgpr11
+    %0:sreg_64 = IMPLICIT_DEF
+    %1:sgpr_128 = S_LOAD_DWORDX4_IMM %0, 1, 0
+    %2:sreg_64 = S_MOV_B64 %1.sub2_sub3
+    %3:sreg_64 = S_MOV_B64 2, implicit $m0
+    %4:sreg_64 = S_MOV_B64 3, implicit $m0
+    %5:vgpr_32 = V_MOV_B32_e32 %1.sub0, implicit $exec, implicit %3, implicit %4
+    %6:vreg_64 = V_MOV_B64_PSEUDO %2, implicit $exec
+    S_NOP 0, implicit %1.sub0, implicit %1.sub3
+...