[AMDGPU] Change FLAT Scratch SADDR to VADDR form in moveToVALU

[rampitec]

Summary

Extend the legalization of global SADDR loads and stores
with changing to VADDR to the FLAT scratch instructions.

[arsenm]

llvm/lib/Target/AMDGPU/SIInstrInfo.h

1228

These functions should probably get comments. It's hard to parse out what this is without already knowing exactly what this is for

[rampitec]

llvm/lib/Target/AMDGPU/SIInstrInfo.h

1228

Ack, will do.

[Joe_Nash]

LGTM

[rampitec]

Added comments to helper functions.

[rampitec]

Rebased.

[rampitec]

Rebased.

[rampitec]

Rebased.

[rampitec]

Rebased.

llvm/lib/Target/AMDGPU/SIInstrInfo.h

namespace AMDGPU {
   /// \returns VADDR form of a FLAT Global instruction given an \p Opcode
   /// of a SADDR form.
   LLVM_READONLY
   int getGlobalVaddrOp(uint16_t Opcode);
 
   LLVM_READONLY
   int getVCMPXNoSDstOp(uint16_t Opcode);
 
+  /// \returns ST form with only immediate offset of a FLAT Scratch instruction
+  /// given an \p Opcode of an SS (SADDR) form.
   LLVM_READONLY
   int getFlatScratchInstSTfromSS(uint16_t Opcode);
 
+  /// \returns SS (SADDR) form of a FLAT Scratch instruction given an \p Opcode
+  /// of an SV (VADDR) form.
   LLVM_READONLY
   int getFlatScratchInstSSfromSV(uint16_t Opcode);
 
+  /// \returns SV (VADDR) form of a FLAT Scratch instruction given an \p Opcode
+  /// of an SS (SADDR) form.
+  LLVM_READONLY
+  int getFlatScratchInstSVfromSS(uint16_t Opcode);
+
   const uint64_t RSRC_DATA_FORMAT = 0xf00000000000LL;
   const uint64_t RSRC_ELEMENT_SIZE_SHIFT = (32 + 19);
   const uint64_t RSRC_INDEX_STRIDE_SHIFT = (32 + 21);
   const uint64_t RSRC_TID_ENABLE = UINT64_C(1) << (32 + 23);
 
 } // end namespace AMDGPU
 
 namespace SI {

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

bool SIInstrInfo::moveFlatAddrToVGPR(MachineInstr &Inst) const {
   int OldSAddrIdx = AMDGPU::getNamedOperandIdx(Opc, AMDGPU::OpName::saddr);
   if (OldSAddrIdx < 0)
     return false;
 
   assert(isSegmentSpecificFLAT(Inst));
 
   int NewOpc = AMDGPU::getGlobalVaddrOp(Opc);
   if (NewOpc < 0)
+    NewOpc = AMDGPU::getFlatScratchInstSVfromSS(Opc);
+  if (NewOpc < 0)
     return false;
 
   MachineRegisterInfo &MRI = Inst.getMF()->getRegInfo();
   MachineOperand &SAddr = Inst.getOperand(OldSAddrIdx);
   if (RI.isSGPRReg(MRI, SAddr.getReg()))
     return false;
 
   int NewVAddrIdx = AMDGPU::getNamedOperandIdx(NewOpc, AMDGPU::OpName::vaddr);
   if (NewVAddrIdx < 0)
     return false;
 
   int OldVAddrIdx = AMDGPU::getNamedOperandIdx(Opc, AMDGPU::OpName::vaddr);
-  assert(OldVAddrIdx >= 0);
 
-  // Check vaddr, it shall be zero
+  // Check vaddr, it shall be zero or absent.
+  MachineInstr *VAddrDef = nullptr;
+  if (OldVAddrIdx >= 0) {
     MachineOperand &VAddr = Inst.getOperand(OldVAddrIdx);
-  MachineInstr *VAddrDef = MRI.getUniqueVRegDef(VAddr.getReg());
+    VAddrDef = MRI.getUniqueVRegDef(VAddr.getReg());
     if (!VAddrDef || VAddrDef->getOpcode() != AMDGPU::V_MOV_B32_e32 ||
-      !VAddrDef->getOperand(1).isImm() || VAddrDef->getOperand(1).getImm() != 0)
+        !VAddrDef->getOperand(1).isImm() ||
+        VAddrDef->getOperand(1).getImm() != 0)
       return false;
+  }
 
   const MCInstrDesc &NewDesc = get(NewOpc);
   Inst.setDesc(NewDesc);
 
   // Callers expect interator to be valid after this call, so modify the
   // instruction in place.
   if (OldVAddrIdx == NewVAddrIdx) {
     MachineOperand &NewVAddr = Inst.getOperand(NewVAddrIdx);
     // Clear use list from the old vaddr holding a zero register.
     MRI.removeRegOperandFromUseList(&NewVAddr);
     MRI.moveOperands(&NewVAddr, &SAddr, 1);
     Inst.RemoveOperand(OldSAddrIdx);
     // Update the use list with the pointer we have just moved from vaddr to
     // saddr poisition. Otherwise new vaddr will be missing from the use list.
     MRI.removeRegOperandFromUseList(&NewVAddr);
     MRI.addRegOperandToUseList(&NewVAddr);
   } else {
     assert(OldSAddrIdx == NewVAddrIdx);
+
+    if (OldVAddrIdx >= 0)
       Inst.RemoveOperand(OldVAddrIdx);
   }
 
-  if (MRI.use_nodbg_empty(VAddrDef->getOperand(0).getReg()))
+  if (VAddrDef && MRI.use_nodbg_empty(VAddrDef->getOperand(0).getReg()))
     VAddrDef->eraseFromParent();
 
   return true;
 }
 
 // FIXME: Remove this when SelectionDAG is obsoleted.
 void SIInstrInfo::legalizeOperandsFLAT(MachineRegisterInfo &MRI,
                                        MachineInstr &MI) const {

llvm/lib/Target/AMDGPU/SIInstrInfo.td

 def getFlatScratchInstSSfromSV : InstrMapping {
   let FilterClass = "FlatScratchInst";
   let RowFields = ["SVOp"];
   let ColFields = ["Mode"];
   let KeyCol = ["SV"];
   let ValueCols = [["SS"]];
 }
 
+def getFlatScratchInstSVfromSS : InstrMapping {
+  let FilterClass = "FlatScratchInst";
+  let RowFields = ["SVOp"];
+  let ColFields = ["Mode"];
+  let KeyCol = ["SS"];
+  let ValueCols = [["SV"]];
+}
+
 include "SIInstructions.td"
 
 include "DSInstructions.td"
 include "MIMGInstructions.td"

llvm/test/CodeGen/AMDGPU/move-load-addr-to-valu.mir

body:             |
   ; GCN-LABEL: name: scratch_load_saddr_to_valu
   ; GCN: bb.0:
   ; GCN:   successors: %bb.1(0x80000000)
   ; GCN:   liveins: $vgpr0
   ; GCN:   [[COPY:%[0-9]+]]:vgpr_32 = COPY $vgpr0
   ; GCN: bb.1:
   ; GCN:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
   ; GCN:   [[PHI:%[0-9]+]]:vgpr_32 = PHI [[COPY]], %bb.0, %6, %bb.1
-  ; GCN:   [[V_READFIRSTLANE_B32_:%[0-9]+]]:sgpr_32 = V_READFIRSTLANE_B32 [[PHI]], implicit $exec
-  ; GCN:   [[SCRATCH_LOAD_DWORD_SADDR:%[0-9]+]]:vgpr_32 = SCRATCH_LOAD_DWORD_SADDR [[V_READFIRSTLANE_B32_]], 0, 0, implicit $exec, implicit $flat_scr
+  ; GCN:   [[SCRATCH_LOAD_DWORD:%[0-9]+]]:vgpr_32 = SCRATCH_LOAD_DWORD [[PHI]], 0, 0, implicit $exec, implicit $flat_scr
   ; GCN:   [[V_AND_B32_e64_:%[0-9]+]]:vgpr_32 = V_AND_B32_e64 [[PHI]], 1, implicit $exec
   ; GCN:   V_CMP_NE_U32_e32 0, [[V_AND_B32_e64_]], implicit-def $vcc, implicit $exec
   ; GCN:   [[COPY1:%[0-9]+]]:vgpr_32 = COPY [[V_AND_B32_e64_]], implicit $exec
   ; GCN:   $vcc = S_AND_B64 $exec, $vcc, implicit-def $scc
   ; GCN:   S_CBRANCH_VCCNZ %bb.1, implicit $vcc
   ; GCN: bb.2:
   ; GCN:   S_ENDPGM 0
   bb.0:
body:             |
   ; GCN: bb.0:
   ; GCN:   successors: %bb.1(0x80000000)
   ; GCN:   liveins: $vgpr0
   ; GCN:   [[COPY:%[0-9]+]]:vgpr_32 = COPY $vgpr0
   ; GCN: bb.1:
   ; GCN:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
   ; GCN:   [[PHI:%[0-9]+]]:vgpr_32 = PHI [[COPY]], %bb.0, %6, %bb.1
   ; GCN:   [[DEF:%[0-9]+]]:vgpr_32 = IMPLICIT_DEF
-  ; GCN:   [[V_READFIRSTLANE_B32_:%[0-9]+]]:sgpr_32 = V_READFIRSTLANE_B32 [[PHI]], implicit $exec
-  ; GCN:   SCRATCH_STORE_DWORD_SADDR [[DEF]], [[V_READFIRSTLANE_B32_]], 0, 0, implicit $exec, implicit $flat_scr
+  ; GCN:   SCRATCH_STORE_DWORD [[DEF]], [[PHI]], 0, 0, implicit $exec, implicit $flat_scr
   ; GCN:   [[V_AND_B32_e64_:%[0-9]+]]:vgpr_32 = V_AND_B32_e64 [[PHI]], 1, implicit $exec
   ; GCN:   V_CMP_NE_U32_e32 0, [[V_AND_B32_e64_]], implicit-def $vcc, implicit $exec
   ; GCN:   [[COPY1:%[0-9]+]]:vgpr_32 = COPY [[V_AND_B32_e64_]], implicit $exec
   ; GCN:   $vcc = S_AND_B64 $exec, $vcc, implicit-def $scc
   ; GCN:   S_CBRANCH_VCCNZ %bb.1, implicit $vcc
   ; GCN: bb.2:
   ; GCN:   S_ENDPGM 0
   bb.0: