[AMDGPU] Add optional tied-op for wwm-register's epilog spill restore

[cdevadas]

Summary

The COPY inserted in the epilog block before return instruction as part
of ABI lowering, can get optimized during machine copy propagation if
the same register is used earlier in a wwm operation that demands the
prolog/epilog wwm-spill store/restore to preserve its inactive lanes.
With the spill restore in the epilog, the preceding COPY appears to be
dead during machine-cp. To avoid it, mark the same register as a tied-op
in the spill restore instruction to ensure a usage for the COPY.

[arsenm]

I noticed in the context we still have some forward scavenging. Those should have been replaced when we switched to using PEI in reverse.

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

1298

Exact register equality is almost always wrong

1669

Should replace this with a liveness query instead

[cdevadas]

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

1669

Using LiveRegs here won't help. The CSRs (including wwm-regs) are marked LiveOut by initLiveRegs() for an epilog block. In that sense, the liveness query will always hold true and we add the tied-op unconditionally for all spill reloads.
These registers are, in fact, liveout only after the insertion of the spill-restore. But we are in the process of inserting one, and we need to look one level backward to get their actual liveness. LiveRegs in its present form won't be helpful for that.

The intention of this patch is to tie the return value register to its own spill restore (if have one) to avoid an incorrect optimization currently performed on it due to the presence of the wwm-restore inserted afterward.
Instead of walking through the operands of RETURN instruction, I would prefer if we have a utility function that returns a regmask or the list of return value registers allocated for this function. But not sure such a thing exists.

[arsenm]

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

1669

I'm having trouble following this. The live-outness is already represented by the $vgpr0 use on the return instruction. You're special casing looking for the return instruction in the low level return utility. Is this just to filter out spills in other contexts, and that's why the liveness query doesn't work? If you performed the liveness query in buildEpilogRestore and passed in that you needed the tied operand, would that let you use the liveness query?

llvm/test/CodeGen/AMDGPU/tied-op-for-wwm-scratch-reg-spill-restore.mir

37

Can you add some additional tests with returned register tuples, and one where the def'd register is in the CSR range?

Also, test where the spill isn't used as the return value

[cdevadas]

Used MachineInstr::readsRegister to check if RETURN instruction reads the spill-reg.
Also, added more tests.

[arsenm]

LGTM with nit

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

1669

MBB.isReturnBlock is redundant with MI->isReturn. isReturnBlock just checks if the last instruction is a return

[cdevadas]

