AMDGPU: Add combine for short vector extract_vector_elts

[arsenm]

Summary

Try to access pieces 4 bytes at a time. This helps
various hasOneUse extract_vector_elt combines, such
as load width reductions.

Avoids test regressions in a future commit.

[rampitec]

LGTM

[arsenm]

r334836

lib/Target/AMDGPU/SIISelLowering.cpp

SDValue SITargetLowering::performCvtPkRTZCombine(SDNode *N,
   if (Src0.isUndef() && Src1.isUndef())
     return DCI.DAG.getUNDEF(N->getValueType(0));
   return SDValue();
 }
 
 SDValue SITargetLowering::performExtractVectorEltCombine(
   SDNode *N, DAGCombinerInfo &DCI) const {
   SDValue Vec = N->getOperand(0);
-
   SelectionDAG &DAG = DCI.DAG;
+
+  EVT VecVT = Vec.getValueType();
+  EVT EltVT = VecVT.getVectorElementType();
+
   if ((Vec.getOpcode() == ISD::FNEG ||
        Vec.getOpcode() == ISD::FABS) && allUsesHaveSourceMods(N)) {
     SDLoc SL(N);
     EVT EltVT = N->getValueType(0);
     SDValue Idx = N->getOperand(1);
     SDValue Elt = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL, EltVT,
                               Vec.getOperand(0), Idx);
     return DAG.getNode(Vec.getOpcode(), SL, EltVT, Elt);
if (Vec.hasOneUse() && DCI.isBeforeLegalize()) {
     case ISD::FMINNUM:
       return DAG.getNode(Opc, SL, EltVT,
                          DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL, EltVT,
                                      Vec.getOperand(0), Idx),
                          DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL, EltVT,
                                      Vec.getOperand(1), Idx));
     }
   }
