[AMDGPU][CodeGen] Support (base | offset) SMEM loads.

[kosarev]

Summary

Prevents generation of unnecessary s_or_b32 instructions.

[kosarev]

Reduces code size for ~770 corpus shaders by ~36 bytes in average.

[foad]

As a follow up could we update all calls to getBaseWithConstantOffset to pass in a non-null KnownBits?

llvm/lib/Target/AMDGPU/AMDGPUGlobalISelUtils.cpp

48

Test KnownBits first because it's cheapest.

llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp

1990–1991

Nice catch!

llvm/test/CodeGen/AMDGPU/amdgcn-load-offset-from-reg.ll

138

Please precommit this new test.

[kosarev]

Rebase on top of D133021 and change the if() as suggested.

[kosarev]

Rebased.

llvm/lib/Target/AMDGPU/AMDGPUGlobalISelUtils.h

 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/CodeGen/Register.h"
 #include <utility>
 
 namespace llvm {
 
 class MachineRegisterInfo;
 class GCNSubtarget;
+class GISelKnownBits;
 class LLT;
 
 namespace AMDGPU {
 
 /// Returns base register and constant offset.
 std::pair<Register, unsigned>
-getBaseWithConstantOffset(MachineRegisterInfo &MRI, Register Reg);
+getBaseWithConstantOffset(MachineRegisterInfo &MRI, Register Reg,
+                          GISelKnownBits *KnownBits = nullptr);
 
 bool isLegalVOP3PShuffleMask(ArrayRef<int> Mask);
 bool hasAtomicFaddRtnForTy(const GCNSubtarget &Subtarget, const LLT &Ty);
 }
 }
 
 #endif

llvm/lib/Target/AMDGPU/AMDGPUGlobalISelUtils.cpp

 //===- AMDGPUGlobalISelUtils.cpp ---------------------------------*- C++ -*-==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPUGlobalISelUtils.h"
 #include "GCNSubtarget.h"
+#include "llvm/CodeGen/GlobalISel/GISelKnownBits.h"
 #include "llvm/CodeGen/GlobalISel/MIPatternMatch.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/Support/LowLevelTypeImpl.h"
 
 using namespace llvm;
 using namespace MIPatternMatch;
 
 std::pair<Register, unsigned>
