AMDGPU: Merge BUFFER_STORE_DWORD_OFFEN/OFFSET into x2, x4

[mareko]

Summary

Only 56 shaders (out of 48486) are affected.

Totals from affected shaders (changed stats only):
SGPRS: 2420 -> 2460 (1.65 %)
Spilled VGPRs: 94 -> 112 (19.15 %)
Scratch size: 524 -> 528 (0.76 %) dwords per thread
Code Size: 187400 -> 184992 (-1.28 %) bytes

One DiRT Showdown shader spills 6 more VGPRs.
One Grid Autosport shader spills 12 more VGPRs.

The other 54 shaders only have a decrease in code size.
(I'm ignoring the SGPR noise)

[nhaehnle]

One minor comment inline.

The volatile / hasOrderedMemoryRef issue is worrying, and I don't think the change should go in as-is. Besides, it's probably not as big a win as the other patches, especially due to the added VGPR spills. Do you have all these patches in a branch somewhere?

lib/Target/AMDGPU/SILoadStoreOptimizer.cpp

768 ↗

(On Diff #119350)

Remove the assignment "MIB = MIB" here and below, it shouldn't be necessary. The return value is just a convenience to allow chains of .add(..).add(..) etc.

test/CodeGen/AMDGPU/fmul-2-combine-multi-use.ll

42–51 ↗

(On Diff #119350)

In this case, the stores are explicitly marked as volatile in the input IR. Merging those seems like a very bad idea.

56–66 ↗

(On Diff #119350)

Same here.

[mareko]

In D39012#903643, @nhaehnle wrote:

One minor comment inline.

The volatile / hasOrderedMemoryRef issue is worrying, and I don't think the change should go in as-is. Besides, it's probably not as big a win as the other patches, especially due to the added VGPR spills. Do you have all these patches in a branch somewhere?

The spills don't bother me because the shaders already spill VGPRs and any random codegen change tends to make it better or worse randomly. I will push my patches to some branch later.

[mareko]

• don't ignore MOVolatile

llvm/trunk/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,
+    BUFFER_STORE_OFFEN,
+    BUFFER_STORE_OFFSET,
   };
 
   struct CombineInfo {
     MachineBasicBlock::iterator I;
     MachineBasicBlock::iterator Paired;
     unsigned EltSize;
     unsigned Offset0;
     unsigned Offset1;
private:
 
   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 mergeBufferLoadPair(CombineInfo &CI);
+  unsigned promoteBufferStoreOpcode(const MachineInstr &I, bool &IsX2,
+                                    bool &IsOffen) const;
+  MachineBasicBlock::iterator mergeBufferStorePair(CombineInfo &CI);
 
 public:
   static char ID;
 
   SILoadStoreOptimizer() : MachineFunctionPass(ID) {
     initializeSILoadStoreOptimizerPass(*PassRegistry::getPassRegistry());
   }
 
bool SILoadStoreOptimizer::offsetsCanBeCombined(CombineInfo &CI) {
   if ((CI.Offset0 % CI.EltSize != 0) || (CI.Offset1 % CI.EltSize != 0))
     return false;
 
   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_OFFSET) {
+  // Handle SMEM and VMEM instructions.
+  if (CI.InstClass != DS_READ_WRITE) {
     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) {
   switch (CI.InstClass) {
   case DS_READ_WRITE:
     AddrOpName[NumAddresses++] = AMDGPU::OpName::addr;
     break;
   case S_BUFFER_LOAD_IMM:
     AddrOpName[NumAddresses++] = AMDGPU::OpName::sbase;
     break;
   case BUFFER_LOAD_OFFEN:
+  case BUFFER_STORE_OFFEN:
     AddrOpName[NumAddresses++] = AMDGPU::OpName::srsrc;
     AddrOpName[NumAddresses++] = AMDGPU::OpName::vaddr;
     AddrOpName[NumAddresses++] = AMDGPU::OpName::soffset;
     break;
   case BUFFER_LOAD_OFFSET:
+  case BUFFER_STORE_OFFSET:
     AddrOpName[NumAddresses++] = AMDGPU::OpName::srsrc;
     AddrOpName[NumAddresses++] = AMDGPU::OpName::soffset;
     break;
   default:
     llvm_unreachable("invalid InstClass");
   }
 
   for (unsigned i = 0; i < NumAddresses; i++) {
MachineBasicBlock::iterator SILoadStoreOptimizer::mergeBufferLoadPair(
   moveInstsAfter(Copy1, CI.InstsToMove);
 
   MachineBasicBlock::iterator Next = std::next(CI.I);
   CI.I->eraseFromParent();
   CI.Paired->eraseFromParent();
   return Next;
 }
 
+unsigned SILoadStoreOptimizer::promoteBufferStoreOpcode(
+  const MachineInstr &I, bool &IsX2, bool &IsOffen) const {
+  IsX2 = false;
+  IsOffen = false;
+
+  switch (I.getOpcode()) {
+  case AMDGPU::BUFFER_STORE_DWORD_OFFEN:
+    IsOffen = true;
+    return AMDGPU::BUFFER_STORE_DWORDX2_OFFEN;
+  case AMDGPU::BUFFER_STORE_DWORD_OFFEN_exact:
+    IsOffen = true;
+    return AMDGPU::BUFFER_STORE_DWORDX2_OFFEN_exact;
+  case AMDGPU::BUFFER_STORE_DWORDX2_OFFEN:
+    IsX2 = true;
+    IsOffen = true;
+    return AMDGPU::BUFFER_STORE_DWORDX4_OFFEN;
+  case AMDGPU::BUFFER_STORE_DWORDX2_OFFEN_exact:
+    IsX2 = true;
+    IsOffen = true;
+    return AMDGPU::BUFFER_STORE_DWORDX4_OFFEN_exact;
+  case AMDGPU::BUFFER_STORE_DWORD_OFFSET:
+    return AMDGPU::BUFFER_STORE_DWORDX2_OFFSET;
+  case AMDGPU::BUFFER_STORE_DWORD_OFFSET_exact:
+    return AMDGPU::BUFFER_STORE_DWORDX2_OFFSET_exact;
+  case AMDGPU::BUFFER_STORE_DWORDX2_OFFSET:
+    IsX2 = true;
+    return AMDGPU::BUFFER_STORE_DWORDX4_OFFSET;
+  case AMDGPU::BUFFER_STORE_DWORDX2_OFFSET_exact:
+    IsX2 = true;
+    return AMDGPU::BUFFER_STORE_DWORDX4_OFFSET_exact;
+  }
+  return 0;
+}
+
+MachineBasicBlock::iterator SILoadStoreOptimizer::mergeBufferStorePair(
+  CombineInfo &CI) {
+  MachineBasicBlock *MBB = CI.I->getParent();
+  DebugLoc DL = CI.I->getDebugLoc();
+  bool Unused1, Unused2;
+  unsigned Opcode = promoteBufferStoreOpcode(*CI.I, Unused1, Unused2);
+
+  unsigned SubRegIdx0 = CI.IsX2 ? AMDGPU::sub0_sub1 : AMDGPU::sub0;
+  unsigned SubRegIdx1 = CI.IsX2 ? AMDGPU::sub2_sub3 : AMDGPU::sub1;
+
+  // Handle descending offsets
+  if (CI.Offset0 > CI.Offset1)
+    std::swap(SubRegIdx0, SubRegIdx1);
+
+  // Copy to the new source register.
+  const TargetRegisterClass *SuperRC =
+    CI.IsX2 ? &AMDGPU::VReg_128RegClass : &AMDGPU::VReg_64RegClass;
+  unsigned SrcReg = MRI->createVirtualRegister(SuperRC);
+
+  const auto *Src0 = TII->getNamedOperand(*CI.I, AMDGPU::OpName::vdata);
+  const auto *Src1 = TII->getNamedOperand(*CI.Paired, AMDGPU::OpName::vdata);
+
+  BuildMI(*MBB, CI.Paired, DL, TII->get(AMDGPU::REG_SEQUENCE), SrcReg)
+      .add(*Src0)
+      .addImm(SubRegIdx0)
+      .add(*Src1)
+      .addImm(SubRegIdx1);
+
+  auto MIB = BuildMI(*MBB, CI.Paired, DL, TII->get(Opcode))
+      .addReg(SrcReg, RegState::Kill);
+
+  if (CI.InstClass == BUFFER_STORE_OFFEN)
+      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(std::min(CI.Offset0, CI.Offset1)) // offset
+      .addImm(CI.GLC0)      // glc
+      .addImm(CI.SLC0)      // slc
+      .addImm(0)            // tfe
+      .setMemRefs(CI.I->mergeMemRefsWith(*CI.Paired));
+
+  moveInstsAfter(MIB, CI.InstsToMove);
+
+  MachineBasicBlock::iterator Next = std::next(CI.I);
+  CI.I->eraseFromParent();
+  CI.Paired->eraseFromParent();
+  return Next;
+}
+
 // Scan through looking for adjacent LDS operations with constant offsets from
 // the same base register. We rely on the scheduler to do the hard work of
 // clustering nearby loads, and assume these are all adjacent.
 bool SILoadStoreOptimizer::optimizeBlock(MachineBasicBlock &MBB) {
   bool Modified = false;
 
   for (MachineBasicBlock::iterator I = MBB.begin(), E = MBB.end(); I != E;) {
     MachineInstr &MI = *I;
if (Opc == AMDGPU::BUFFER_LOAD_DWORD_OFFEN ||
         if (!CI.IsX2)
           CreatedX2++;
       } else {
         ++I;
       }
       continue;
     }
 
+    bool StoreIsX2, IsOffen;
+    if (promoteBufferStoreOpcode(*I, StoreIsX2, IsOffen)) {
+      CI.InstClass = IsOffen ? BUFFER_STORE_OFFEN : BUFFER_STORE_OFFSET;
+      CI.EltSize = 4;
+      CI.IsX2 = StoreIsX2;
+      if (findMatchingInst(CI)) {
+        Modified = true;
+        I = mergeBufferStorePair(CI);
+        if (!CI.IsX2)
+          CreatedX2++;
+      } else {
+        ++I;
+      }
+      continue;
+    }
+
     ++I;
   }
 
   return Modified;
 }
 
 bool SILoadStoreOptimizer::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(*MF.getFunction()))

llvm/trunk/test/CodeGen/AMDGPU/llvm.amdgcn.buffer.store.ll

 ;CHECK-NOT: s_waitcnt
 ;CHECK: buffer_store_dwordx2 v[0:1], v2, s[0:3], 0 idxen
 define amdgpu_ps void @buffer_store_x2(<4 x i32> inreg %rsrc, <2 x float> %data, i32 %index) #0 {
 main_body:
   call void @llvm.amdgcn.buffer.store.v2f32(<2 x float> %data, <4 x i32> %rsrc, i32 %index, i32 0, i1 0, i1 0)
   ret void
 }
 
+;CHECK-LABEL: {{^}}buffer_store_x1_offen_merged:
+;CHECK-NOT: s_waitcnt
+;CHECK-DAG: buffer_store_dwordx4 v[{{[0-9]}}:{{[0-9]}}], v0, s[0:3], 0 offen offset:4
+;CHECK-DAG: buffer_store_dwordx2 v[{{[0-9]}}:{{[0-9]}}], v0, s[0:3], 0 offen offset:28
+define amdgpu_ps void @buffer_store_x1_offen_merged(<4 x i32> inreg %rsrc, i32 %a, float %v1, float %v2, float %v3, float %v4, float %v5, float %v6) {
+  %a1 = add i32 %a, 4
+  %a2 = add i32 %a, 8
+  %a3 = add i32 %a, 12
+  %a4 = add i32 %a, 16
+  %a5 = add i32 %a, 28
+  %a6 = add i32 %a, 32
+  call void @llvm.amdgcn.buffer.store.f32(float %v1, <4 x i32> %rsrc, i32 0, i32 %a1, i1 0, i1 0)
+  call void @llvm.amdgcn.buffer.store.f32(float %v2, <4 x i32> %rsrc, i32 0, i32 %a2, i1 0, i1 0)
+  call void @llvm.amdgcn.buffer.store.f32(float %v3, <4 x i32> %rsrc, i32 0, i32 %a3, i1 0, i1 0)
+  call void @llvm.amdgcn.buffer.store.f32(float %v4, <4 x i32> %rsrc, i32 0, i32 %a4, i1 0, i1 0)
+  call void @llvm.amdgcn.buffer.store.f32(float %v5, <4 x i32> %rsrc, i32 0, i32 %a5, i1 0, i1 0)
+  call void @llvm.amdgcn.buffer.store.f32(float %v6, <4 x i32> %rsrc, i32 0, i32 %a6, i1 0, i1 0)
+  ret void
+}
+
+;CHECK-LABEL: {{^}}buffer_store_x1_offen_merged_glc_slc:
+;CHECK-NOT: s_waitcnt
+;CHECK-DAG: buffer_store_dwordx2 v[{{[0-9]}}:{{[0-9]}}], v0, s[0:3], 0 offen offset:4{{$}}
+;CHECK-DAG: buffer_store_dwordx2 v[{{[0-9]}}:{{[0-9]}}], v0, s[0:3], 0 offen offset:12 glc{{$}}
+;CHECK-DAG: buffer_store_dwordx2 v[{{[0-9]}}:{{[0-9]}}], v0, s[0:3], 0 offen offset:28 glc slc{{$}}
+define amdgpu_ps void @buffer_store_x1_offen_merged_glc_slc(<4 x i32> inreg %rsrc, i32 %a, float %v1, float %v2, float %v3, float %v4, float %v5, float %v6) {
+  %a1 = add i32 %a, 4
+  %a2 = add i32 %a, 8
+  %a3 = add i32 %a, 12
+  %a4 = add i32 %a, 16
+  %a5 = add i32 %a, 28
+  %a6 = add i32 %a, 32
+  call void @llvm.amdgcn.buffer.store.f32(float %v1, <4 x i32> %rsrc, i32 0, i32 %a1, i1 0, i1 0)
+  call void @llvm.amdgcn.buffer.store.f32(float %v2, <4 x i32> %rsrc, i32 0, i32 %a2, i1 0, i1 0)
+  call void @llvm.amdgcn.buffer.store.f32(float %v3, <4 x i32> %rsrc, i32 0, i32 %a3, i1 1, i1 0)
+  call void @llvm.amdgcn.buffer.store.f32(float %v4, <4 x i32> %rsrc, i32 0, i32 %a4, i1 1, i1 0)
+  call void @llvm.amdgcn.buffer.store.f32(float %v5, <4 x i32> %rsrc, i32 0, i32 %a5, i1 1, i1 1)
+  call void @llvm.amdgcn.buffer.store.f32(float %v6, <4 x i32> %rsrc, i32 0, i32 %a6, i1 1, i1 1)
+  ret void
+}
+
+;CHECK-LABEL: {{^}}buffer_store_x2_offen_merged:
+;CHECK-NOT: s_waitcnt
+;CHECK: buffer_store_dwordx4 v[{{[0-9]}}:{{[0-9]}}], v0, s[0:3], 0 offen offset:4
+define amdgpu_ps void @buffer_store_x2_offen_merged(<4 x i32> inreg %rsrc, i32 %a, <2 x float> %v1, <2 x float> %v2) {
+  %a1 = add i32 %a, 4
+  %a2 = add i32 %a, 12
+  call void @llvm.amdgcn.buffer.store.v2f32(<2 x float> %v1, <4 x i32> %rsrc, i32 0, i32 %a1, i1 0, i1 0)
+  call void @llvm.amdgcn.buffer.store.v2f32(<2 x float> %v2, <4 x i32> %rsrc, i32 0, i32 %a2, i1 0, i1 0)
+  ret void
+}
+
+;CHECK-LABEL: {{^}}buffer_store_x1_offset_merged:
+;CHECK-NOT: s_waitcnt
+;CHECK-DAG: buffer_store_dwordx4 v[{{[0-9]}}:{{[0-9]}}], off, s[0:3], 0 offset:4
+;CHECK-DAG: buffer_store_dwordx2 v[{{[0-9]}}:{{[0-9]}}], off, s[0:3], 0 offset:28
+define amdgpu_ps void @buffer_store_x1_offset_merged(<4 x i32> inreg %rsrc, float %v1, float %v2, float %v3, float %v4, float %v5, float %v6) {
+  call void @llvm.amdgcn.buffer.store.f32(float %v1, <4 x i32> %rsrc, i32 0, i32 4, i1 0, i1 0)
+  call void @llvm.amdgcn.buffer.store.f32(float %v2, <4 x i32> %rsrc, i32 0, i32 8, i1 0, i1 0)
+  call void @llvm.amdgcn.buffer.store.f32(float %v3, <4 x i32> %rsrc, i32 0, i32 12, i1 0, i1 0)
+  call void @llvm.amdgcn.buffer.store.f32(float %v4, <4 x i32> %rsrc, i32 0, i32 16, i1 0, i1 0)
+  call void @llvm.amdgcn.buffer.store.f32(float %v5, <4 x i32> %rsrc, i32 0, i32 28, i1 0, i1 0)
+  call void @llvm.amdgcn.buffer.store.f32(float %v6, <4 x i32> %rsrc, i32 0, i32 32, i1 0, i1 0)
+  ret void
+}
+
+;CHECK-LABEL: {{^}}buffer_store_x2_offset_merged:
+;CHECK-NOT: s_waitcnt
+;CHECK: buffer_store_dwordx4 v[{{[0-9]}}:{{[0-9]}}], off, s[0:3], 0 offset:4
+define amdgpu_ps void @buffer_store_x2_offset_merged(<4 x i32> inreg %rsrc, <2 x float> %v1,<2 x float> %v2) {
+  call void @llvm.amdgcn.buffer.store.v2f32(<2 x float> %v1, <4 x i32> %rsrc, i32 0, i32 4, i1 0, i1 0)
+  call void @llvm.amdgcn.buffer.store.v2f32(<2 x float> %v2, <4 x i32> %rsrc, i32 0, i32 12, i1 0, i1 0)
+  ret void
+}
+
 declare void @llvm.amdgcn.buffer.store.f32(float, <4 x i32>, i32, i32, i1, i1) #0
 declare void @llvm.amdgcn.buffer.store.v2f32(<2 x float>, <4 x i32>, i32, i32, i1, i1) #0
 declare void @llvm.amdgcn.buffer.store.v4f32(<4 x float>, <4 x i32>, i32, i32, i1, i1) #0
 declare <4 x float> @llvm.amdgcn.buffer.load.v4f32(<4 x i32>, i32, i32, i1, i1) #1
 
 attributes #0 = { nounwind }
 attributes #1 = { nounwind readonly }

llvm/trunk/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 {
 
   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_dwordx2 v
-; GCN: buffer_store_dword v
-; GCN: buffer_store_dword v
+; GCN: buffer_store_dwordx2 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
 
   store i32 %hi, i32 addrspace(1)* %out

 }
 
 ; 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_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
+; GCN: buffer_store_dwordx4 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
   %out.gep.3 = getelementptr i32, i32 addrspace(1)* %out, i32 3
   %in.gep.1 = getelementptr i32, i32 addrspace(1)* %in, i32 1
   %in.gep.2 = getelementptr i32, i32 addrspace(1)* %in, i32 2
   %in.gep.3 = getelementptr i32, i32 addrspace(1)* %in, i32 3