[AMDGPU][GlobalISel] Improve constant offset lookup for llvm.amdgcn.s.buffer

[mbrkusanin]

Summary

Looking up constant offset for this intrinsic is done during RegBankSelect
so the constant might find itself behind a copy from sgpr to vgpr.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUGlobalISelUtils.cpp

51

This shouldn't add a second match pattern for a copy. That should be hidden away somehow. In this case I'd say the problem is regbankselect isn't picking the correct bank for the constant in the first place since we don't consider the users right now. We've been working around this on a case by case basis, but I think we should try to fix this.

Also arguably mi_match should look through copies, although in the presence of cross reg bank copies that's more questionable

[mbrkusanin]

I revisited this again to try to find a way to eliminate these cross reg bank copies but I wasn't able to find an efficient solution. Looking through uses of dst of G_CONSTANT at the time of selecting a bank for that instruction does not help since RegBank is not working backwards.

One idea that does work is to instead of making a cross reg bank copy to a constant to simply recreate a new G_CONSTANT with required bank (change is in RegBankSelect::repairReg). This can help with the current problem but produces less optimal solutions in a lot of tests. In some cases instead of an SGPR, an extra VPGR will be used instead. And in cases where same constant needs to be used as both SGPR and VGPR there can also be an additional s/v_mov instruction.

Here we can at least hide copy check inside matcher.

llvm/include/llvm/CodeGen/GlobalISel/MIPatternMatch.h

 
 template <typename SubPat>
 inline OneNonDBGUse_match<SubPat> m_OneNonDBGUse(const SubPat &SP) {
   return SP;
 }
 
 struct ConstantMatch {
   int64_t &CR;
-  ConstantMatch(int64_t &C) : CR(C) {}
+  bool LookThroughCopy;
+  ConstantMatch(int64_t &C, bool LT) : CR(C), LookThroughCopy(LT) {}
   bool match(const MachineRegisterInfo &MRI, Register Reg) {
+    if (LookThroughCopy) {
+      auto MI = MRI.getVRegDef(Reg);
+      if (MI && MI->getOpcode() == TargetOpcode::COPY)
+        Reg = MI->getOperand(1).getReg();
+    }
     if (auto MaybeCst = getConstantVRegSExtVal(Reg, MRI)) {
       CR = *MaybeCst;
       return true;
     }
     return false;
   }
 };
 
-inline ConstantMatch m_ICst(int64_t &Cst) { return ConstantMatch(Cst); }
+inline ConstantMatch m_ICst(int64_t &Cst, bool LookThroughCopy = false) {
+  return ConstantMatch(Cst, LookThroughCopy);
+}
 
 struct ICstRegMatch {
   Register &CR;
   ICstRegMatch(Register &C) : CR(C) {}
   bool match(const MachineRegisterInfo &MRI, Register Reg) {
     if (auto MaybeCst = getConstantVRegValWithLookThrough(
             Reg, MRI, /*LookThroughInstrs*/ true,
             /*HandleFConstants*/ false)) {

llvm/lib/Target/AMDGPU/AMDGPUGlobalISelUtils.cpp

else
       Offset = Op.getCImm()->getZExtValue();
 
     return std::make_pair(Register(), Offset);
   }
 
   int64_t Offset;
   if (Def->getOpcode() == TargetOpcode::G_ADD) {
     // TODO: Handle G_OR used for add case
-    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))))
+    if (mi_match(Def->getOperand(2).getReg(), MRI, m_ICst(Offset, true)))
       return std::make_pair(Def->getOperand(1).getReg(), 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)))) {
+                 m_GPtrAdd(m_MInstr(Base), m_ICst(Offset, true)))) {
       // 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);
 
       // Register returned here will be of pointer type.
       return std::make_pair(Base->getOperand(0).getReg(), Offset);
     }
   }
 
   return std::make_pair(Reg, 0);
 }
 
 bool AMDGPU::isLegalVOP3PShuffleMask(ArrayRef<int> Mask) {
   assert(Mask.size() == 2);
 
   // If one half is undef, the other is trivially in the same reg.
   if (Mask[0] == -1 || Mask[1] == -1)
     return true;
   return (Mask[0] & 2) == (Mask[1] & 2);
 }

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