+
+  if (!DCI.isBeforeLegalize())
+    return SDValue();
+
+  unsigned VecSize = VecVT.getSizeInBits();
+  unsigned EltSize = EltVT.getSizeInBits();
+
+  // Try to turn sub-dword accesses of vectors into accesses of the same 32-bit
+  // elements. This exposes more load reduction opportunities by replacing
+  // multiple small extract_vector_elements with a single 32-bit extract.
+  auto *Idx = dyn_cast<ConstantSDNode>(N->getOperand(1));
+  if (EltSize <= 16 &&
+      EltVT.isByteSized() &&
+      VecSize > 32 &&
+      VecSize % 32 == 0 &&
+      Idx) {
+    EVT NewVT = getEquivalentMemType(*DAG.getContext(), VecVT);
+
+    unsigned BitIndex = Idx->getZExtValue() * EltSize;
+    unsigned EltIdx = BitIndex / 32;
+    unsigned LeftoverBitIdx = BitIndex % 32;
+    SDLoc SL(N);
+
+    SDValue Cast = DAG.getNode(ISD::BITCAST, SL, NewVT, Vec);
+    DCI.AddToWorklist(Cast.getNode());
+
+    SDValue Elt = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL, MVT::i32, Cast,
+                              DAG.getConstant(EltIdx, SL, MVT::i32));
+    DCI.AddToWorklist(Elt.getNode());
+    SDValue Srl = DAG.getNode(ISD::SRL, SL, MVT::i32, Elt,
+                              DAG.getConstant(LeftoverBitIdx, SL, MVT::i32));
+    DCI.AddToWorklist(Srl.getNode());
+
+    SDValue Trunc = DAG.getNode(ISD::TRUNCATE, SL, EltVT.changeTypeToInteger(), Srl);
+    DCI.AddToWorklist(Trunc.getNode());
+    return DAG.getNode(ISD::BITCAST, SL, EltVT, Trunc);
+  }
+
   return SDValue();
 }
 
 static bool convertBuildVectorCastElt(SelectionDAG &DAG,
                                       SDValue &Lo, SDValue &Hi) {
   if (Hi.getOpcode() == ISD::BITCAST &&
       Hi.getOperand(0).getValueType() == MVT::f16 &&
       (isa<ConstantSDNode>(Lo) || Lo.isUndef())) {

test/CodeGen/AMDGPU/extract_vector_elt-f16.ll

define amdgpu_kernel void @v_insertelement_v4f16_dynamic_vgpr(half addrspace(1)* %out, <4 x half> addrspace(1)* %in) #0 {
   %out.gep = getelementptr inbounds half, half addrspace(1)* %out, i64 %tid.ext
   %idx.val = load volatile i32, i32 addrspace(1)* undef
   %vec = load <4 x half>, <4 x half> addrspace(1)* %in.gep
   %vec.extract = extractelement <4 x half> %vec, i32 %idx.val
   store half %vec.extract, half addrspace(1)* %out.gep
   ret void
 }
 
+; GCN-LABEL: {{^}}reduce_load_vector_v8f16_extract_01:
+; GCN: s_load_dwordx2 [[PTR:s\[[0-9]+:[0-9]+\]]],
+; GCN-NOT: {{s|buffer|flat|global}}_load_
+; GCN: s_load_dword s{{[0-9]+}}, [[PTR]], 0x0
+; GCN-NOT: {{s|buffer|flat|global}}_load_
+; GCN: s_lshr_b32 s{{[0-9]+}}, s{{[0-9]+}}, 16
+define amdgpu_kernel void @reduce_load_vector_v8f16_extract_01(<16 x half> addrspace(4)* %ptr) #0 {
+  %load = load <16 x half>, <16 x half> addrspace(4)* %ptr
+  %elt0 = extractelement <16 x half> %load, i32 0
+  %elt1 = extractelement <16 x half> %load, i32 1
+  store volatile half %elt0, half addrspace(1)* undef, align 2
+  store volatile half %elt1, half addrspace(1)* undef, align 2
+  ret void
+}
+
+; GCN-LABEL: {{^}}reduce_load_vector_v8f16_extract_23:
+; GCN: s_load_dwordx2 [[PTR:s\[[0-9]+:[0-9]+\]]],
+; GCN-NOT: {{s|buffer|flat|global}}_load_
+; GCN: s_load_dword s{{[0-9]+}}, [[PTR]], {{0x1|0x4}}
+; GCN-NOT: {{s|buffer|flat|global}}_load_
+; GCN: s_lshr_b32 s{{[0-9]+}}, s{{[0-9]+}}, 16
+define amdgpu_kernel void @reduce_load_vector_v8f16_extract_23(<16 x half> addrspace(4)* %ptr) #0 {
+  %load = load <16 x half>, <16 x half> addrspace(4)* %ptr
+  %elt2 = extractelement <16 x half> %load, i32 2
+  %elt3 = extractelement <16 x half> %load, i32 3
+  store volatile half %elt2, half addrspace(1)* undef, align 2
+  store volatile half %elt3, half addrspace(1)* undef, align 2
+  ret void
+}
+
 declare i32 @llvm.amdgcn.workitem.id.x() #1
 
 attributes #0 = { nounwind }
 attributes #1 = { nounwind readnone }

test/CodeGen/AMDGPU/extract_vector_elt-i16.ll

define amdgpu_kernel void @v_insertelement_v4i16_dynamic_sgpr(i16 addrspace(1)* %out, <4 x i16> addrspace(1)* %in, i32 %idx) #0 {
   %in.gep = getelementptr inbounds <4 x i16>, <4 x i16> addrspace(1)* %in, i64 %tid.ext
   %out.gep = getelementptr inbounds i16, i16 addrspace(1)* %out, i64 %tid.ext
   %vec = load <4 x i16>, <4 x i16> addrspace(1)* %in.gep
   %vec.extract = extractelement <4 x i16> %vec, i32 %idx
   store i16 %vec.extract, i16 addrspace(1)* %out.gep
   ret void
 }
 
