AMDGPU: Merge BUFFER_LOAD_DWORD_OFFSET into x2, x4

[mareko]

Summary

Only 3 (out of 48486) shaders are affected.

[mareko]

fix up merge-stores.ll

[mareko]

Ping

lib/Target/AMDGPU/SILoadStoreOptimizer.cpp

 
 namespace {
 
 class SILoadStoreOptimizer : public MachineFunctionPass {
   enum InstClassEnum {
     DS_READ_WRITE,
     S_BUFFER_LOAD_IMM,
     BUFFER_LOAD_OFFEN,
+    BUFFER_LOAD_OFFSET,
   };
 
   struct CombineInfo {
     MachineBasicBlock::iterator I;
     MachineBasicBlock::iterator Paired;
     unsigned EltSize;
     unsigned Offset0;
     unsigned Offset1;
private:
   unsigned CreatedX2;
 
   static bool offsetsCanBeCombined(CombineInfo &CI);
 
   bool findMatchingInst(CombineInfo &CI);
   MachineBasicBlock::iterator mergeRead2Pair(CombineInfo &CI);
   MachineBasicBlock::iterator mergeWrite2Pair(CombineInfo &CI);
   MachineBasicBlock::iterator mergeSBufferLoadImmPair(CombineInfo &CI);
-  MachineBasicBlock::iterator mergeBufferLoadOffenPair(CombineInfo &CI);
+  MachineBasicBlock::iterator mergeBufferLoadPair(CombineInfo &CI);
 
 public:
   static char ID;
 
   SILoadStoreOptimizer() : MachineFunctionPass(ID) {
     initializeSILoadStoreOptimizerPass(*PassRegistry::getPassRegistry());
   }
 
bool SILoadStoreOptimizer::offsetsCanBeCombined(CombineInfo &CI) {
 
   unsigned EltOffset0 = CI.Offset0 / CI.EltSize;
   unsigned EltOffset1 = CI.Offset1 / CI.EltSize;
   CI.UseST64 = false;
   CI.BaseOff = 0;
 
   // SMEM offsets must be consecutive.
   if (CI.InstClass == S_BUFFER_LOAD_IMM ||
-      CI.InstClass == BUFFER_LOAD_OFFEN) {
+      CI.InstClass == BUFFER_LOAD_OFFEN ||
+      CI.InstClass == BUFFER_LOAD_OFFSET) {
     unsigned Diff = CI.IsX2 ? 2 : 1;
     return (EltOffset0 + Diff == EltOffset1 ||
             EltOffset1 + Diff == EltOffset0) &&
            CI.GLC0 == CI.GLC1 &&
            (CI.InstClass == S_BUFFER_LOAD_IMM || CI.SLC0 == CI.SLC1);
   }
 
   // If the offset in elements doesn't fit in 8-bits, we might be able to use
bool SILoadStoreOptimizer::findMatchingInst(CombineInfo &CI) {
   case S_BUFFER_LOAD_IMM:
     AddrOpName[NumAddresses++] = AMDGPU::OpName::sbase;
     break;
   case BUFFER_LOAD_OFFEN:
     AddrOpName[NumAddresses++] = AMDGPU::OpName::srsrc;
     AddrOpName[NumAddresses++] = AMDGPU::OpName::vaddr;
     AddrOpName[NumAddresses++] = AMDGPU::OpName::soffset;
     break;
+  case BUFFER_LOAD_OFFSET:
+    AddrOpName[NumAddresses++] = AMDGPU::OpName::srsrc;
+    AddrOpName[NumAddresses++] = AMDGPU::OpName::soffset;
+    break;
   default:
     llvm_unreachable("invalid InstClass");
   }
 
   for (unsigned i = 0; i < NumAddresses; i++) {
     AddrIdx[i] = AMDGPU::getNamedOperandIdx(CI.I->getOpcode(), AddrOpName[i]);
     AddrReg[i] = &CI.I->getOperand(AddrIdx[i]);
 
if (Match) {
       CI.Paired = MBBI;
 
       if (CI.InstClass == DS_READ_WRITE) {
         CI.Offset0 &= 0xffff;
         CI.Offset1 &= 0xffff;
       } else {
         CI.GLC0 = TII->getNamedOperand(*CI.I, AMDGPU::OpName::glc)->getImm();
         CI.GLC1 = TII->getNamedOperand(*MBBI, AMDGPU::OpName::glc)->getImm();
-        if (CI.InstClass == BUFFER_LOAD_OFFEN) {
+        if (CI.InstClass != S_BUFFER_LOAD_IMM) {
           CI.SLC0 = TII->getNamedOperand(*CI.I, AMDGPU::OpName::slc)->getImm();
           CI.SLC1 = TII->getNamedOperand(*MBBI, AMDGPU::OpName::slc)->getImm();
         }
       }
 
       // Check both offsets fit in the reduced range.
       // We also need to go through the list of instructions that we plan to
       // move and make sure they are all safe to move down past the merged
MachineBasicBlock::iterator SILoadStoreOptimizer::mergeSBufferLoadImmPair(
   moveInstsAfter(Copy1, CI.InstsToMove);
 
   MachineBasicBlock::iterator Next = std::next(CI.I);
   CI.I->eraseFromParent();
   CI.Paired->eraseFromParent();
   return Next;
 }
 
-MachineBasicBlock::iterator SILoadStoreOptimizer::mergeBufferLoadOffenPair(
+MachineBasicBlock::iterator SILoadStoreOptimizer::mergeBufferLoadPair(
   CombineInfo &CI) {
   MachineBasicBlock *MBB = CI.I->getParent();
   DebugLoc DL = CI.I->getDebugLoc();
-  unsigned Opcode = CI.IsX2 ? AMDGPU::BUFFER_LOAD_DWORDX4_OFFEN :
-                              AMDGPU::BUFFER_LOAD_DWORDX2_OFFEN;
+  unsigned Opcode;
+
+  if (CI.InstClass == BUFFER_LOAD_OFFEN) {
+    Opcode = CI.IsX2 ? AMDGPU::BUFFER_LOAD_DWORDX4_OFFEN :
+  } else {
+    Opcode = CI.IsX2 ? AMDGPU::BUFFER_LOAD_DWORDX4_OFFSET :
+                       AMDGPU::BUFFER_LOAD_DWORDX2_OFFSET;
+  }
 
   const TargetRegisterClass *SuperRC =
     CI.IsX2 ? &AMDGPU::VReg_128RegClass : &AMDGPU::VReg_64RegClass;
   unsigned DestReg = MRI->createVirtualRegister(SuperRC);
   unsigned MergedOffset = std::min(CI.Offset0, CI.Offset1);
 
-  BuildMI(*MBB, CI.Paired, DL, TII->get(Opcode), DestReg)
-      .add(*TII->getNamedOperand(*CI.I, AMDGPU::OpName::vaddr))
-      .add(*TII->getNamedOperand(*CI.I, AMDGPU::OpName::srsrc))
+  auto MIB = BuildMI(*MBB, CI.Paired, DL, TII->get(Opcode), DestReg);
+
+  if (CI.InstClass == BUFFER_LOAD_OFFEN)
+      MIB = MIB.add(*TII->getNamedOperand(*CI.I, AMDGPU::OpName::vaddr));
+
+  MIB.add(*TII->getNamedOperand(*CI.I, AMDGPU::OpName::srsrc))
       .add(*TII->getNamedOperand(*CI.I, AMDGPU::OpName::soffset))
       .addImm(MergedOffset) // offset
       .addImm(CI.GLC0)      // glc
       .addImm(CI.SLC0)      // slc
       .addImm(0)            // tfe
       .setMemRefs(CI.I->mergeMemRefsWith(*CI.Paired));
 
   unsigned SubRegIdx0 = CI.IsX2 ? AMDGPU::sub0_sub1 : AMDGPU::sub0;
if (STM->hasSBufferLoadStoreAtomicDwordxN() &&
         if (!CI.IsX2)
           CreatedX2++;
       } else {
         ++I;
       }
       continue;
     }
     if (Opc == AMDGPU::BUFFER_LOAD_DWORD_OFFEN ||
-        Opc == AMDGPU::BUFFER_LOAD_DWORDX2_OFFEN) {
+        Opc == AMDGPU::BUFFER_LOAD_DWORDX2_OFFEN ||
+        Opc == AMDGPU::BUFFER_LOAD_DWORD_OFFSET ||
+        Opc == AMDGPU::BUFFER_LOAD_DWORDX2_OFFSET) {
+      if (Opc == AMDGPU::BUFFER_LOAD_DWORD_OFFEN ||
+          Opc == AMDGPU::BUFFER_LOAD_DWORDX2_OFFEN)
         CI.InstClass = BUFFER_LOAD_OFFEN;
+      else
+        CI.InstClass = BUFFER_LOAD_OFFSET;
+
       CI.EltSize = 4;
-      CI.IsX2 = Opc == AMDGPU::BUFFER_LOAD_DWORDX2_OFFEN;
+      CI.IsX2 = Opc == AMDGPU::BUFFER_LOAD_DWORDX2_OFFEN ||
+                Opc == AMDGPU::BUFFER_LOAD_DWORDX2_OFFSET;
       if (findMatchingInst(CI)) {
         Modified = true;
-        I = mergeBufferLoadOffenPair(CI);
+        I = mergeBufferLoadPair(CI);
         if (!CI.IsX2)
           CreatedX2++;
       } else {
         ++I;
       }
       continue;
     }
 

test/CodeGen/AMDGPU/llvm.amdgcn.buffer.load.ll

main_body:
   %r1 = extractelement <2 x float> %vr1, i32 0
   %r2 = extractelement <2 x float> %vr1, i32 1
   %r3 = extractelement <2 x float> %vr2, i32 0
   %r4 = extractelement <2 x float> %vr2, i32 1
   call void @llvm.amdgcn.exp.f32(i32 0, i32 15, float %r1, float %r2, float %r3, float %r4, i1 true, i1 true)
   ret void
 }
 
+;CHECK-LABEL: {{^}}buffer_load_x1_offset_merged:
+;CHECK-NEXT: BB#
+;CHECK-NEXT: buffer_load_dwordx4 v[{{[0-9]}}:{{[0-9]}}], off, s[0:3], 0 offset:4
+;CHECK-NEXT: buffer_load_dwordx2 v[{{[0-9]}}:{{[0-9]}}], off, s[0:3], 0 offset:28
+;CHECK: s_waitcnt
+define amdgpu_ps void @buffer_load_x1_offset_merged(<4 x i32> inreg %rsrc) {
+main_body:
+  %r1 = call float @llvm.amdgcn.buffer.load.f32(<4 x i32> %rsrc, i32 0, i32 4, i1 0, i1 0)
+  %r2 = call float @llvm.amdgcn.buffer.load.f32(<4 x i32> %rsrc, i32 0, i32 8, i1 0, i1 0)
+  %r3 = call float @llvm.amdgcn.buffer.load.f32(<4 x i32> %rsrc, i32 0, i32 12, i1 0, i1 0)
+  %r4 = call float @llvm.amdgcn.buffer.load.f32(<4 x i32> %rsrc, i32 0, i32 16, i1 0, i1 0)
+  %r5 = call float @llvm.amdgcn.buffer.load.f32(<4 x i32> %rsrc, i32 0, i32 28, i1 0, i1 0)
+  %r6 = call float @llvm.amdgcn.buffer.load.f32(<4 x i32> %rsrc, i32 0, i32 32, i1 0, i1 0)
+  call void @llvm.amdgcn.exp.f32(i32 0, i32 15, float %r1, float %r2, float %r3, float %r4, i1 true, i1 true)
+  call void @llvm.amdgcn.exp.f32(i32 0, i32 15, float %r5, float %r6, float undef, float undef, i1 true, i1 true)
+  ret void
+}
+
+;CHECK-LABEL: {{^}}buffer_load_x2_offset_merged:
+;CHECK-NEXT: BB#
+;CHECK-NEXT: buffer_load_dwordx4 v[{{[0-9]}}:{{[0-9]}}], off, s[0:3], 0 offset:4
+;CHECK: s_waitcnt
+define amdgpu_ps void @buffer_load_x2_offset_merged(<4 x i32> inreg %rsrc) {
+main_body:
+  %vr1 = call <2 x float> @llvm.amdgcn.buffer.load.v2f32(<4 x i32> %rsrc, i32 0, i32 4, i1 0, i1 0)
+  %vr2 = call <2 x float> @llvm.amdgcn.buffer.load.v2f32(<4 x i32> %rsrc, i32 0, i32 12, i1 0, i1 0)
+  %r1 = extractelement <2 x float> %vr1, i32 0
+  %r2 = extractelement <2 x float> %vr1, i32 1
+  %r3 = extractelement <2 x float> %vr2, i32 0
+  %r4 = extractelement <2 x float> %vr2, i32 1
+  call void @llvm.amdgcn.exp.f32(i32 0, i32 15, float %r1, float %r2, float %r3, float %r4, i1 true, i1 true)
+  ret void
+}
+
 declare float @llvm.amdgcn.buffer.load.f32(<4 x i32>, i32, i32, i1, i1) #0
 declare <2 x float> @llvm.amdgcn.buffer.load.v2f32(<4 x i32>, i32, i32, i1, i1) #0
 declare <4 x float> @llvm.amdgcn.buffer.load.v4f32(<4 x i32>, i32, i32, i1, i1) #0
 declare void @llvm.amdgcn.exp.f32(i32, i32, float, float, float, float, i1, i1) #0
 
 attributes #0 = { nounwind readonly }

test/CodeGen/AMDGPU/merge-stores.ll

define amdgpu_kernel void @merge_global_store_2_adjacent_loads_i32_nonzero_base(i32 addrspace(1)* %out, i32 addrspace(1)* %in) #0 {
   %hi = load i32, i32 addrspace(1)* %in.gep.1
 
   store i32 %lo, i32 addrspace(1)* %out.gep.0
   store i32 %hi, i32 addrspace(1)* %out.gep.1
   ret void
 }
 
 ; GCN-LABEL: {{^}}merge_global_store_2_adjacent_loads_shuffle_i32:
-; GCN: buffer_load_dword v
-; GCN: buffer_load_dword v
+; GCN: buffer_load_dwordx2 v
 ; GCN: buffer_store_dword v
 ; GCN: buffer_store_dword v
 define amdgpu_kernel void @merge_global_store_2_adjacent_loads_shuffle_i32(i32 addrspace(1)* %out, i32 addrspace(1)* %in) #0 {
   %out.gep.1 = getelementptr i32, i32 addrspace(1)* %out, i32 1
   %in.gep.1 = getelementptr i32, i32 addrspace(1)* %in, i32 1
 
   %lo = load i32, i32 addrspace(1)* %in
   %hi = load i32, i32 addrspace(1)* %in.gep.1
define amdgpu_kernel void @merge_global_store_4_adjacent_loads_inverse_i32(i32 addrspace(1)* %out, i32 addrspace(1)* %in) #0 {
 
   ret void
 }
 
 ; TODO: Re-packing of loaded register required. Maybe an IR pass
 ; should catch this?
 
 ; GCN-LABEL: {{^}}merge_global_store_4_adjacent_loads_shuffle_i32:
-; GCN: buffer_load_dword v
-; GCN: buffer_load_dword v
-; GCN: buffer_load_dword v
-; GCN: buffer_load_dword v
+; GCN: buffer_load_dwordx4 v
 ; GCN: s_barrier
 ; GCN: buffer_store_dword v
 ; GCN: buffer_store_dword v
 ; GCN: buffer_store_dword v
 ; GCN: buffer_store_dword v
 define amdgpu_kernel void @merge_global_store_4_adjacent_loads_shuffle_i32(i32 addrspace(1)* %out, i32 addrspace(1)* %in) #0 {
   %out.gep.1 = getelementptr i32, i32 addrspace(1)* %out, i32 1
   %out.gep.2 = getelementptr i32, i32 addrspace(1)* %out, i32 2