static unsigned setBufferOffsets(MachineIRBuilder &B,
   }
 
   Register Base;
   unsigned Offset;
 
   std::tie(Base, Offset) =
       AMDGPU::getBaseWithConstantOffset(*MRI, CombinedOffset);
 
+  // If BaseReg is a pointer, convert it to int.
+  if (MRI->getType(Base).isPointer()) {
+    const RegisterBank *BaseBank = RBI.getRegBank(Base, *MRI, *RBI.TRI);
+    Base = B.buildPtrToInt(MRI->getType(CombinedOffset), Base).getReg(0);
+    MRI->setRegBank(Base, *BaseBank);
+  }
+
   uint32_t SOffset, ImmOffset;
   if ((int)Offset > 0 && AMDGPU::splitMUBUFOffset(Offset, SOffset, ImmOffset,
                                                   &RBI.Subtarget, Alignment)) {
     if (RBI.getRegBank(Base, *MRI, *RBI.TRI) == &AMDGPU::VGPRRegBank) {
       VOffsetReg = Base;
       SOffsetReg = B.buildConstant(S32, SOffset).getReg(0);
       B.getMRI()->setRegBank(SOffsetReg, AMDGPU::SGPRRegBank);
       InstOffsetVal = ImmOffset;

llvm/test/CodeGen/AMDGPU/GlobalISel/merge-s-buffer-loads.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -global-isel -march=amdgcn -verify-machineinstrs -o - %s | FileCheck %s
+
+declare i32 @llvm.amdgcn.s.buffer.load.i32(<4 x i32>, i32, i32 immarg)
+declare void @llvm.amdgcn.exp.i32(i32 immarg, i32 immarg, i32, i32, i32, i32, i1, i1)
+
+define amdgpu_cs void @test1(i32 %index, <4 x i32> inreg %desc, <4 x i32> addrspace(6)* inreg %array) {
+; CHECK-LABEL: test1:
+; CHECK:       ; %bb.0: ; %.entry
+; CHECK-NEXT:    v_lshlrev_b32_e32 v0, 4, v0
+; CHECK-NEXT:    v_add_i32_e32 v0, vcc, s4, v0
+; CHECK-NEXT:    buffer_load_dwordx4 v[0:3], v0, s[0:3], 0 offen
+; CHECK-NEXT:    s_waitcnt vmcnt(0)
+; CHECK-NEXT:    exp mrt0 off, off, off, off
+; CHECK-NEXT:    s_endpgm
+.entry:
+  %ep11 = getelementptr <4 x i32>, <4 x i32> addrspace(6)* %array, i32 %index, i32 0
+  %ei11 = ptrtoint i32 addrspace(6)* %ep11 to i32
+  %el11 = call i32 @llvm.amdgcn.s.buffer.load.i32(<4 x i32> %desc, i32 %ei11, i32 0)
+
+  %ep12 = getelementptr <4 x i32>, <4 x i32> addrspace(6)* %array, i32 %index, i32 1
+  %ei12 = ptrtoint i32 addrspace(6)* %ep12 to i32
+  %el12 = call i32 @llvm.amdgcn.s.buffer.load.i32(<4 x i32> %desc, i32 %ei12, i32 0)
+
+  %ep13 = getelementptr <4 x i32>, <4 x i32> addrspace(6)* %array, i32 %index, i32 2
+  %ei13 = ptrtoint i32 addrspace(6)* %ep13 to i32
+  %el13 = call i32 @llvm.amdgcn.s.buffer.load.i32(<4 x i32> %desc, i32 %ei13, i32 0)
+
+  %ep14 = getelementptr <4 x i32>, <4 x i32> addrspace(6)* %array, i32 %index, i32 3
+  %ei14 = ptrtoint i32 addrspace(6)* %ep14 to i32
+  %el14 = call i32 @llvm.amdgcn.s.buffer.load.i32(<4 x i32> %desc, i32 %ei14, i32 0)
+
+  call void @llvm.amdgcn.exp.i32(i32 0, i32 0, i32 %el11, i32 %el12, i32 %el13, i32 %el14, i1 false, i1 false)
+
+  ret void
+}
+
+define amdgpu_cs void @test2(i32 %index, <4 x i32> addrspace(6)* inreg %array) {
+; CHECK-LABEL: test2:
+; CHECK:       ; %bb.0: ; %.entry
+; CHECK-NEXT:    v_lshlrev_b32_e32 v0, 4, v0
+; CHECK-NEXT:    v_add_i32_e32 v0, vcc, s0, v0
+; CHECK-NEXT:    buffer_load_dwordx4 v[0:3], v0, s[0:3], 0 offen
+; CHECK-NEXT:    s_waitcnt vmcnt(0)
+; CHECK-NEXT:    exp mrt0 off, off, off, off
+; CHECK-NEXT:    s_endpgm
+.entry:
+  %ep11 = getelementptr <4 x i32>, <4 x i32> addrspace(6)* %array, i32 %index, i32 0
+  %ei11 = ptrtoint i32 addrspace(6)* %ep11 to i32
+  %el11 = call i32 @llvm.amdgcn.s.buffer.load.i32(<4 x i32> undef, i32 %ei11, i32 0)
+
+  %ep12 = getelementptr <4 x i32>, <4 x i32> addrspace(6)* %array, i32 %index, i32 1
+  %ei12 = ptrtoint i32 addrspace(6)* %ep12 to i32
+  %el12 = call i32 @llvm.amdgcn.s.buffer.load.i32(<4 x i32> undef, i32 %ei12, i32 0)
+
+  %ep13 = getelementptr <4 x i32>, <4 x i32> addrspace(6)* %array, i32 %index, i32 2
+  %ei13 = ptrtoint i32 addrspace(6)* %ep13 to i32
+  %el13 = call i32 @llvm.amdgcn.s.buffer.load.i32(<4 x i32> undef, i32 %ei13, i32 0)
+
+  %ep14 = getelementptr <4 x i32>, <4 x i32> addrspace(6)* %array, i32 %index, i32 3
+  %ei14 = ptrtoint i32 addrspace(6)* %ep14 to i32
+  %el14 = call i32 @llvm.amdgcn.s.buffer.load.i32(<4 x i32> undef, i32 %ei14, i32 0)
+
+  call void @llvm.amdgcn.exp.i32(i32 0, i32 0, i32 %el11, i32 %el12, i32 %el13, i32 %el14, i1 false, i1 false)
+
+  ret void
+}

llvm/unittests/CodeGen/GlobalISel/PatternMatchTest.cpp

if (!TM)
     return;
   auto MIBCst = B.buildConstant(LLT::scalar(64), 42);
   int64_t Cst;
   bool match = mi_match(MIBCst.getReg(0), *MRI, m_ICst(Cst));
   EXPECT_TRUE(match);
   EXPECT_EQ(Cst, 42);
 }
 
+TEST_F(AArch64GISelMITest, MatchIntConstantWithLookThrough) {
+  setUp();
+  if (!TM)
+    return;
+  auto MIBCst = B.buildConstant(LLT::scalar(64), 42);
+  auto MIBCpy = B.buildCopy(LLT::scalar(64), MIBCst);
+  int64_t Cst;
+  bool match = mi_match(MIBCpy.getReg(0), *MRI, m_ICst(Cst, true));
+  EXPECT_TRUE(match);
+  EXPECT_EQ(Cst, 42);
+}
+
 TEST_F(AArch64GISelMITest, MatchIntConstantRegister) {
   setUp();
   if (!TM)
     return;
   auto MIBCst = B.buildConstant(LLT::scalar(64), 42);
   Register Src0;
   bool match = mi_match(MIBCst.getReg(0), *MRI, m_ICst(Src0));
   EXPECT_TRUE(match);