AMDGPU/SDAG: Improve {extract,insert}_subvector lowering for 16-bit vectors

[foad]

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

1440

Start != 1 is redundant now.

[nhaehnle]

How about this? :)

[nhaehnle]

More test changes

[Joe_Nash]

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

1440

I'm not sure the Start == NumElt case will ever be hit.
But it looks like that can be removed in a follow up patch anyway.
LGTM

[nhaehnle]

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

1440

It's an extract, so Start == NumElt should mean that we're extracting the high half of the source vector. Unless I'm extremely confused right now.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

1440

I think this is one of those cases that theoretically exists but is really hard to observe since the combiner would always just fold it to the source operand

[Joe_Nash]

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

1440

Oh, you are right. I was doing NumElt->NumSrcElt in my head

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

SDValue AMDGPUTargetLowering::LowerCONCAT_VECTORS(SDValue Op,
   for (const SDUse &U : Op->ops())
     DAG.ExtractVectorElements(U.get(), Args);
 
   return DAG.getBuildVector(Op.getValueType(), SL, Args);
 }
 
 SDValue AMDGPUTargetLowering::LowerEXTRACT_SUBVECTOR(SDValue Op,
                                                      SelectionDAG &DAG) const {
-
+  SDLoc SL(Op);
   SmallVector<SDValue, 8> Args;
   unsigned Start = cast<ConstantSDNode>(Op.getOperand(1))->getZExtValue();
   EVT VT = Op.getValueType();
   EVT SrcVT = Op.getOperand(0).getValueType();
 
-  // For these types, we have some TableGen patterns except if the index is 1
-  if (((SrcVT == MVT::v4f16 && VT == MVT::v2f16) ||
-       (SrcVT == MVT::v4i16 && VT == MVT::v2i16)) &&
-      Start != 1)
+  if (VT.getScalarSizeInBits() == 16 && Start % 2 == 0) {
+    unsigned NumElt = VT.getVectorNumElements();
+    unsigned NumSrcElt = SrcVT.getVectorNumElements();
+    assert(NumElt % 2 == 0 && NumSrcElt % 2 == 0 && "expect legal types");
+
+    // We have some TableGen patterns for when the extracted vector is exactly
+    // the low or high half of the operand.
+    if ((NumSrcElt == 2 * NumElt) && (Start == 0 || Start == NumElt))
       return Op;
 
-  if (((SrcVT == MVT::v8f16 && VT == MVT::v4f16) ||
-       (SrcVT == MVT::v8i16 && VT == MVT::v4i16)) &&
-      (Start == 0 || Start == 4))
-    return Op;
+    // Extract 32-bit registers at a time.
+    EVT NewSrcVT = EVT::getVectorVT(*DAG.getContext(), MVT::i32, NumSrcElt / 2);
+    EVT NewVT = NumElt == 2
+                    ? MVT::i32
+                    : EVT::getVectorVT(*DAG.getContext(), MVT::i32, NumElt / 2);
+    SDValue Tmp = DAG.getNode(ISD::BITCAST, SL, NewSrcVT, Op.getOperand(0));
+
+    DAG.ExtractVectorElements(Tmp, Args, Start / 2, NumElt / 2);
+    if (NumElt == 2)
+      Tmp = Args[0];
+    else
+      Tmp = DAG.getBuildVector(NewVT, SL, Args);
 
-  if (((SrcVT == MVT::v16f16 && VT == MVT::v8f16) ||
-       (SrcVT == MVT::v16i16 && VT == MVT::v8i16)) &&
-      (Start == 0 || Start == 8))
-    return Op;
+    return DAG.getNode(ISD::BITCAST, SL, VT, Tmp);
+  }
 
   DAG.ExtractVectorElements(Op.getOperand(0), Args, Start,
                             VT.getVectorNumElements());
 
-  return DAG.getBuildVector(Op.getValueType(), SDLoc(Op), Args);
+  return DAG.getBuildVector(Op.getValueType(), SL, Args);
 }
 
 // TODO: Handle fabs too
 static SDValue peekFNeg(SDValue Val) {
   if (Val.getOpcode() == ISD::FNEG)
     return Val.getOperand(0);
 
   return Val;

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

SDValue SITargetLowering::lowerINSERT_SUBVECTOR(SDValue Op,
   SDValue Idx = Op.getOperand(2);
   EVT VecVT = Vec.getValueType();
   EVT InsVT = Ins.getValueType();
   EVT EltVT = VecVT.getVectorElementType();
   unsigned InsNumElts = InsVT.getVectorNumElements();
   unsigned IdxVal = cast<ConstantSDNode>(Idx)->getZExtValue();
   SDLoc SL(Op);
 
+  if (EltVT.getScalarSizeInBits() == 16 && IdxVal % 2 == 0) {
+    // Insert 32-bit registers at a time.
+    assert(InsNumElts % 2 == 0 && "expect legal vector types");
+
+    unsigned VecNumElts = VecVT.getVectorNumElements();
+    EVT NewVecVT =
+        EVT::getVectorVT(*DAG.getContext(), MVT::i32, VecNumElts / 2);
+    EVT NewInsVT = InsNumElts == 2 ? MVT::i32
+                                   : EVT::getVectorVT(*DAG.getContext(),
+                                                      MVT::i32, InsNumElts / 2);
+
+    Vec = DAG.getNode(ISD::BITCAST, SL, NewVecVT, Vec);
+    Ins = DAG.getNode(ISD::BITCAST, SL, NewInsVT, Ins);
+
+    for (unsigned I = 0; I != InsNumElts / 2; ++I) {
+      SDValue Elt;
+      if (InsNumElts == 2) {
+        Elt = Ins;
+      } else {
+        Elt = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL, MVT::i32, Ins,
+                          DAG.getConstant(I, SL, MVT::i32));
+      }
+      Vec = DAG.getNode(ISD::INSERT_VECTOR_ELT, SL, NewVecVT, Vec, Elt,
+                        DAG.getConstant(IdxVal / 2 + I, SL, MVT::i32));
+    }
+
+    return DAG.getNode(ISD::BITCAST, SL, VecVT, Vec);
+  }
+
   for (unsigned I = 0; I != InsNumElts; ++I) {
     SDValue Elt = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL, EltVT, Ins,
                               DAG.getConstant(I, SL, MVT::i32));
     Vec = DAG.getNode(ISD::INSERT_VECTOR_ELT, SL, VecVT, Vec, Elt,
                       DAG.getConstant(IdxVal + I, SL, MVT::i32));
   }
   return Vec;
 }