Removed the isReturn() blk check. Also, added a validity check on the iterator.

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

 void SIRegisterInfo::buildSpillLoadStore(
     MachineBasicBlock &MBB, MachineBasicBlock::iterator MI, const DebugLoc &DL,
     unsigned LoadStoreOp, int Index, Register ValueReg, bool IsKill,
     MCRegister ScratchOffsetReg, int64_t InstOffset, MachineMemOperand *MMO,
     RegScavenger *RS, LivePhysRegs *LiveRegs) const {
   assert((!RS || !LiveRegs) && "Only RS or LiveRegs can be set but not both");
 
   MachineFunction *MF = MBB.getParent();
   const SIInstrInfo *TII = ST.getInstrInfo();
   const MachineFrameInfo &MFI = MF->getFrameInfo();
   const SIMachineFunctionInfo *FuncInfo = MF->getInfo<SIMachineFunctionInfo>();
 
   const MCInstrDesc *Desc = &TII->get(LoadStoreOp);
   bool IsStore = Desc->mayStore();
   bool IsFlat = TII->isFLATScratch(LoadStoreOp);
 
   bool CanClobberSCC = false;
if (!IsStore && IsAGPR && TmpIntermediateVGPR != AMDGPU::NoRegister) {
       MIB = BuildMI(MBB, MI, DL, TII->get(AMDGPU::V_ACCVGPR_WRITE_B32_e64),
                     FinalReg)
                 .addReg(TmpIntermediateVGPR, RegState::Kill);
       MIB->setAsmPrinterFlag(MachineInstr::ReloadReuse);
     }
 
     if (NeedSuperRegImpOperand && (IsFirstSubReg || IsLastSubReg))
       MIB.addReg(ValueReg, RegState::Implicit | SrcDstRegState);
+
+    // The epilog restore of a wwm-scratch register can cause undesired
+    // optimization during machine-cp post PrologEpilogInserter if the same
+    // register was assigned for return value ABI lowering with a COPY
+    // instruction. As given below, with the epilog reload, the earlier COPY
+    // appeared to be dead during machine-cp.
+    // ...
+    // v0 in WWM operation, needs the WWM spill at prolog/epilog.
+    // $vgpr0 = V_WRITELANE_B32 $sgpr20, 0, $vgpr0
+    // ...
+    // Epilog block:
+    // $vgpr0 = COPY $vgpr1 // outgoing value moved to v0
+    // ...
+    // WWM spill restore to preserve the inactive lanes of v0.
+    // $sgpr4_sgpr5 = S_XOR_SAVEEXEC_B64 -1
+    // $vgpr0 = BUFFER_LOAD $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 0, 0, 0
+    // $exec = S_MOV_B64 killed $sgpr4_sgpr5
+    // ...
+    // SI_RETURN implicit $vgpr0
+    // ...
+    // To fix it, mark the same reg as a tied op for such restore instructions
+    // so that it marks a usage for the preceding COPY.
+    if (!IsStore && MI != MBB.end() && MI->isReturn() &&
+        MI->readsRegister(SubReg, this)) {
+      MIB.addReg(SubReg, RegState::Implicit);
+      MIB->tieOperands(0, MIB->getNumOperands() - 1);
+    }
   }
 
   if (ScratchOffsetRegDelta != 0) {
     // Subtract the offset we added to the ScratchOffset register.
     BuildMI(MBB, MI, DL, TII->get(AMDGPU::S_ADD_I32), SOffset)
         .addReg(SOffset)
         .addImm(-ScratchOffsetRegDelta);
   }

llvm/test/CodeGen/AMDGPU/preserve-only-inactive-lane.mir

bb.0:
     ; GCN-NEXT: $sgpr4_sgpr5 = S_XOR_SAVEEXEC_B64 -1, implicit-def $exec, implicit-def dead $scc, implicit $exec
     ; GCN-NEXT: BUFFER_STORE_DWORD_OFFSET $vgpr0, $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 0, 0, 0, implicit $exec :: (store (s32) into %stack.0, addrspace 5)
     ; GCN-NEXT: $exec = S_MOV_B64 killed $sgpr4_sgpr5
     ; GCN-NEXT: renamable $vgpr0 = V_WRITELANE_B32 $sgpr35, 0, killed $vgpr0
     ; GCN-NEXT: $sgpr35 = S_MOV_B32 5
     ; GCN-NEXT: $sgpr35 = V_READLANE_B32 $vgpr0, 0
     ; GCN-NEXT: renamable $vgpr0 = V_MOV_B32_e32 10, implicit $exec
     ; GCN-NEXT: $sgpr4_sgpr5 = S_XOR_SAVEEXEC_B64 -1, implicit-def $exec, implicit-def dead $scc, implicit $exec
-    ; GCN-NEXT: $vgpr0 = BUFFER_LOAD_DWORD_OFFSET $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 0, 0, 0, implicit $exec :: (load (s32) from %stack.0, addrspace 5)
+    ; GCN-NEXT: $vgpr0 = BUFFER_LOAD_DWORD_OFFSET $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 0, 0, 0, implicit $exec, implicit $vgpr0(tied-def 0) :: (load (s32) from %stack.0, addrspace 5)
     ; GCN-NEXT: $exec = S_MOV_B64 killed $sgpr4_sgpr5
     ; GCN-NEXT: S_SETPC_B64_return killed renamable $sgpr30_sgpr31, implicit $vgpr0
     renamable $vgpr0 = V_WRITELANE_B32 $sgpr35, 0, killed $vgpr0
     $sgpr35 = S_MOV_B32 5
     $sgpr35 = V_READLANE_B32 $vgpr0, 0
     renamable $vgpr0 = V_MOV_B32_e32 10, implicit $exec
     S_SETPC_B64_return killed renamable $sgpr30_sgpr31, implicit $vgpr0
 ...