+; GCN-LABEL: {{^}}reduce_load_vector_v8i16_extract_01:
+; GCN: s_load_dwordx2 [[PTR:s\[[0-9]+:[0-9]+\]]],
+; GCN-NOT: {{s|buffer|flat|global}}_load_
+; GCN: s_load_dword s{{[0-9]+}}, [[PTR]], 0x0
+; GCN-NOT: {{s|buffer|flat|global}}_load_
+; GCN: s_lshr_b32 s{{[0-9]+}}, s{{[0-9]+}}, 16
+define amdgpu_kernel void @reduce_load_vector_v8i16_extract_01(<16 x i16> addrspace(4)* %ptr) #0 {
+  %load = load <16 x i16>, <16 x i16> addrspace(4)* %ptr
+  %elt0 = extractelement <16 x i16> %load, i32 0
+  %elt1 = extractelement <16 x i16> %load, i32 1
+  store volatile i16 %elt0, i16 addrspace(1)* undef, align 2
+  store volatile i16 %elt1, i16 addrspace(1)* undef, align 2
+  ret void
+}
+
+; GCN-LABEL: {{^}}reduce_load_vector_v8i16_extract_23:
+; GCN: s_load_dwordx2 [[PTR:s\[[0-9]+:[0-9]+\]]],
+; GCN-NOT: {{s|buffer|flat|global}}_load_
+; GCN: s_load_dword s{{[0-9]+}}, [[PTR]], {{0x1|0x4}}
+; GCN-NOT: {{s|buffer|flat|global}}_load_
+; GCN: s_lshr_b32 s{{[0-9]+}}, s{{[0-9]+}}, 16
+define amdgpu_kernel void @reduce_load_vector_v8i16_extract_23(<16 x i16> addrspace(4)* %ptr) #0 {
+  %load = load <16 x i16>, <16 x i16> addrspace(4)* %ptr
+  %elt2 = extractelement <16 x i16> %load, i32 2
+  %elt3 = extractelement <16 x i16> %load, i32 3
+  store volatile i16 %elt2, i16 addrspace(1)* undef, align 2
+  store volatile i16 %elt3, i16 addrspace(1)* undef, align 2
+  ret void
+}
+
 declare i32 @llvm.amdgcn.workitem.id.x() #1
 
 attributes #0 = { nounwind }
 attributes #1 = { nounwind readnone }

test/CodeGen/AMDGPU/extract_vector_elt-i8.ll

 define amdgpu_kernel void @dynamic_extract_vector_elt_v8i8(i8 addrspace(1)* %out, <8 x i8> addrspace(4)* %vec.ptr, i32 %idx) #0 {
   %vec = load <8 x i8>, <8 x i8> addrspace(4)* %vec.ptr
   %p0 = extractelement <8 x i8> %vec, i32 %idx
   %out1 = getelementptr i8, i8 addrspace(1)* %out, i32 1
   store i8 %p0, i8 addrspace(1)* %out
   ret void
 }
 