llvm/test/CodeGen/AMDGPU/cross-block-use-is-not-abi-copy.ll

 ; GCN-NEXT:    s_add_u32 s18, s18, func_v3i16@rel32@lo+4
 ; GCN-NEXT:    s_addc_u32 s19, s19, func_v3i16@rel32@hi+12
 ; GCN-NEXT:    s_swappc_b64 s[30:31], s[18:19]
 ; GCN-NEXT:    s_branch .LBB4_3
 ; GCN-NEXT:  .LBB4_2:
 ; GCN-NEXT:    v_mov_b32_e32 v0, 0
 ; GCN-NEXT:    v_mov_b32_e32 v1, 0
 ; GCN-NEXT:  .LBB4_3: ; %if.end
-; GCN-NEXT:    s_mov_b32 s4, 0xffff
-; GCN-NEXT:    v_bfi_b32 v0, s4, v0, v0
 ; GCN-NEXT:    global_store_short v[0:1], v1, off
 ; GCN-NEXT:    global_store_dword v[0:1], v0, off
 ; GCN-NEXT:    s_endpgm
 entry:
   br i1 %cond, label %if.then, label %if.else
 
 if.then:                                          ; preds = %entry
   br label %if.end

llvm/test/CodeGen/AMDGPU/extract-subvector-16bit.ll

 ; SI-NEXT:    s_setpc_b64 s[30:31]
 ;
 ; GFX9-LABEL: large_vector:
 ; GFX9:       ; %bb.0:
 ; GFX9-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX9-NEXT:    v_lshl_add_u32 v2, v1, 5, v0
 ; GFX9-NEXT:    ds_read2_b32 v[0:1], v2 offset1:1
 ; GFX9-NEXT:    ds_read2_b32 v[2:3], v2 offset0:2 offset1:3
