[AMDGPU] Supported ds_read_b128 generation; Widened vector length for local address-space

[FarhanaAleen]

Summary

Starting from GCN 2nd generation, ISA supports ds_read_b128 on top of ds_read_b64. This patch supports ds_read_b128 instruction pattern and generation of this instruction.

In the vectorizer, this patch also widen the vector length so that vectorizer generates 128 bit loads for local address-space which gets translated to ds_read_b128.

[nhaehnle]

Have you tested this on real hardware? I remember reading that there is a hardware bug on gfx7 with this instruction. The bug may apply only to early gfx7 chips.

test/CodeGen/AMDGPU/reorder-stores.ll

2 ↗

(On Diff #137397)

What does SEA mean? We usually use CI for Sea Islands.

[arsenm]

I've implemented this before: https://github.com/arsenm/llvm/tree/ds-128

This looks mostly the same. It's not clear to me it's always better to use this. I don't think this executes any faster, and at least for ds_write_b128, this has an additional constraint that the inputs must now be in a contiguous 128-bit register instead of 2 independent 64-bit pairs, which increases register pressure and may require copies. It might be better to defer forming this until later, like in the LoadStoreOptimizer pass. Jeff had a benchmark he wanted to try with this.

lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

272–274

It might be OK to say 128 anyway. You could still do adjacent ds_read2_b64 even when not using ds_read_b128. I don't think we try to do the same trick we do with 4-byte aligned 8 byte reads for the 64-bit equivalent, but you might want to look into that.

Anything you change here would also equally apply to REGION_ADDRESS

lib/Target/AMDGPU/SIISelLowering.cpp

5434

This should be hidden inside a Subtarget->hasDS128() check

5435

You don't need the isAligned16 helper. You just need to check that the alignment is >= 16, not % 16

test/CodeGen/AMDGPU/ds_read2_superreg.ll

100–135 ↗

(On Diff #137397)

These tests have the unfortunate side effect of breaking what this test intended, which was the pass forming the read2. Maybe change all of these to reduce the alignment so you still get read2?

[rampitec]

Given performance benefit is somewhat unclear can you put it under an option?

lib/Target/AMDGPU/SIISelLowering.cpp

5435

Second to that.

[FarhanaAleen]

Enabled ds_read_b128 under a switch and incorporated additional comments.

[rampitec]

lib/Target/AMDGPU/SIISelLowering.cpp

5434

You only have pattern for v4i32, but enable operation for all 128 bit. Will it work with v8i16 for example?

[FarhanaAleen]

lib/Target/AMDGPU/SIISelLowering.cpp

5434

Yes, it works for i16/i8.

During dag combine, AMDGPU loadCombiner combines vector types of 8/16/64 to vector types of 32 bit type.

[rampitec]

LGTM. Thanks.

lib/Target/AMDGPU/AMDGPUInstructions.td

 //===----------------------------------------------------------------------===//
 // Load/Store Pattern Fragments
 //===----------------------------------------------------------------------===//
 
 class Aligned8Bytes <dag ops, dag frag> : PatFrag <ops, frag, [{
   return cast<MemSDNode>(N)->getAlignment() % 8 == 0;
 }]>;
 
+class Aligned16Bytes <dag ops, dag frag> : PatFrag <ops, frag, [{
+  return cast<MemSDNode>(N)->getAlignment() >= 16;
+}]>;
+
 class LoadFrag <SDPatternOperator op> : PatFrag<(ops node:$ptr), (op node:$ptr)>;
 
 class StoreFrag<SDPatternOperator op> : PatFrag <
   (ops node:$value, node:$ptr), (op node:$value, node:$ptr)
 >;
 
 class StoreHi16<SDPatternOperator op> : PatFrag <
   (ops node:$value, node:$ptr), (op (srl node:$value, (i32 16)), node:$ptr)

 def truncstorei16_local : LocalStore <truncstorei16>;
 def store_local_hi16 : StoreHi16 <truncstorei16>, LocalAddress;
 def truncstorei8_local_hi16 : StoreHi16<truncstorei8>, LocalAddress;
 
 def load_align8_local : Aligned8Bytes <
   (ops node:$ptr), (load_local node:$ptr)
 >;
 
+def load_align16_local : Aligned16Bytes <
+  (ops node:$ptr), (load_local node:$ptr)
+>;
+
 def store_align8_local : Aligned8Bytes <
   (ops node:$val, node:$ptr), (store_local node:$val, node:$ptr)
 >;
 
 
 def load_flat          : FlatLoad <load>;
 def az_extloadi8_flat  : FlatLoad <az_extloadi8>;
 def sextloadi8_flat    : FlatLoad <sextloadi8>;

lib/Target/AMDGPU/AMDGPUSubtarget.h

public:
   bool enableIEEEBit(const MachineFunction &MF) const {
     return AMDGPU::isCompute(MF.getFunction().getCallingConv());
   }
 
   bool useFlatForGlobal() const {
     return FlatForGlobal;
   }
 
+  /// \returns If target supports ds_read/write_b128 and user enables generation
+  /// of ds_read/write_b128.
+  bool useDS128(bool UserEnable) const {
+    return CIInsts && UserEnable;
+  }
+
   /// \returns If MUBUF instructions always perform range checking, even for
   /// buffer resources used for private memory access.
   bool privateMemoryResourceIsRangeChecked() const {
     return getGeneration() < AMDGPUSubtarget::GFX9;
   }
 
   bool hasAutoWaitcntBeforeBarrier() const {
     return AutoWaitcntBeforeBarrier;

lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

unsigned AMDGPUTTIImpl::getLoadStoreVecRegBitWidth(unsigned AddrSpace) const {
   if (AddrSpace == AS.GLOBAL_ADDRESS ||
       AddrSpace == AS.CONSTANT_ADDRESS ||
       AddrSpace == AS.CONSTANT_ADDRESS_32BIT) {
     if (ST->getGeneration() <= AMDGPUSubtarget::NORTHERN_ISLANDS)
       return 128;
     return 512;
   }
 
-  if (AddrSpace == AS.FLAT_ADDRESS)
-    return 128;
-  if (AddrSpace == AS.LOCAL_ADDRESS ||
+  if (AddrSpace == AS.FLAT_ADDRESS ||
+      AddrSpace == AS.LOCAL_ADDRESS ||
       AddrSpace == AS.REGION_ADDRESS)
-    return 64;
+    return 128;
+
   if (AddrSpace == AS.PRIVATE_ADDRESS)
     return 8 * ST->getMaxPrivateElementSize();
 
   if (ST->getGeneration() <= AMDGPUSubtarget::NORTHERN_ISLANDS &&
       (AddrSpace == AS.PARAM_D_ADDRESS ||
       AddrSpace == AS.PARAM_I_ADDRESS ||
        (AddrSpace >= AS.CONSTANT_BUFFER_0 &&
         AddrSpace <= AS.CONSTANT_BUFFER_15)))
     return 128;
   llvm_unreachable("unhandled address space");

lib/Target/AMDGPU/DSInstructions.td

 defm : DSReadPat_mc <DS_READ_I16, i32, "sextloadi16_local">;
 defm : DSReadPat_mc <DS_READ_U16, i32, "az_extloadi16_local">;
 defm : DSReadPat_mc <DS_READ_U16, i16, "load_local">;
 defm : DSReadPat_mc <DS_READ_B32, i32, "load_local">;
 
 let AddedComplexity = 100 in {
 
 defm : DSReadPat_mc <DS_READ_B64, v2i32, "load_align8_local">;
+defm : DSReadPat_mc <DS_READ_B128, v4i32, "load_align16_local">;
 
 } // End AddedComplexity = 100
 
 let OtherPredicates = [HasD16LoadStore] in {
 let AddedComplexity = 100 in {
 defm : DSReadPat_Hi16<DS_READ_U16_D16_HI, load_local>;
 defm : DSReadPat_Hi16<DS_READ_U8_D16_HI, az_extloadi8_local>;
 defm : DSReadPat_Hi16<DS_READ_I8_D16_HI, sextloadi8_local>;

lib/Target/AMDGPU/SIISelLowering.cpp

 
 STATISTIC(NumTailCalls, "Number of tail calls");
 
 static cl::opt<bool> EnableVGPRIndexMode(
   "amdgpu-vgpr-index-mode",
   cl::desc("Use GPR indexing mode instead of movrel for vector indexing"),
   cl::init(false));
 
+static cl::opt<bool> EnableDS128(
+  "amdgpu-ds128",
+  cl::desc("Use DS_read/write_b128"),
+  cl::init(false));
+
 static cl::opt<unsigned> AssumeFrameIndexHighZeroBits(
   "amdgpu-frame-index-zero-bits",
   cl::desc("High bits of frame index assumed to be zero"),
   cl::init(5),
   cl::ReallyHidden);
 
 static unsigned findFirstFreeSGPR(CCState &CCInfo) {
   unsigned NumSGPRs = AMDGPU::SGPR_32RegClass.getNumRegs();
case 16:
       // Same as global/flat
       if (NumElements > 4)
         return SplitVectorLoad(Op, DAG);
       return SDValue();
     default:
       llvm_unreachable("unsupported private_element_size");
     }
   } else if (AS == AMDGPUASI.LOCAL_ADDRESS) {
-    if (NumElements > 2)
-      return SplitVectorLoad(Op, DAG);
-
-    if (NumElements == 2)
+    // Use ds_read_b128 if possible.
+    if (Subtarget->useDS128(EnableDS128) && Load->getAlignment() >= 16 &&
+        MemVT.getStoreSize() == 16)
       return SDValue();
 
-    // If properly aligned, if we split we might be able to use ds_read_b64.
+    if (NumElements > 2)
       return SplitVectorLoad(Op, DAG);
   }
   return SDValue();
 }
 
 SDValue SITargetLowering::LowerSELECT(SDValue Op, SelectionDAG &DAG) const {
   if (Op.getValueType() != MVT::i64)
     return SDValue();
 

lib/Target/AMDGPU/SIInstrInfo.td

 
 def sextloadi16_glue : PatFrag<(ops node:$ptr), (sextload_glue node:$ptr), [{
   return cast<LoadSDNode>(N)->getMemoryVT() == MVT::i16;
 }]>;
 
 def load_glue_align8 : Aligned8Bytes <
   (ops node:$ptr), (load_glue node:$ptr)
 >;
+def load_glue_align16 : Aligned16Bytes <
+  (ops node:$ptr), (load_glue node:$ptr)
+>;
 
 
 def load_local_m0 : LoadFrag<load_glue>, LocalAddress;
 def sextloadi8_local_m0 : LoadFrag<sextloadi8_glue>, LocalAddress;
 def sextloadi16_local_m0 : LoadFrag<sextloadi16_glue>, LocalAddress;
 def az_extloadi8_local_m0 : LoadFrag<az_extloadi8_glue>, LocalAddress;
 def az_extloadi16_local_m0 : LoadFrag<az_extloadi16_glue>, LocalAddress;
 def load_align8_local_m0 : LoadFrag <load_glue_align8>, LocalAddress;
+def load_align16_local_m0 : LoadFrag <load_glue_align16>, LocalAddress;
 
 
 def AMDGPUst_glue : SDNode <"ISD::STORE", SDTStore,
   [SDNPHasChain, SDNPMayStore, SDNPMemOperand, SDNPInGlue]
 >;
 
 def unindexedstore_glue : PatFrag<(ops node:$val, node:$ptr),
                                    (AMDGPUst_glue node:$val, node:$ptr), [{

test/CodeGen/AMDGPU/load-local-i32.ll

 ; RUN: llc -march=amdgcn -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,SICIVI,FUNC %s
 ; RUN: llc -march=amdgcn -mcpu=tonga -mattr=-flat-for-global -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,SICIVI,FUNC %s
 ; RUN: llc -march=amdgcn -mcpu=gfx900 -mattr=-flat-for-global -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,VI,FUNC %s
 ; RUN: llc -march=r600 -mcpu=redwood < %s | FileCheck -check-prefix=EG -check-prefix=FUNC %s
 
+; RUN: llc -march=amdgcn -mcpu=tahiti -amdgpu-ds128 < %s | FileCheck -check-prefixes=SI,FUNC %s
+; RUN: llc -march=amdgcn -mcpu=tonga -amdgpu-ds128 < %s | FileCheck -check-prefixes=CIVI,FUNC %s
+; RUN: llc -march=amdgcn -mcpu=gfx900 -amdgpu-ds128 < %s | FileCheck -check-prefixes=CIVI,FUNC %s
+
 ; FUNC-LABEL: {{^}}local_load_i32:
 ; GCN-NOT: s_wqm_b64
 ; SICIVI: s_mov_b32 m0, -1
 ; GFX9-NOT: m0
 ; GCN: ds_read_b32
 
 ; EG: LDS_READ_RET
 define amdgpu_kernel void @local_load_i32(i32 addrspace(3)* %out, i32 addrspace(3)* %in) #0 {

 
 define amdgpu_kernel void @local_sextload_v4i32_to_v4i64(<4 x i64> addrspace(3)* %out, <4 x i32> addrspace(3)* %in) #0 {
   %ld = load <4 x i32>, <4 x i32> addrspace(3)* %in
   %ext = sext <4 x i32> %ld to <4 x i64>
   store <4 x i64> %ext, <4 x i64> addrspace(3)* %out
   ret void
 }
 
+; FUNC-LABEL: {{^}}local_v4i32_to_128:
+; SI-NOT: ds_read_b128
+; CIVI: ds_read_b128
+; EG: LDS_READ_RET
+; EG: LDS_READ_RET
+; EG: LDS_READ_RET
+; EG: LDS_READ_RET
+define amdgpu_kernel void @local_v4i32_to_128(<4 x i32> addrspace(3)* %out, <4 x i32> addrspace(3)* %in) {
+  %ld = load <4 x i32>, <4 x i32> addrspace(3)* %in, align 16
+  store <4 x i32> %ld, <4 x i32> addrspace(3)* %out
+  ret void
+}
+
 ; FUNC-LABEL: {{^}}local_zextload_v8i32_to_v8i64:
 ; SICIVI: s_mov_b32 m0, -1
 ; GFX9-NOT: m0
 
 define amdgpu_kernel void @local_zextload_v8i32_to_v8i64(<8 x i64> addrspace(3)* %out, <8 x i32> addrspace(3)* %in) #0 {
   %ld = load <8 x i32>, <8 x i32> addrspace(3)* %in
   %ext = zext <8 x i32> %ld to <8 x i64>
   store <8 x i64> %ext, <8 x i64> addrspace(3)* %out

test/CodeGen/AMDGPU/load-local-i64.ll

 ; RUN: llc -march=amdgcn -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,SICIVI,FUNC %s
 ; RUN: llc -mtriple=amdgcn--amdhsa -mcpu=kaveri -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,SICIVI,FUNC %s
 ; RUN: llc -march=amdgcn -mcpu=tonga -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,SICIVI,FUNC %s
 ; RUN: llc -march=amdgcn -mcpu=gfx900 -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,GFX9,FUNC %s
 ; RUN: llc -march=r600 -mcpu=redwood < %s | FileCheck -check-prefixes=EG,FUNC %s
 
+; RUN: llc -march=amdgcn -mcpu=tonga -verify-machineinstrs -amdgpu-ds128 < %s | FileCheck -check-prefixes=CIVI,FUNC %s
+; RUN: llc -march=amdgcn -mcpu=gfx900 -verify-machineinstrs -amdgpu-ds128 < %s | FileCheck -check-prefixes=CIVI,FUNC %s
+
 ; FUNC-LABEL: {{^}}local_load_i64:
 ; SICIVI: s_mov_b32 m0
 ; GFX9-NOT: m0
 
 ; GCN: ds_read_b64 [[VAL:v\[[0-9]+:[0-9]+\]]], v{{[0-9]+}}{{$}}
 ; GCN: ds_write_b64 v{{[0-9]+}}, [[VAL]]
 
 ; EG: LDS_READ_RET

 ; EG: LDS_READ_RET
 define amdgpu_kernel void @local_load_v2i64(<2 x i64> addrspace(3)* %out, <2 x i64> addrspace(3)* %in) #0 {
 entry:
   %ld = load <2 x i64>, <2 x i64> addrspace(3)* %in
   store <2 x i64> %ld, <2 x i64> addrspace(3)* %out
   ret void
 }
 
+; FUNC-LABEL: {{^}}local_load_v2i64_to_128:
+; CIVI: ds_read_b128
+define amdgpu_kernel void @local_load_v2i64_to_128(<2 x i64> addrspace(3)* %out, <2 x i64> addrspace(3)* %in) #0 {
+entry:
+  %ld = load <2 x i64>, <2 x i64> addrspace(3)* %in
+  store <2 x i64> %ld, <2 x i64> addrspace(3)* %out
+  ret void
+}
+
 ; FUNC-LABEL: {{^}}local_load_v3i64:
 ; SICIVI: s_mov_b32 m0
 ; GFX9-NOT: m0
 
 ; GCN-DAG: ds_read2_b64
 ; GCN-DAG: ds_read_b64
 
 ; EG: LDS_READ_RET

test/Transforms/LoadStoreVectorizer/AMDGPU/merge-stores.ll

define amdgpu_kernel void @merge_local_store_2_constants_i32_align_2(i32 addrspace(3)* %out) #0 {
   %out.gep.1 = getelementptr i32, i32 addrspace(3)* %out, i32 1
 
   store i32 123, i32 addrspace(3)* %out.gep.1, align 2
   store i32 456, i32 addrspace(3)* %out, align 2
   ret void
 }
 
 ; CHECK-LABEL: @merge_local_store_4_constants_i32
-; CHECK: store <2 x i32> <i32 456, i32 333>, <2 x i32> addrspace(3)*
-; CHECK: store <2 x i32> <i32 1234, i32 123>, <2 x i32> addrspace(3)*
+; CHECK: store <4 x i32> <i32 1234, i32 123, i32 456, i32 333>, <4 x i32> addrspace(3)*
 define amdgpu_kernel void @merge_local_store_4_constants_i32(i32 addrspace(3)* %out) #0 {
   %out.gep.1 = getelementptr i32, i32 addrspace(3)* %out, i32 1
   %out.gep.2 = getelementptr i32, i32 addrspace(3)* %out, i32 2
   %out.gep.3 = getelementptr i32, i32 addrspace(3)* %out, i32 3
 
   store i32 123, i32 addrspace(3)* %out.gep.1
   store i32 456, i32 addrspace(3)* %out.gep.2
   store i32 333, i32 addrspace(3)* %out.gep.3

test/Transforms/LoadStoreVectorizer/AMDGPU/multiple_tails.ll

define amdgpu_kernel void @no_crash(i32 %arg) {
 
   ret void
 }
 
 ; Check adjiacent memory locations are properly matched and the
 ; longest chain vectorized
 
 ; CHECK-LABEL: @interleave_get_longest
-; CHECK: load <2 x i32>
+; CHECK: load <4 x i32>
 ; CHECK: load i32
 ; CHECK: store <2 x i32> zeroinitializer
 ; CHECK: load i32
-; CHECK: load <2 x i32>
 ; CHECK: load i32
 ; CHECK: load i32
 
 define amdgpu_kernel void @interleave_get_longest(i32 %arg) {
   %a1 = add i32 %arg, 1
   %a2 = add i32 %arg, 2
   %a3 = add i32 %arg, 3
   %a4 = add i32 %arg, 4