[AMDGPU][GlobalISel] Handle G_PTR_ADD when looking for constant offset

[mbrkusanin]

Summary

Look through G_PTRTOINT and G_PTR_ADD nodes when looking for constant
offset for buffer stores. This also helps with merging of these instructions
later on.

[arsenm]

I don't like the test name, because there isn't a .merge intrinsic

[mbrkusanin]

• Rebase and rename test

llvm/lib/Target/AMDGPU/AMDGPUGlobalISelUtils.cpp

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);
   }
 
+  // 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);
+
+      // 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/AMDGPULegalizerInfo.cpp

 // offset and figures out how to split it between voffset and immoffset.
 std::tuple<Register, unsigned, unsigned>
 AMDGPULegalizerInfo::splitBufferOffsets(MachineIRBuilder &B,
                                         Register OrigOffset) const {
   const unsigned MaxImm = 4095;
   Register BaseReg;
   unsigned TotalConstOffset;
   const LLT S32 = LLT::scalar(32);
+  MachineRegisterInfo &MRI = *B.getMRI();
 
   std::tie(BaseReg, TotalConstOffset) =
-      AMDGPU::getBaseWithConstantOffset(*B.getMRI(), OrigOffset);
+      AMDGPU::getBaseWithConstantOffset(MRI, OrigOffset);
 
   unsigned ImmOffset = TotalConstOffset;
 
+  // If BaseReg is a pointer, convert it to int.
+  if (MRI.getType(BaseReg).isPointer())
+    BaseReg = B.buildPtrToInt(MRI.getType(OrigOffset), BaseReg).getReg(0);
+
   // If the immediate value is too big for the immoffset field, put the value
   // and -4096 into the immoffset field so that the value that is copied/added
   // for the voffset field is a multiple of 4096, and it stands more chance
   // of being CSEd with the copy/add for another similar load/store.
   // However, do not do that rounding down to a multiple of 4096 if that is a
   // negative number, as it appears to be illegal to have a negative offset
   // in the vgpr, even if adding the immediate offset makes it positive.
   unsigned Overflow = ImmOffset & ~MaxImm;

llvm/test/CodeGen/AMDGPU/GlobalISel/merge-buffer-stores.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
+
+define amdgpu_cs void @test1(i32 %arg1, <4 x i32> inreg %arg2, i32, i32 addrspace(6)* inreg %arg3) {
+; CHECK-LABEL: test1:
+; CHECK:       ; %bb.0: ; %.entry
+; CHECK-NEXT:    v_and_b32_e32 v3, 0x3ffffffc, v0
+; CHECK-NEXT:    v_mov_b32_e32 v0, 11
+; CHECK-NEXT:    v_mov_b32_e32 v1, 22
+; CHECK-NEXT:    v_mov_b32_e32 v2, 33
+; CHECK-NEXT:    v_lshlrev_b32_e32 v3, 2, v3
+; CHECK-NEXT:    v_add_i32_e32 v4, vcc, s4, v3
+; CHECK-NEXT:    v_mov_b32_e32 v3, 44
+; CHECK-NEXT:    buffer_store_dwordx4 v[0:3], v4, s[0:3], 0 offen
+; CHECK-NEXT:    s_endpgm
+.entry:
+  %bs1 = and i32 %arg1, 1073741820
+  %ep1 = getelementptr i32, i32 addrspace(6)* %arg3, i32 %bs1
+  %ad1 = ptrtoint i32 addrspace(6)* %ep1 to i32
+  call void @llvm.amdgcn.raw.buffer.store.i32(i32 11, <4 x i32> %arg2, i32 %ad1, i32 0, i32 0)
+
+  %bs2 = or i32 %bs1, 1
+  %ep2 = getelementptr i32, i32 addrspace(6)* %arg3, i32 %bs2
+  %ad2 = ptrtoint i32 addrspace(6)* %ep2 to i32
+  call void @llvm.amdgcn.raw.buffer.store.i32(i32 22, <4 x i32> %arg2, i32 %ad2, i32 0, i32 0)
+
+  %bs3 = or i32 %bs1, 2
+  %ep3 = getelementptr i32, i32 addrspace(6)* %arg3, i32 %bs3
+  %ad3 = ptrtoint i32 addrspace(6)* %ep3 to i32
+  call void @llvm.amdgcn.raw.buffer.store.i32(i32 33, <4 x i32> %arg2, i32 %ad3, i32 0, i32 0)
+
+  %bs4 = or i32 %bs1, 3
+  %ep4 = getelementptr i32, i32 addrspace(6)* %arg3, i32 %bs4
+  %ad4 = ptrtoint i32 addrspace(6)* %ep4 to i32
+  call void @llvm.amdgcn.raw.buffer.store.i32(i32 44, <4 x i32> %arg2, i32 %ad4, i32 0, i32 0)
+
+  ret void
+}
+
+define amdgpu_cs void @test2(i32 %arg1, <4 x i32> inreg %arg2) {
+; CHECK-LABEL: test2:
+; CHECK:       ; %bb.0: ; %.entry
+; CHECK-NEXT:    v_and_b32_e32 v3, 0x3ffffffc, v0
+; CHECK-NEXT:    v_mov_b32_e32 v0, 11
+; CHECK-NEXT:    v_mov_b32_e32 v1, 22
+; CHECK-NEXT:    v_mov_b32_e32 v2, 33
+; CHECK-NEXT:    v_lshlrev_b32_e32 v4, 2, v3
+; CHECK-NEXT:    v_mov_b32_e32 v3, 44
+; CHECK-NEXT:    buffer_store_dwordx4 v[0:3], v4, s[0:3], 0 offen
+; CHECK-NEXT:    s_endpgm
+.entry:
+  %bs1 = and i32 %arg1, 1073741820
+  %ep1 = getelementptr <{ [64 x i32] }>, <{ [64 x i32] }> addrspace(6)* null, i32 0, i32 0, i32 %bs1
+  %ad1 = ptrtoint i32 addrspace(6)* %ep1 to i32
+  call void @llvm.amdgcn.raw.buffer.store.i32(i32 11, <4 x i32> %arg2, i32 %ad1, i32 0, i32 0)
+
+  %bs2 = or i32 %bs1, 1
+  %ep2 = getelementptr <{ [64 x i32] }>, <{ [64 x i32] }> addrspace(6)* null, i32 0, i32 0, i32 %bs2
+  %ad2 = ptrtoint i32 addrspace(6)* %ep2 to i32
+  call void @llvm.amdgcn.raw.buffer.store.i32(i32 22, <4 x i32> %arg2, i32 %ad2, i32 0, i32 0)
+
+  %bs3 = or i32 %bs1, 2
+  %ep3 = getelementptr <{ [64 x i32] }>, <{ [64 x i32] }> addrspace(6)* null, i32 0, i32 0, i32 %bs3
+  %ad3 = ptrtoint i32 addrspace(6)* %ep3 to i32
+  call void @llvm.amdgcn.raw.buffer.store.i32(i32 33, <4 x i32> %arg2, i32 %ad3, i32 0, i32 0)
+
+  %bs4 = or i32 %bs1, 3
+  %ep4 = getelementptr <{ [64 x i32] }>, <{ [64 x i32] }> addrspace(6)* null, i32 0, i32 0, i32 %bs4
+  %ad4 = ptrtoint i32 addrspace(6)* %ep4 to i32
+  call void @llvm.amdgcn.raw.buffer.store.i32(i32 44, <4 x i32> %arg2, i32 %ad4, i32 0, i32 0)
+
+  ret void
+}
+
+declare void @llvm.amdgcn.raw.buffer.store.i32(i32, <4 x i32>, i32, i32, i32 immarg)