-; GFX9-NEXT:    s_mov_b32 s4, 0xffff
-; GFX9-NEXT:    s_waitcnt lgkmcnt(1)
-; GFX9-NEXT:    v_bfi_b32 v4, s4, v0, v0
 ; GFX9-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX9-NEXT:    v_bfi_b32 v5, s4, v2, v2
-; GFX9-NEXT:    v_bfi_b32 v4, s4, v0, v4
-; GFX9-NEXT:    v_bfi_b32 v4, s4, v0, v4
-; GFX9-NEXT:    v_bfi_b32 v5, s4, v2, v5
-; GFX9-NEXT:    v_bfi_b32 v2, s4, v2, v5
-; GFX9-NEXT:    v_bfi_b32 v0, s4, v0, v4
 ; GFX9-NEXT:    s_setpc_b64 s[30:31]
   %idx = shl i32 %idxp, 4
 
   %i.0 = or i32 %idx, 0
   %p.0 = getelementptr half, ptr addrspace(3) %p, i32 %i.0
   %x.0 = load i16, ptr addrspace(3) %p.0, align 4
   %v0p = insertelement <8 x i16> poison, i16 %x.0, i32 0
   %i.1 = or i32 %idx, 1

llvm/test/CodeGen/AMDGPU/scalar_to_vector.v8i16.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -verify-machineinstrs -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck -check-prefix=GFX900 %s
 ; RUN: llc -verify-machineinstrs -mtriple=amdgcn-amd-amdhsa -mcpu=gfx906 < %s | FileCheck -check-prefixes=GFX906 %s
 ; RUN: llc -verify-machineinstrs -mtriple=amdgcn-amd-amdhsa -mcpu=gfx908 < %s | FileCheck -check-prefixes=GFX908 %s
 ; RUN: llc -verify-machineinstrs -mtriple=amdgcn-amd-amdhsa -mcpu=gfx90a < %s | FileCheck -check-prefixes=GFX90A %s
 
 define amdgpu_kernel void @scalar_to_vector_v8i16(<2 x i32> %in, ptr %out) #0 {
 ; GFX900-LABEL: scalar_to_vector_v8i16:
 ; GFX900:       ; %bb.0: ; %entry
 ; GFX900-NEXT:    s_load_dwordx4 s[0:3], s[4:5], 0x0
 ; GFX900-NEXT:    v_lshlrev_b32_e32 v4, 4, v0
 ; GFX900-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX900-NEXT:    s_pack_lh_b32_b16 s4, s0, s0
 ; GFX900-NEXT:    v_mov_b32_e32 v5, s3
 ; GFX900-NEXT:    v_add_co_u32_e32 v4, vcc, s2, v4
+; GFX900-NEXT:    v_mov_b32_e32 v0, s0
 ; GFX900-NEXT:    v_mov_b32_e32 v1, s1
+; GFX900-NEXT:    v_mov_b32_e32 v2, s0
 ; GFX900-NEXT:    v_mov_b32_e32 v3, s0
-; GFX900-NEXT:    v_mov_b32_e32 v0, s4
-; GFX900-NEXT:    v_mov_b32_e32 v2, s4
 ; GFX900-NEXT:    v_addc_co_u32_e32 v5, vcc, 0, v5, vcc
 ; GFX900-NEXT:    flat_store_dwordx4 v[4:5], v[0:3]
 ; GFX900-NEXT:    s_endpgm
 ;
 ; GFX906-LABEL: scalar_to_vector_v8i16:
 ; GFX906:       ; %bb.0: ; %entry
 ; GFX906-NEXT:    s_load_dwordx4 s[0:3], s[4:5], 0x0
 ; GFX906-NEXT:    v_lshlrev_b32_e32 v4, 4, v0
 ; GFX906-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX906-NEXT:    s_pack_lh_b32_b16 s4, s0, s0
 ; GFX906-NEXT:    v_mov_b32_e32 v5, s3
 ; GFX906-NEXT:    v_add_co_u32_e32 v4, vcc, s2, v4
