[AMDGPU][GlobalISel] Support register offsets for SMRDs.

[arsenm]

Should add some MIR select tests

llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp

3862–3873

Isn't GEPInfo.Imm > 0 redundant with isUInt<32>?

3877

Should use matchZeroExtendFromS32 to handle all the extended cases. Also this isn't checking the match return

[kosarev]

Updated to use matchZeroExtendFromS32().

[kosarev]

llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp

3862–3873

I believe isUInt<32>() allows zero?

3877

Changed.

Also this isn't checking the match return

Yes, intentionally. In the previous version, as long as the operand originates from a 32-bit scalar, we don't really care if zero-expansion matched. Though I couldn't provide a reproducer for the case without expansion.

[arsenm]

Should add some selector MIR tests

llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp

3862–3873

Clearer to change to != 0 then

[kosarev]

Added an MIR instruction selection test and changed the if () condition
as suggested.

[foad]

LGTM with a tweak to inst-select-load-smrd.mir.

llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp

3240

Not your fault, but this seems to be missing checks (or asserts) on the number of operands and the type of the operands of the G_MERGE_VALUES.

3854–3855

I don't understand this comment. What cases are we still missing?

llvm/test/CodeGen/AMDGPU/GlobalISel/inst-select-load-smrd.mir

232

Please put these checks before the function, not after. (But wow the existing checks in this file are really messy.)

[kosarev]

Moved the check lines in the test before the function.

[kosarev]

llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp

3240

Indeed. Will address this with a separate patch.

3854–3855

I guess this refers to i64 values that are known to fit 32 bits, such as (x & 3) << 2, etc.

[foad]

LGTM.

llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp

 static Register matchZeroExtendFromS32(MachineRegisterInfo &MRI, Register Reg) {
   Register ZExtSrc;
   if (mi_match(Reg, MRI, m_GZExt(m_Reg(ZExtSrc))))
     return MRI.getType(ZExtSrc) == LLT::scalar(32) ? ZExtSrc : Register();
 
   // Match legalized form %zext = G_MERGE_VALUES (s32 %x), (s32 0)
   const MachineInstr *Def = getDefIgnoringCopies(Reg, MRI);
   if (Def->getOpcode() != AMDGPU::G_MERGE_VALUES)
-    return false;
+    return Register();
 
   if (mi_match(Def->getOperand(2).getReg(), MRI, m_ZeroInt())) {
     return Def->getOperand(1).getReg();
   }
 
   return Register();
 }
 
 bool AMDGPUInstructionSelector::selectGlobalLoadLds(MachineInstr &MI) const{
   unsigned Opc;

 AMDGPUInstructionSelector::selectSmrdSgpr(MachineOperand &Root) const {
   MachineInstr *MI = Root.getParent();
   MachineBasicBlock *MBB = MI->getParent();
 
   SmallVector<GEPInfo, 4> AddrInfo;
   getAddrModeInfo(*MI, *MRI, AddrInfo);
 
   // FIXME: We should shrink the GEP if the offset is known to be <= 32-bits,
-  // then we can select all ptr + 32-bit offsets not just immediate offsets.
-  if (AddrInfo.empty() || AddrInfo[0].SgprParts.size() != 1)
+  // then we can select all ptr + 32-bit offsets.
+  if (AddrInfo.empty())
     return None;
 
   const GEPInfo &GEPInfo = AddrInfo[0];
-  // SGPR offset is unsigned.
-  if (!GEPInfo.Imm || GEPInfo.Imm < 0 || !isUInt<32>(GEPInfo.Imm))
-    return None;
+  Register PtrReg = GEPInfo.SgprParts[0];
 
+  // SGPR offset is unsigned.
+  if (AddrInfo[0].SgprParts.size() == 1 && isUInt<32>(GEPInfo.Imm) &&
+      GEPInfo.Imm != 0) {
     // If we make it this far we have a load with an 32-bit immediate offset.
     // It is OK to select this using a sgpr offset, because we have already
     // failed trying to select this load into one of the _IMM variants since
     // the _IMM Patterns are considered before the _SGPR patterns.
-  Register PtrReg = GEPInfo.SgprParts[0];
     Register OffsetReg = MRI->createVirtualRegister(&AMDGPU::SReg_32RegClass);
     BuildMI(*MBB, MI, MI->getDebugLoc(), TII.get(AMDGPU::S_MOV_B32), OffsetReg)
-          .addImm(GEPInfo.Imm);
-  return {{
-    [=](MachineInstrBuilder &MIB) { MIB.addReg(PtrReg); },
-    [=](MachineInstrBuilder &MIB) { MIB.addReg(OffsetReg); }
-  }};
+    return {{[=](MachineInstrBuilder &MIB) { MIB.addReg(PtrReg); },
+             [=](MachineInstrBuilder &MIB) { MIB.addReg(OffsetReg); }}};
+  }
+
+  if (AddrInfo[0].SgprParts.size() == 2 && GEPInfo.Imm == 0) {
+    if (Register OffsetReg =
+            matchZeroExtendFromS32(*MRI, GEPInfo.SgprParts[1])) {
+      return {{[=](MachineInstrBuilder &MIB) { MIB.addReg(PtrReg); },
+               [=](MachineInstrBuilder &MIB) { MIB.addReg(OffsetReg); }}};
+    }
+  }
+
+  return None;
 }
 
 std::pair<Register, int>
 AMDGPUInstructionSelector::selectFlatOffsetImpl(MachineOperand &Root,
                                                 uint64_t FlatVariant) const {
   MachineInstr *MI = Root.getParent();
 
   auto Default = std::make_pair(Root.getReg(), 0);

llvm/test/CodeGen/AMDGPU/GlobalISel/inst-select-load-smrd.mir

 # RUN: llc -march=amdgcn -mcpu=tahiti -run-pass=instruction-select -verify-machineinstrs %s -o - | FileCheck %s -check-prefixes=GCN,SI,SICI,SIVI
 # RUN: llc -march=amdgcn -mcpu=hawaii -run-pass=instruction-select -verify-machineinstrs %s -o - | FileCheck %s -check-prefixes=GCN,CI,SICI
 # RUN: llc -march=amdgcn -mcpu=fiji -run-pass=instruction-select -verify-machineinstrs %s -o - | FileCheck %s -check-prefixes=GCN,VI,SIVI
 
 --- |
   define amdgpu_kernel void @smrd_imm(i32 addrspace(4)* %const0) { ret void }
   define amdgpu_kernel void @smrd_wide() { ret void }
   define amdgpu_kernel void @constant_address_positive() { ret void }
+  define amdgpu_kernel void @smrd_sgpr() { ret void }
 ...
 ---
 
 name:            smrd_imm
 legalized:       true
 regBankSelected: true
 
 # GCN: body:
body: |
   bb.0:
     liveins: $sgpr0_sgpr1, $vgpr2_vgpr3
     %0:sgpr(p4) = G_CONSTANT i64 44
     %1:sgpr(s64) = G_CONSTANT i64 64
     %2:sgpr(p4) = G_PTR_ADD %0, %1
     %3:sgpr(s32) = G_LOAD %2 :: (dereferenceable invariant load (s32), align 4, addrspace 4)
     S_ENDPGM 0, implicit %3
 ...
+
+---
+
+# Test a load with a register offset.
+# GCN-LABEL: name: smrd_sgpr{{$}}
+# GCN: S_LOAD_DWORD_SGPR %0, %1, 0
+
+name:            smrd_sgpr
+legalized:       true
+regBankSelected: true
+
+body: |
+  bb.0:
+    liveins: $sgpr0_sgpr1, $sgpr2
+    %0:sgpr(p4) = COPY $sgpr0_sgpr1
+    %1:sgpr(s32) = COPY $sgpr2
+    %2:sgpr(s64) = G_ZEXT %1:sgpr(s32)
+    %4:sgpr(p4) = G_PTR_ADD %0, %2
+    %5:sgpr(s32) = G_LOAD %4 :: (dereferenceable invariant load (s32), align 4, addrspace 4)
+    S_ENDPGM 0, implicit %5
+...

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

 ; Test that DAG->DAG ISel is able to pick up the S_LOAD_DWORDX4_SGPR instruction that fetches the offset
 ; from a register.
 
 ; RUN: llc -march=amdgcn -verify-machineinstrs -stop-after=amdgpu-isel -o - %s | FileCheck -check-prefix=GCN %s
+; RUN: llc -march=amdgcn -global-isel -verify-machineinstrs -stop-after=amdgpu-isel -o - %s | FileCheck -check-prefix=GISEL %s
 
 ; GCN: %[[OFFSET:[0-9]+]]:sreg_32 = S_MOV_B32 target-flags(amdgpu-abs32-lo) @DescriptorBuffer
 ; GCN: %{{[0-9]+}}:sgpr_128 = S_LOAD_DWORDX4_SGPR killed %{{[0-9]+}}, killed %[[OFFSET]], 0 :: (invariant load (s128) from %ir.13, addrspace 4)
 
+; GISEL: $[[OFFSET:.*]] = S_MOV_B32 target-flags(amdgpu-abs32-lo) @DescriptorBuffer
+; GISEL: S_LOAD_DWORDX4_SGPR killed renamable {{.*}}, killed renamable $[[OFFSET]], 0 :: (invariant load (<4 x s32>) from {{.*}}, addrspace 4)
+
 define amdgpu_cs void @test_load_zext(i32 inreg %0, i32 inreg %1, i32 inreg %resNode0, i32 inreg %resNode1, <3 x i32> inreg %2, i32 inreg %3, <3 x i32> %4) local_unnamed_addr #2 {
 .entry:
   %5 = call i64 @llvm.amdgcn.s.getpc() #3
   %6 = bitcast i64 %5 to <2 x i32>
   %7 = insertelement <2 x i32> %6, i32 %resNode0, i32 0
   %8 = bitcast <2 x i32> %7 to i64
   %9 = inttoptr i64 %8 to [4294967295 x i8] addrspace(4)*
   %10 = call i32 @llvm.amdgcn.reloc.constant(metadata !4)