AMDGPU: Fold immediate offset into BUFFER_LOAD_DWORD lowered from SMEM

[mareko]

Summary

-5.3% code size in affected shaders.

Changed stats only:

48486 shaders in 30489 tests
Totals:
SGPRS: 2086406 -> 2072430 (-0.67 %)
VGPRS: 1626872 -> 1627960 (0.07 %)
Spilled SGPRs: 7865 -> 7912 (0.60 %)
Code Size: 60978060 -> 60188764 (-1.29 %) bytes
Max Waves: 374530 -> 374342 (-0.05 %)

Totals from affected shaders:
SGPRS: 299664 -> 285688 (-4.66 %)
VGPRS: 233844 -> 234932 (0.47 %)
Spilled SGPRs: 3959 -> 4006 (1.19 %)
Code Size: 14905272 -> 14115976 (-5.30 %) bytes
Max Waves: 46202 -> 46014 (-0.41 %)

[mareko]

Rebase.

[nhaehnle]

Some small comments inline. More generally, I think stuff like this belongs in a MIR-level InstCombine pass. We just don't have the infrastructure for that yet, and this patch is a nice improvement in the meantime.

lib/Target/AMDGPU/SIInstrInfo.cpp

3727–3732 ↗

(On Diff #119335)

Could you please reduce the nesting here a bit by wrapping this first part in if (Src && Src->isReg())?

3742 ↗

(On Diff #119335)

Not a huge deal, but it seems a bit silly not to handle the case of an immediate Offset == 0. Such stuff should mostly be caught earlier, but who knows what kind of nonsense intermediate passes generate again.

[nhaehnle]

P.S.: It would help to use ReviewBoard's feature of defining dependent (parent/child) revisions, like I've just done now. It sucks that this is not done automatically, unfortunately I haven't found a good way to do proper reviews of series with it yet. (This is my mine gripe with all those shiny tools written by web hipsters...)

[mareko]

Another improvement that can be done here is that LLVM sometimes generates OR instead of ADD, and I don't know how to tell when OR means ADD.

[mareko]

• address feedback
• Offset == 0 also means that the operand is not an immediate.

llvm/trunk/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp

bool AMDGPUDAGToDAGISel::SelectDS64Bit4ByteAligned(SDValue Addr, SDValue &Base,
   // FIXME: This is broken on SI where we still need to check if the base
   // pointer is positive here.
   Base = Addr;
   Offset0 = CurDAG->getTargetConstant(0, DL, MVT::i8);
   Offset1 = CurDAG->getTargetConstant(1, DL, MVT::i8);
   return true;
 }
 
-static bool isLegalMUBUFImmOffset(unsigned Imm) {
-  return isUInt<12>(Imm);
-}
-
-static bool isLegalMUBUFImmOffset(const ConstantSDNode *Imm) {
-  return isLegalMUBUFImmOffset(Imm->getZExtValue());
-}
-
 bool AMDGPUDAGToDAGISel::SelectMUBUF(SDValue Addr, SDValue &Ptr,
                                      SDValue &VAddr, SDValue &SOffset,
                                      SDValue &Offset, SDValue &Offen,
                                      SDValue &Idxen, SDValue &Addr64,
                                      SDValue &GLC, SDValue &SLC,
                                      SDValue &TFE) const {
   // Subtarget prefers to use flat instruction
   if (Subtarget->useFlatForGlobal())
if (N0.getOpcode() == ISD::ADD) {
       Ptr = N2;
       VAddr = N3;
     } else {
       // (add N0, C1) -> offset
       VAddr = CurDAG->getTargetConstant(0, DL, MVT::i32);
       Ptr = N0;
     }
 
-    if (isLegalMUBUFImmOffset(C1)) {
+    if (SIInstrInfo::isLegalMUBUFImmOffset(C1->getZExtValue())) {
       Offset = CurDAG->getTargetConstant(C1->getZExtValue(), DL, MVT::i16);
       return true;
     }
 
     if (isUInt<32>(C1->getZExtValue())) {
       // Illegal offset, store it in soffset.
       Offset = CurDAG->getTargetConstant(0, DL, MVT::i16);
       SOffset = SDValue(CurDAG->getMachineNode(AMDGPU::S_MOV_B32, DL, MVT::i32,
bool AMDGPUDAGToDAGISel::SelectMUBUFScratchOffen(SDNode *Parent,
   SDLoc DL(Addr);
   MachineFunction &MF = CurDAG->getMachineFunction();
   const SIMachineFunctionInfo *Info = MF.getInfo<SIMachineFunctionInfo>();
 
   Rsrc = CurDAG->getRegister(Info->getScratchRSrcReg(), MVT::v4i32);
 
   if (ConstantSDNode *CAddr = dyn_cast<ConstantSDNode>(Addr)) {
     unsigned Imm = CAddr->getZExtValue();
-    assert(!isLegalMUBUFImmOffset(Imm) &&
+    assert(!SIInstrInfo::isLegalMUBUFImmOffset(Imm) &&
            "should have been selected by other pattern");
 
     SDValue HighBits = CurDAG->getTargetConstant(Imm & ~4095, DL, MVT::i32);
     MachineSDNode *MovHighBits = CurDAG->getMachineNode(AMDGPU::V_MOV_B32_e32,
                                                         DL, MVT::i32, HighBits);
     VAddr = SDValue(MovHighBits, 0);
 
     // In a call sequence, stores to the argument stack area are relative to the
bool AMDGPUDAGToDAGISel::SelectMUBUFScratchOffen(SDNode *Parent,
   if (CurDAG->isBaseWithConstantOffset(Addr)) {
     // (add n0, c1)
 
     SDValue N0 = Addr.getOperand(0);
     SDValue N1 = Addr.getOperand(1);
 
     // Offsets in vaddr must be positive.
     ConstantSDNode *C1 = cast<ConstantSDNode>(N1);
-    if (isLegalMUBUFImmOffset(C1)) {
+    if (SIInstrInfo::isLegalMUBUFImmOffset(C1->getZExtValue())) {
       std::tie(VAddr, SOffset) = foldFrameIndex(N0);
       ImmOffset = CurDAG->getTargetConstant(C1->getZExtValue(), DL, MVT::i16);
       return true;
     }
   }
 
   // (node)
   std::tie(VAddr, SOffset) = foldFrameIndex(Addr);
   ImmOffset = CurDAG->getTargetConstant(0, DL, MVT::i16);
   return true;
 }
 
 bool AMDGPUDAGToDAGISel::SelectMUBUFScratchOffset(SDNode *Parent,
                                                   SDValue Addr,
                                                   SDValue &SRsrc,
                                                   SDValue &SOffset,
                                                   SDValue &Offset) const {
   ConstantSDNode *CAddr = dyn_cast<ConstantSDNode>(Addr);
-  if (!CAddr || !isLegalMUBUFImmOffset(CAddr))
+  if (!CAddr || !SIInstrInfo::isLegalMUBUFImmOffset(CAddr->getZExtValue()))
     return false;
 
   SDLoc DL(Addr);
   MachineFunction &MF = CurDAG->getMachineFunction();
   const SIMachineFunctionInfo *Info = MF.getInfo<SIMachineFunctionInfo>();
 
   SRsrc = CurDAG->getRegister(Info->getScratchRSrcReg(), MVT::v4i32);
 

llvm/trunk/lib/Target/AMDGPU/SIInstrInfo.h

public:
   /// TODO: After GFX9 it should return a no-carry operation.
   MachineInstrBuilder getAddNoCarry(MachineBasicBlock &MBB,
                                     MachineBasicBlock::iterator I,
                                     const DebugLoc &DL,
                                     unsigned DestReg) const;
 
   static bool isKillTerminator(unsigned Opcode);
   const MCInstrDesc &getKillTerminatorFromPseudo(unsigned Opcode) const;
+
+  static bool isLegalMUBUFImmOffset(unsigned Imm) {
+    return isUInt<12>(Imm);
+  }
 };
 
 namespace AMDGPU {
 
   LLVM_READONLY
   int getVOPe64(uint16_t Opcode);
 
   LLVM_READONLY

llvm/trunk/lib/Target/AMDGPU/SIInstrInfo.cpp

while (!Worklist.empty()) {
 
     case AMDGPU::S_XNOR_B64:
       splitScalar64BitBinaryOp(Worklist, Inst, AMDGPU::S_XNOR_B32);
       Inst.eraseFromParent();
       continue;
 
     case AMDGPU::S_BUFFER_LOAD_DWORD_SGPR: {
       unsigned VDst = MRI.createVirtualRegister(&AMDGPU::VGPR_32RegClass);
+      const MachineOperand *VAddr = getNamedOperand(Inst, AMDGPU::OpName::soff);
+      auto Add = MRI.getUniqueVRegDef(VAddr->getReg());
+      unsigned Offset = 0;
+
+      // See if we can extract an immediate offset by recognizing one of these:
+      //   V_ADD_I32_e32 dst, imm, src1
+      //   V_ADD_I32_e32 dst, (S_MOV_B32 imm), src1
+      // V_ADD will be removed by "Remove dead machine instructions".
+      if (Add && Add->getOpcode() == AMDGPU::V_ADD_I32_e32) {
+        const MachineOperand *Src =
+          getNamedOperand(*Add, AMDGPU::OpName::src0);
+
+        if (Src && Src->isReg()) {
+          auto Mov = MRI.getUniqueVRegDef(Src->getReg());
+          if (Mov && Mov->getOpcode() == AMDGPU::S_MOV_B32)
+            Src = &Mov->getOperand(1);
+        }
+
+        if (Src) {
+          if (Src->isImm())
+            Offset = Src->getImm();
+          else if (Src->isCImm())
+            Offset = Src->getCImm()->getZExtValue();
+        }
+
+        if (Offset && isLegalMUBUFImmOffset(Offset))
+          VAddr = getNamedOperand(*Add, AMDGPU::OpName::src1);
+        else
+          Offset = 0;
+      }
 
       BuildMI(*MBB, Inst, Inst.getDebugLoc(),
               get(AMDGPU::BUFFER_LOAD_DWORD_OFFEN), VDst)
-        .add(*getNamedOperand(Inst, AMDGPU::OpName::soff)) // vaddr
+        .add(*VAddr) // vaddr
         .add(*getNamedOperand(Inst, AMDGPU::OpName::sbase)) // srsrc
         .addImm(0) // soffset
-        .addImm(0) // offset
+        .addImm(Offset) // offset
         .addImm(getNamedOperand(Inst, AMDGPU::OpName::glc)->getImm())
         .addImm(0) // slc
         .addImm(0) // tfe
         .setMemRefs(Inst.memoperands_begin(), Inst.memoperands_end());
 
       MRI.replaceRegWith(getNamedOperand(Inst, AMDGPU::OpName::sdst)->getReg(),
                          VDst);
       addUsersToMoveToVALUWorklist(VDst, MRI, Worklist);

llvm/trunk/test/CodeGen/AMDGPU/smrd.ll

 ; GCN-LABEL: {{^}}smrd_vgpr_offset:
 ; GCN: buffer_load_dword v{{[0-9]}}, v0, s[0:3], 0 offen ;
 define amdgpu_ps float @smrd_vgpr_offset(<4 x i32> inreg %desc, i32 %offset) #0 {
 main_body:
   %r = call float @llvm.SI.load.const.v4i32(<4 x i32> %desc, i32 %offset)
   ret float %r
 }
 
+; GCN-LABEL: {{^}}smrd_vgpr_offset_imm:
+; GCN-NEXT: BB#
+; GCN-NEXT: buffer_load_dword v{{[0-9]}}, v0, s[0:3], 0 offen offset:4095 ;
+define amdgpu_ps float @smrd_vgpr_offset_imm(<4 x i32> inreg %desc, i32 %offset) #0 {
+main_body:
+  %off = add i32 %offset, 4095
+  %r = call float @llvm.SI.load.const.v4i32(<4 x i32> %desc, i32 %off)
+  ret float %r
+}
+
+; GCN-LABEL: {{^}}smrd_vgpr_offset_imm_too_large:
+; GCN-NEXT: BB#
+; GCN-NEXT: v_add_i32_e32 v0, vcc, 0x1000, v0
+; GCN-NEXT: buffer_load_dword v{{[0-9]}}, v0, s[0:3], 0 offen ;
+define amdgpu_ps float @smrd_vgpr_offset_imm_too_large(<4 x i32> inreg %desc, i32 %offset) #0 {
+main_body:
+  %off = add i32 %offset, 4096
+  %r = call float @llvm.SI.load.const.v4i32(<4 x i32> %desc, i32 %off)
+  ret float %r
+}
+
 declare void @llvm.amdgcn.exp.f32(i32, i32, float, float, float, float, i1, i1) #0
 declare float @llvm.SI.load.const.v4i32(<4 x i32>, i32) #1
 
 attributes #0 = { nounwind }
 attributes #1 = { nounwind readnone }