+; GFX906-NEXT:    v_mov_b32_e32 v0, s0
 ; GFX906-NEXT:    v_mov_b32_e32 v1, s1
+; GFX906-NEXT:    v_mov_b32_e32 v2, s0
 ; GFX906-NEXT:    v_mov_b32_e32 v3, s0
-; GFX906-NEXT:    v_mov_b32_e32 v0, s4
-; GFX906-NEXT:    v_mov_b32_e32 v2, s4
 ; GFX906-NEXT:    v_addc_co_u32_e32 v5, vcc, 0, v5, vcc
 ; GFX906-NEXT:    flat_store_dwordx4 v[4:5], v[0:3]
 ; GFX906-NEXT:    s_endpgm
 ;
 ; GFX908-LABEL: scalar_to_vector_v8i16:
 ; GFX908:       ; %bb.0: ; %entry
 ; GFX908-NEXT:    s_load_dwordx4 s[0:3], s[4:5], 0x0
 ; GFX908-NEXT:    v_lshlrev_b32_e32 v4, 4, v0
 ; GFX908-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX908-NEXT:    s_pack_lh_b32_b16 s4, s0, s0
 ; GFX908-NEXT:    v_mov_b32_e32 v5, s3
 ; GFX908-NEXT:    v_add_co_u32_e32 v4, vcc, s2, v4
+; GFX908-NEXT:    v_mov_b32_e32 v0, s0
 ; GFX908-NEXT:    v_mov_b32_e32 v1, s1
+; GFX908-NEXT:    v_mov_b32_e32 v2, s0
 ; GFX908-NEXT:    v_mov_b32_e32 v3, s0
-; GFX908-NEXT:    v_mov_b32_e32 v0, s4
-; GFX908-NEXT:    v_mov_b32_e32 v2, s4
 ; GFX908-NEXT:    v_addc_co_u32_e32 v5, vcc, 0, v5, vcc
 ; GFX908-NEXT:    flat_store_dwordx4 v[4:5], v[0:3]
 ; GFX908-NEXT:    s_endpgm
 ;
 ; GFX90A-LABEL: scalar_to_vector_v8i16:
 ; GFX90A:       ; %bb.0: ; %entry
 ; GFX90A-NEXT:    s_load_dwordx4 s[0:3], s[4:5], 0x0
 ; GFX90A-NEXT:    v_lshlrev_b32_e32 v4, 4, v0
 ; GFX90A-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX90A-NEXT:    s_pack_lh_b32_b16 s4, s0, s0
 ; GFX90A-NEXT:    v_mov_b32_e32 v5, s3
 ; GFX90A-NEXT:    v_add_co_u32_e32 v4, vcc, s2, v4
+; GFX90A-NEXT:    v_mov_b32_e32 v0, s0
 ; GFX90A-NEXT:    v_mov_b32_e32 v1, s1
+; GFX90A-NEXT:    v_mov_b32_e32 v2, s0
 ; GFX90A-NEXT:    v_mov_b32_e32 v3, s0
-; GFX90A-NEXT:    v_mov_b32_e32 v0, s4
-; GFX90A-NEXT:    v_mov_b32_e32 v2, s4
 ; GFX90A-NEXT:    v_addc_co_u32_e32 v5, vcc, 0, v5, vcc
 ; GFX90A-NEXT:    flat_store_dwordx4 v[4:5], v[0:3]
 ; GFX90A-NEXT:    s_endpgm
 entry:
   %val.1.i32 = extractelement <2 x i32> %in, i64 0
   %val.2.vec2.i16 = bitcast i32 %val.1.i32 to <2 x i16>
   %val.3.vec8.i16 = shufflevector <2 x i16> %val.2.vec2.i16, <2 x i16> %val.2.vec2.i16, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 0, i32 1, i32 2, i32 3>