-AMDGPU::getBaseWithConstantOffset(MachineRegisterInfo &MRI, Register Reg) {
+AMDGPU::getBaseWithConstantOffset(MachineRegisterInfo &MRI, Register Reg,
+                                  GISelKnownBits *KnownBits) {
   MachineInstr *Def = getDefIgnoringCopies(Reg, MRI);
   if (!Def)
     return std::make_pair(Reg, 0);
 
   if (Def->getOpcode() == TargetOpcode::G_CONSTANT) {
     unsigned Offset;
     const MachineOperand &Op = Def->getOperand(1);
     if (Op.isImm())
if (Def->getOpcode() == TargetOpcode::G_ADD) {
     if (mi_match(Def->getOperand(2).getReg(), MRI, m_ICst(Offset)))
       return std::make_pair(Def->getOperand(1).getReg(), Offset);
 
     // FIXME: matcher should ignore copies
     if (mi_match(Def->getOperand(2).getReg(), MRI, m_Copy(m_ICst(Offset))))
       return std::make_pair(Def->getOperand(1).getReg(), Offset);
   }
 
+  Register Base;
+  if (KnownBits && mi_match(Reg, MRI, m_GOr(m_Reg(Base), m_ICst(Offset))) &&
+      KnownBits->maskedValueIsZero(Base, APInt(32, Offset)))
+    return std::make_pair(Base, Offset);
+
   // Handle G_PTRTOINT (G_PTR_ADD base, const) case
   if (Def->getOpcode() == TargetOpcode::G_PTRTOINT) {
     MachineInstr *Base;
     if (mi_match(Def->getOperand(1).getReg(), MRI,
                  m_GPtrAdd(m_MInstr(Base), m_ICst(Offset)))) {
       // If Base was int converted to pointer, simply return int and offset.
       if (Base->getOpcode() == TargetOpcode::G_INTTOPTR)
         return std::make_pair(Base->getOperand(1).getReg(), Offset);

llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp

if (SOffset && Offset) {
     assert(!Imm32Only && !IsBuffer);
     SDValue B;
     return SelectSMRDBaseOffset(Addr, B, nullptr, Offset) &&
            SelectSMRDBaseOffset(B, SBase, SOffset, nullptr);
   }
 
   // A 32-bit (address + offset) should not cause unsigned 32-bit integer
   // wraparound, because s_load instructions perform the addition in 64 bits.
-  if (Addr.getValueType() == MVT::i32 && !Addr->getFlags().hasNoUnsignedWrap())
+  if (Addr.getValueType() == MVT::i32 && Addr.getOpcode() == ISD::ADD &&
+      !Addr->getFlags().hasNoUnsignedWrap())
     return false;
 
   SDValue N0, N1;
   // Extract the base and offset if possible.
   if (CurDAG->isBaseWithConstantOffset(Addr) || Addr.getOpcode() == ISD::ADD) {
     N0 = Addr.getOperand(0);
     N1 = Addr.getOperand(1);
   } else if (getBaseWithOffsetUsingSplitOR(*CurDAG, Addr, N0, N1)) {

llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp

if (OffsetDef->getOpcode() == AMDGPU::G_CONSTANT) {
     // default -1 only set the low 16-bits, we could leave it as-is and add 1 to
     // the immediate offset.
 
     ImmOffset = OffsetDef->getOperand(1).getCImm()->getZExtValue();
     BuildMI(*MBB, &MI, DL, TII.get(AMDGPU::S_MOV_B32), AMDGPU::M0)
       .addImm(0);
   } else {
     std::tie(BaseOffset, ImmOffset) =
-        AMDGPU::getBaseWithConstantOffset(*MRI, BaseOffset);
+        AMDGPU::getBaseWithConstantOffset(*MRI, BaseOffset, KnownBits);
 
     if (Readfirstlane) {
       // We have the constant offset now, so put the readfirstlane back on the
       // variable component.
       if (!RBI.constrainGenericRegister(BaseOffset, AMDGPU::VGPR_32RegClass, *MRI))
         return false;
 
       Readfirstlane->getOperand(1).setReg(BaseOffset);
BuildMI(*BB, &I, DL, TII.get(AMDGPU::REG_SEQUENCE), DstReg)
     .addImm(AMDGPU::sub1);
   I.eraseFromParent();
   return true;
 }
 
 /// Return the register to use for the index value, and the subregister to use
 /// for the indirectly accessed register.
 static std::pair<Register, unsigned>
-computeIndirectRegIndex(MachineRegisterInfo &MRI,
-                        const SIRegisterInfo &TRI,
-                        const TargetRegisterClass *SuperRC,
-                        Register IdxReg,
-                        unsigned EltSize) {
+computeIndirectRegIndex(MachineRegisterInfo &MRI, const SIRegisterInfo &TRI,
+                        const TargetRegisterClass *SuperRC, Register IdxReg,
+                        unsigned EltSize, GISelKnownBits &KnownBits) {
   Register IdxBaseReg;
   int Offset;
 
-  std::tie(IdxBaseReg, Offset) = AMDGPU::getBaseWithConstantOffset(MRI, IdxReg);
+  std::tie(IdxBaseReg, Offset) =
+      AMDGPU::getBaseWithConstantOffset(MRI, IdxReg, &KnownBits);
   if (IdxBaseReg == AMDGPU::NoRegister) {
     // This will happen if the index is a known constant. This should ordinarily
     // be legalized out, but handle it as a register just in case.
     assert(Offset == 0);
     IdxBaseReg = IdxReg;
   }
 
   ArrayRef<int16_t> SubRegs = TRI.getRegSplitParts(SuperRC, EltSize);
if (!RBI.constrainGenericRegister(SrcReg, *SrcRC, *MRI) ||
       !RBI.constrainGenericRegister(IdxReg, AMDGPU::SReg_32RegClass, *MRI))
     return false;
 
   MachineBasicBlock *BB = MI.getParent();
   const DebugLoc &DL = MI.getDebugLoc();
   const bool Is64 = DstTy.getSizeInBits() == 64;
 
   unsigned SubReg;
-  std::tie(IdxReg, SubReg) = computeIndirectRegIndex(*MRI, TRI, SrcRC, IdxReg,
-                                                     DstTy.getSizeInBits() / 8);
+  std::tie(IdxReg, SubReg) = computeIndirectRegIndex(
+      *MRI, TRI, SrcRC, IdxReg, DstTy.getSizeInBits() / 8, *KnownBits);
 
   if (SrcRB->getID() == AMDGPU::SGPRRegBankID) {
     if (DstTy.getSizeInBits() != 32 && !Is64)
       return false;
 
     BuildMI(*BB, &MI, DL, TII.get(AMDGPU::COPY), AMDGPU::M0)
       .addReg(IdxReg);
 
if (!RBI.constrainGenericRegister(VecReg, *VecRC, *MRI) ||
       !RBI.constrainGenericRegister(IdxReg, AMDGPU::SReg_32RegClass, *MRI))
     return false;
 
   if (VecRB->getID() == AMDGPU::VGPRRegBankID && ValSize != 32)
     return false;
 
   unsigned SubReg;
   std::tie(IdxReg, SubReg) = computeIndirectRegIndex(*MRI, TRI, VecRC, IdxReg,
-                                                     ValSize / 8);
+                                                     ValSize / 8, *KnownBits);
 
   const bool IndexMode = VecRB->getID() == AMDGPU::VGPRRegBankID &&
                          STI.useVGPRIndexMode();
 
   MachineBasicBlock *BB = MI.getParent();
   const DebugLoc &DL = MI.getDebugLoc();
 
   if (!IndexMode) {

 
 InstructionSelector::ComplexRendererFns
 AMDGPUInstructionSelector::selectSMRDBufferSgprImm(MachineOperand &Root) const {
   // Match the (soffset + offset) pair as a 32-bit register base and
   // an immediate offset.
   Register SOffset;
   unsigned Offset;
   std::tie(SOffset, Offset) =
-      AMDGPU::getBaseWithConstantOffset(*MRI, Root.getReg());
+      AMDGPU::getBaseWithConstantOffset(*MRI, Root.getReg(), KnownBits);
   if (!SOffset)
     return None;
 
   Optional<int64_t> EncodedOffset =
       AMDGPU::getSMRDEncodedOffset(STI, Offset, /* IsBuffer */ true);
   if (!EncodedOffset)
     return None;
 

llvm/test/CodeGen/AMDGPU/amdgcn-load-offset-from-reg.ll

 ; GISEL: S_BUFFER_LOAD_DWORD_SGPR_IMM %[[BASE]], %[[OFFSET]], 77,
 define amdgpu_cs void @test_buffer_load_sgpr_plus_imm_offset(<4 x i32> inreg %base, i32 inreg %i, i32 addrspace(1)* inreg %out) {
   %off = add nuw nsw i32 %i, 77
   %v = call i32 @llvm.amdgcn.s.buffer.load.i32(<4 x i32> %base, i32 %off, i32 0)
   store i32 %v, i32 addrspace(1)* %out, align 4
   ret void
 }
 
-; TODO: Select S_BUFFER_LOAD_DWORD_SGPR_IMM here.
 ; GCN-LABEL: name: test_buffer_load_sgpr_or_imm_offset
 ; SDAG-DAG: %[[BASE0:.*]]:sgpr_32 = COPY $sgpr0
 ; SDAG-DAG: %[[BASE1:.*]]:sgpr_32 = COPY $sgpr1
 ; SDAG-DAG: %[[BASE2:.*]]:sgpr_32 = COPY $sgpr2
 ; SDAG-DAG: %[[BASE3:.*]]:sgpr_32 = COPY $sgpr3
 ; SDAG-DAG: %[[INDEX:.*]]:sgpr_32 = COPY $sgpr4
 ; SDAG-DAG: %[[SHIFT:.*]]:sreg_32 = S_LSHL_B32 %[[INDEX]],
-; SDAG-DAG: %[[OR:.*]]:sreg_32 = S_OR_B32 killed %[[SHIFT]],
 ; SDAG-DAG: %[[BASE:.*]]:sgpr_128 = REG_SEQUENCE %[[BASE0]], %subreg.sub0, %[[BASE1]], %subreg.sub1, %[[BASE2]], %subreg.sub2, %[[BASE3]], %subreg.sub3
-; SDAG: S_BUFFER_LOAD_DWORD_SGPR killed %[[BASE]], killed %[[OR]], 0
+; SDAG: S_BUFFER_LOAD_DWORD_SGPR_IMM killed %[[BASE]], killed %[[SHIFT]], 5,
 ; GISEL-DAG: %[[BASE0:.*]]:sreg_32 = COPY $sgpr0
 ; GISEL-DAG: %[[BASE1:.*]]:sreg_32 = COPY $sgpr1
 ; GISEL-DAG: %[[BASE2:.*]]:sreg_32 = COPY $sgpr2
 ; GISEL-DAG: %[[BASE3:.*]]:sreg_32 = COPY $sgpr3
 ; GISEL-DAG: %[[INDEX:.*]]:sreg_32 = COPY $sgpr4
 ; GISEL-DAG: %[[SHIFT:.*]]:sreg_32 = S_LSHL_B32 %[[INDEX]],
-; GISEL-DAG: %[[OR:.*]]:sreg_32 = S_OR_B32 %[[SHIFT]],
 ; GISEL-DAG: %[[BASE:.*]]:sgpr_128 = REG_SEQUENCE %[[BASE0]], %subreg.sub0, %[[BASE1]], %subreg.sub1, %[[BASE2]], %subreg.sub2, %[[BASE3]], %subreg.sub3
-; GISEL: S_BUFFER_LOAD_DWORD_SGPR_IMM %[[BASE]], %[[OR]], 0,
+; GISEL: S_BUFFER_LOAD_DWORD_SGPR_IMM %[[BASE]], %[[SHIFT]], 5,
 define amdgpu_cs void @test_buffer_load_sgpr_or_imm_offset(<4 x i32> inreg %base, i32 inreg %i, i32 addrspace(1)* inreg %out) {
   %shift = shl i32 %i, 7
   %off = or i32 %shift, 5
   %v = call i32 @llvm.amdgcn.s.buffer.load.i32(<4 x i32> %base, i32 %off, i32 0)
   store i32 %v, i32 addrspace(1)* %out, align 4
   ret void
 }
 
 declare void @llvm.amdgcn.raw.buffer.store.v4i32(<4 x i32>, <4 x i32>, i32, i32, i32 immarg) #1
 
 declare i32 @llvm.amdgcn.s.buffer.load.i32(<4 x i32>, i32, i32 immarg) nounwind readnone willreturn
 
 ; Function Attrs: nounwind readnone speculatable
 declare i32 @llvm.amdgcn.reloc.constant(metadata) #3
 
 ; Function Attrs: nounwind readnone speculatable
 declare i64 @llvm.amdgcn.s.getpc() #3