+; GCN-LABEL: {{^}}reduce_load_vector_v8i8_extract_0123:
+; GCN-NOT: {{s|buffer|flat|global}}_load_
+; GCN: s_load_dword s
+; GCN-NOT: {{s|buffer|flat|global}}_load_
+; GCN: s_lshr_b32 s{{[0-9]+}}, s{{[0-9]+}}, 8
+; GCN: s_lshr_b32 s{{[0-9]+}}, s{{[0-9]+}}, 16
+; GCN: s_lshr_b32 s{{[0-9]+}}, s{{[0-9]+}}, 24
+define amdgpu_kernel void @reduce_load_vector_v8i8_extract_0123() #0 {
+  %load = load <8 x i8>, <8 x i8> addrspace(4)* null
+  %elt0 = extractelement <8 x i8> %load, i32 0
+  %elt1 = extractelement <8 x i8> %load, i32 1
+  %elt2 = extractelement <8 x i8> %load, i32 2
+  %elt3 = extractelement <8 x i8> %load, i32 3
+  store volatile i8 %elt0, i8 addrspace(1)* undef, align 1
+  store volatile i8 %elt1, i8 addrspace(1)* undef, align 1
+  store volatile i8 %elt2, i8 addrspace(1)* undef, align 1
+  store volatile i8 %elt3, i8 addrspace(1)* undef, align 1
+  ret void
+}
+
+; GCN-LABEL: {{^}}reduce_load_vector_v8i8_extract_0145:
+; GCN-NOT: {{s|buffer|flat|global}}_load_
+; GCN: s_load_dwordx2
+; GCN-NOT: {{s|buffer|flat|global}}_load_
+; GCN: s_lshr_b32 s{{[0-9]+}}, s{{[0-9]+}}, 8
+; GCN: s_lshr_b32 s{{[0-9]+}}, s{{[0-9]+}}, 8
+define amdgpu_kernel void @reduce_load_vector_v8i8_extract_0145() #0 {
+  %load = load <8 x i8>, <8 x i8> addrspace(4)* null
+  %elt0 = extractelement <8 x i8> %load, i32 0
+  %elt1 = extractelement <8 x i8> %load, i32 1
+  %elt4 = extractelement <8 x i8> %load, i32 4
+  %elt5 = extractelement <8 x i8> %load, i32 5
+  store volatile i8 %elt0, i8 addrspace(1)* undef, align 1
+  store volatile i8 %elt1, i8 addrspace(1)* undef, align 1
+  store volatile i8 %elt4, i8 addrspace(1)* undef, align 1
+  store volatile i8 %elt5, i8 addrspace(1)* undef, align 1
+  ret void
+}
+
+; GCN-LABEL: {{^}}reduce_load_vector_v8i8_extract_45:
+; GCN-NOT: {{s|buffer|flat|global}}_load_
+; GCN: s_mov_b64 [[PTR:s\[[0-9]+:[0-9]+\]]], 4{{$}}
+; GCN: s_load_dword s{{[0-9]+}}, [[PTR]], 0x0{{$}}
+; GCN-NOT: {{s|buffer|flat|global}}_load_
+; GCN: s_lshr_b32 s{{[0-9]+}}, s{{[0-9]+}}, 8
+define amdgpu_kernel void @reduce_load_vector_v8i8_extract_45() #0 {
+  %load = load <8 x i8>, <8 x i8> addrspace(4)* null
+  %elt4 = extractelement <8 x i8> %load, i32 4
+  %elt5 = extractelement <8 x i8> %load, i32 5
+  store volatile i8 %elt4, i8 addrspace(1)* undef, align 1
+  store volatile i8 %elt5, i8 addrspace(1)* undef, align 1
+  ret void
+}
+
+; FIXME: ought to be able to eliminate high half of load
+; GCN-LABEL: {{^}}reduce_load_vector_v16i8_extract_0145:
+; GCN-NOT: {{s|buffer|flat|global}}_load_
+; GCN: s_load_dwordx4
+; GCN-NOT: {{s|buffer|flat|global}}_load_
+; GCN: s_lshr_b32 s{{[0-9]+}}, s{{[0-9]+}}, 8
+; GCN: s_lshr_b32 s{{[0-9]+}}, s{{[0-9]+}}, 8
+define amdgpu_kernel void @reduce_load_vector_v16i8_extract_0145() #0 {
+  %load = load <16 x i8>, <16 x i8> addrspace(4)* null
+  %elt0 = extractelement <16 x i8> %load, i32 0
+  %elt1 = extractelement <16 x i8> %load, i32 1
+  %elt4 = extractelement <16 x i8> %load, i32 4
+  %elt5 = extractelement <16 x i8> %load, i32 5
+  store volatile i8 %elt0, i8 addrspace(1)* undef, align 1
+  store volatile i8 %elt1, i8 addrspace(1)* undef, align 1
+  store volatile i8 %elt4, i8 addrspace(1)* undef, align 1
+  store volatile i8 %elt5, i8 addrspace(1)* undef, align 1
+  ret void
+}
+
 attributes #0 = { nounwind }