llvm/test/CodeGen/AMDGPU/tied-op-for-wwm-scratch-reg-spill-restore.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx906 -run-pass=prologepilog,machine-cp -verify-machineinstrs %s -o - | FileCheck -check-prefix=GCN %s
+
+# The COPY that moves the return value to VGPR0 should not be removed during machine-cp. The spill restore of the same register that follows,
+# meant to only reload its inactive lanes. By marking the reg itself as the tied-op in the spill reload prevents the undesired optimization.
+
+---
+name:            wwm_scratch_reg_spill_reload_of_outgoing_reg
+tracksRegLiveness: true
+machineFunctionInfo:
+  wwmReservedRegs: ['$vgpr0']
+  isEntryFunction: false
+  scratchRSrcReg:  '$sgpr0_sgpr1_sgpr2_sgpr3'
+  stackPtrOffsetReg: '$sgpr32'
+  frameOffsetReg: '$sgpr33'
+body:             |
+  bb.0:
+    liveins: $sgpr20, $vgpr1
+    ; GCN-LABEL: name: wwm_scratch_reg_spill_reload_of_outgoing_reg
+    ; GCN: liveins: $sgpr20, $vgpr0, $vgpr1
+    ; GCN-NEXT: {{  $}}
+    ; GCN-NEXT: $sgpr4_sgpr5 = S_XOR_SAVEEXEC_B64 -1, implicit-def $exec, implicit-def dead $scc, implicit $exec
+    ; GCN-NEXT: BUFFER_STORE_DWORD_OFFSET $vgpr0, $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 0, 0, 0, implicit $exec :: (store (s32) into %stack.0, addrspace 5)
+    ; GCN-NEXT: $exec = S_MOV_B64 killed $sgpr4_sgpr5
+    ; GCN-NEXT: $vgpr0 = IMPLICIT_DEF
+    ; GCN-NEXT: $vgpr0 = V_WRITELANE_B32 killed $sgpr20, 0, $vgpr0
+    ; GCN-NEXT: $vgpr0 = COPY killed renamable $vgpr1, implicit $exec
+    ; GCN-NEXT: $sgpr4_sgpr5 = S_XOR_SAVEEXEC_B64 -1, implicit-def $exec, implicit-def dead $scc, implicit $exec
+    ; GCN-NEXT: $vgpr0 = BUFFER_LOAD_DWORD_OFFSET $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 0, 0, 0, implicit $exec, implicit $vgpr0(tied-def 0) :: (load (s32) from %stack.0, addrspace 5)
+    ; GCN-NEXT: $exec = S_MOV_B64 killed $sgpr4_sgpr5
+    ; GCN-NEXT: SI_RETURN implicit $vgpr0
+    $vgpr0 = IMPLICIT_DEF
+    $vgpr0 = V_WRITELANE_B32 killed $sgpr20, 0, $vgpr0
+    $vgpr0 = COPY killed renamable $vgpr1, implicit $exec
+    SI_RETURN implicit $vgpr0
+...
+
+# The reload of vgpr0 require the tied-op as it is a subreg in the outgoing tuple register vgpr0_vgpr1.
+# The vgpr2 doesn't need the tied-op in the reload as it isn't holding any return value.
+---
+name:            wwm_scratch_reg_spill_reload_of_outgoing_tuple_subreg
+tracksRegLiveness: true
+machineFunctionInfo:
+  wwmReservedRegs: ['$vgpr0', '$vgpr2']
+  isEntryFunction: false
+  scratchRSrcReg:  '$sgpr0_sgpr1_sgpr2_sgpr3'
+  stackPtrOffsetReg: '$sgpr32'
+  frameOffsetReg: '$sgpr33'
+body:             |
+  bb.0:
+    liveins: $sgpr20, $sgpr21, $vgpr1
+    ; GCN-LABEL: name: wwm_scratch_reg_spill_reload_of_outgoing_tuple_subreg
+    ; GCN: liveins: $sgpr20, $sgpr21, $vgpr0, $vgpr1, $vgpr2
+    ; GCN-NEXT: {{  $}}
+    ; GCN-NEXT: $sgpr4_sgpr5 = S_XOR_SAVEEXEC_B64 -1, implicit-def $exec, implicit-def dead $scc, implicit $exec
+    ; GCN-NEXT: BUFFER_STORE_DWORD_OFFSET $vgpr0, $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 0, 0, 0, implicit $exec :: (store (s32) into %stack.0, addrspace 5)
+    ; GCN-NEXT: BUFFER_STORE_DWORD_OFFSET $vgpr2, $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 4, 0, 0, implicit $exec :: (store (s32) into %stack.1, addrspace 5)
+    ; GCN-NEXT: $exec = S_MOV_B64 killed $sgpr4_sgpr5
+    ; GCN-NEXT: $vgpr0 = IMPLICIT_DEF
+    ; GCN-NEXT: $vgpr2 = IMPLICIT_DEF
+    ; GCN-NEXT: $vgpr0 = V_WRITELANE_B32 killed $sgpr20, 0, $vgpr0
+    ; GCN-NEXT: $vgpr2 = V_WRITELANE_B32 killed $sgpr21, 0, $vgpr2
+    ; GCN-NEXT: $vgpr0 = COPY $vgpr1, implicit $exec
+    ; GCN-NEXT: $sgpr4_sgpr5 = S_XOR_SAVEEXEC_B64 -1, implicit-def $exec, implicit-def dead $scc, implicit $exec
+    ; GCN-NEXT: $vgpr0 = BUFFER_LOAD_DWORD_OFFSET $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 0, 0, 0, implicit $exec, implicit $vgpr0(tied-def 0) :: (load (s32) from %stack.0, addrspace 5)
+    ; GCN-NEXT: $vgpr2 = BUFFER_LOAD_DWORD_OFFSET $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 4, 0, 0, implicit $exec :: (load (s32) from %stack.1, addrspace 5)
+    ; GCN-NEXT: $exec = S_MOV_B64 killed $sgpr4_sgpr5
+    ; GCN-NEXT: SI_RETURN implicit $vgpr0_vgpr1
+    $vgpr0 = IMPLICIT_DEF
+    $vgpr2 = IMPLICIT_DEF
+    $vgpr0 = V_WRITELANE_B32 killed $sgpr20, 0, $vgpr0
+    $vgpr2 = V_WRITELANE_B32 killed $sgpr21, 0, $vgpr2
+    $vgpr0 = COPY $vgpr1, implicit $exec
+    SI_RETURN implicit $vgpr0_vgpr1
+...
+
+# Tied op not required in the spill reload of vgpr2.
+
+---
+name:            wwm_scratch_reg_spill_reload_different_outgoing_reg
+tracksRegLiveness: true
+machineFunctionInfo:
+  wwmReservedRegs: ['$vgpr2']
+  isEntryFunction: false
+  scratchRSrcReg:  '$sgpr0_sgpr1_sgpr2_sgpr3'
+  stackPtrOffsetReg: '$sgpr32'
+  frameOffsetReg: '$sgpr33'
+body:             |
+  bb.0:
+    liveins: $sgpr20, $vgpr1
+    ; GCN-LABEL: name: wwm_scratch_reg_spill_reload_different_outgoing_reg
+    ; GCN: liveins: $sgpr20, $vgpr1, $vgpr2
+    ; GCN-NEXT: {{  $}}
+    ; GCN-NEXT: $sgpr4_sgpr5 = S_XOR_SAVEEXEC_B64 -1, implicit-def $exec, implicit-def dead $scc, implicit $exec
+    ; GCN-NEXT: BUFFER_STORE_DWORD_OFFSET $vgpr2, $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 0, 0, 0, implicit $exec :: (store (s32) into %stack.0, addrspace 5)
+    ; GCN-NEXT: $exec = S_MOV_B64 killed $sgpr4_sgpr5
+    ; GCN-NEXT: $vgpr2 = IMPLICIT_DEF
+    ; GCN-NEXT: $vgpr2 = V_WRITELANE_B32 killed $sgpr20, 0, $vgpr2
+    ; GCN-NEXT: $vgpr0 = COPY $vgpr1, implicit $exec
+    ; GCN-NEXT: $sgpr4_sgpr5 = S_XOR_SAVEEXEC_B64 -1, implicit-def $exec, implicit-def dead $scc, implicit $exec
+    ; GCN-NEXT: $vgpr2 = BUFFER_LOAD_DWORD_OFFSET $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 0, 0, 0, implicit $exec :: (load (s32) from %stack.0, addrspace 5)
+    ; GCN-NEXT: $exec = S_MOV_B64 killed $sgpr4_sgpr5
+    ; GCN-NEXT: SI_RETURN implicit $vgpr0_vgpr1
+    $vgpr2 = IMPLICIT_DEF
+    $vgpr2 = V_WRITELANE_B32 killed $sgpr20, 0, $vgpr2
+    $vgpr0 = COPY $vgpr1, implicit $exec
+    SI_RETURN implicit $vgpr0_vgpr1
+...
+
+# Tied op not required in the spill reload of vgpr40 which is in the CSR range.
+---
+name:            wwm_csr_spill_reload
+tracksRegLiveness: true
+machineFunctionInfo:
+  wwmReservedRegs: ['$vgpr40']
+  isEntryFunction: false
+  scratchRSrcReg:  '$sgpr0_sgpr1_sgpr2_sgpr3'
+  stackPtrOffsetReg: '$sgpr32'
+  frameOffsetReg: '$sgpr33'
+body:             |
+  bb.0:
+    liveins: $sgpr20, $vgpr1
+    ; GCN-LABEL: name: wwm_csr_spill_reload
+    ; GCN: liveins: $sgpr20, $vgpr1, $vgpr40
+    ; GCN-NEXT: {{  $}}
+    ; GCN-NEXT: $sgpr4_sgpr5 = S_OR_SAVEEXEC_B64 -1, implicit-def $exec, implicit-def dead $scc, implicit $exec
+    ; GCN-NEXT: BUFFER_STORE_DWORD_OFFSET $vgpr40, $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 0, 0, 0, implicit $exec :: (store (s32) into %stack.0, addrspace 5)
+    ; GCN-NEXT: $exec = S_MOV_B64 killed $sgpr4_sgpr5
+    ; GCN-NEXT: $vgpr40 = IMPLICIT_DEF
+    ; GCN-NEXT: $vgpr40 = V_WRITELANE_B32 killed $sgpr20, 0, $vgpr40
+    ; GCN-NEXT: $sgpr20 = V_READLANE_B32 $vgpr40, 0, implicit $exec
+    ; GCN-NEXT: $vgpr0 = COPY killed $vgpr1, implicit $exec
+    ; GCN-NEXT: $sgpr4_sgpr5 = S_OR_SAVEEXEC_B64 -1, implicit-def $exec, implicit-def dead $scc, implicit $exec
+    ; GCN-NEXT: $vgpr40 = BUFFER_LOAD_DWORD_OFFSET $sgpr0_sgpr1_sgpr2_sgpr3, $sgpr32, 0, 0, 0, implicit $exec :: (load (s32) from %stack.0, addrspace 5)
+    ; GCN-NEXT: $exec = S_MOV_B64 killed $sgpr4_sgpr5
+    ; GCN-NEXT: SI_RETURN implicit $vgpr0
+    $vgpr40 = IMPLICIT_DEF
+    $vgpr40 = V_WRITELANE_B32 killed $sgpr20, 0, $vgpr40
+    $sgpr20 = V_READLANE_B32 $vgpr40, 0, implicit $exec
+    $vgpr0 = COPY killed $vgpr1, implicit $exec
+    SI_RETURN implicit $vgpr0
+...