AMDGPU: Fix assert on trunc from bitcast of build_vector

[arsenm]

Summary

The v2i64 argument is lowered to a bitcast of v4i32 build_vector.
This would then attempt to use the i32-element as the source of the
vector truncate. This really would need to collect 2 elements from the
build_vector to produce the intended truncate.

[kzhuravl]

LGTM

[arsenm]

r353202

lib/Target/AMDGPU/AMDGPUISelLowering.cpp

 SDValue AMDGPUTargetLowering::performTruncateCombine(
   SDNode *N, DAGCombinerInfo &DCI) const {
   SDLoc SL(N);
   SelectionDAG &DAG = DCI.DAG;
   EVT VT = N->getValueType(0);
   SDValue Src = N->getOperand(0);
 
   // vt1 (truncate (bitcast (build_vector vt0:x, ...))) -> vt1 (bitcast vt0:x)
-  if (Src.getOpcode() == ISD::BITCAST) {
+  if (Src.getOpcode() == ISD::BITCAST && !VT.isVector()) {
     SDValue Vec = Src.getOperand(0);
     if (Vec.getOpcode() == ISD::BUILD_VECTOR) {
       SDValue Elt0 = Vec.getOperand(0);
       EVT EltVT = Elt0.getValueType();
       if (VT.getSizeInBits() <= EltVT.getSizeInBits()) {
         if (EltVT.isFloatingPoint()) {
           Elt0 = DAG.getNode(ISD::BITCAST, SL,
                              EltVT.changeTypeToInteger(), Elt0);

test/CodeGen/AMDGPU/trunc-combine.ll

+; RUN: llc -march=amdgcn -mcpu=tahiti -verify-machineinstrs< %s | FileCheck -enable-var-scope -check-prefixes=GCN,SI %s
 ; RUN: llc -march=amdgcn -mcpu=fiji -verify-machineinstrs< %s | FileCheck -enable-var-scope -check-prefixes=GCN,VI %s
 
 ; Make sure high constant 0 isn't pointlessly materialized
 ; GCN-LABEL: {{^}}trunc_bitcast_i64_lshr_32_i16:
 ; GCN: s_waitcnt
 ; GCN-NEXT: v_mov_b32_e32 v0, v1
 ; GCN-NEXT: s_setpc_b64
 define i16 @trunc_bitcast_i64_lshr_32_i16(i64 %bar) {
define i32 @trunc_bitcast_i64_lshr_32_i32(i64 %bar) {
   %trunc = trunc i64 %srl to i32
   ret i32 %trunc
 }
 
 ; GCN-LABEL: {{^}}trunc_bitcast_v2i32_to_i16:
 ; GCN: _load_dword
 ; GCN-NOT: _load_dword
 ; GCN-NOT: v_mov_b32
-; GCN: v_add_u16_e32 v0, 4, v0
+; VI: v_add_u16_e32 v0, 4, v0
 define i16 @trunc_bitcast_v2i32_to_i16(<2 x i32> %bar) {
   %load0 = load i32, i32 addrspace(1)* undef
   %load1 = load i32, i32 addrspace(1)* null
   %insert.0 = insertelement <2 x i32> undef, i32 %load0, i32 0
   %insert.1 = insertelement <2 x i32> %insert.0, i32 99, i32 1
   %bc = bitcast <2 x i32> %insert.1 to i64
   %trunc = trunc i64 %bc to i16
   %add = add i16 %trunc, 4
   ret i16 %add
 }
 
 ; Make sure there's no crash if the source vector type is FP
 ; GCN-LABEL: {{^}}trunc_bitcast_v2f32_to_i16:
 ; GCN: _load_dword
 ; GCN-NOT: _load_dword
 ; GCN-NOT: v_mov_b32
-; GCN: v_add_u16_e32 v0, 4, v0
+; VI: v_add_u16_e32 v0, 4, v0
 define i16 @trunc_bitcast_v2f32_to_i16(<2 x float> %bar) {
   %load0 = load float, float addrspace(1)* undef
   %load1 = load float, float addrspace(1)* null
   %insert.0 = insertelement <2 x float> undef, float %load0, i32 0
   %insert.1 = insertelement <2 x float> %insert.0, float 4.0, i32 1
   %bc = bitcast <2 x float> %insert.1 to i64
   %trunc = trunc i64 %bc to i16
   %add = add i16 %trunc, 4
bb:
   %tmp11 = insertelement <2 x i32> undef, i32 %tmp10, i32 0
   %tmp12 = insertelement <2 x i32> %tmp11, i32 undef, i32 1
   %tmp13 = lshr <2 x i32> %tmp12, <i32 16, i32 16>
   %tmp14 = trunc <2 x i32> %tmp13 to <2 x i16>
   %tmp15 = getelementptr inbounds <2 x i16>, <2 x i16> addrspace(1)* %arg2, i64 undef
   store <2 x i16> %tmp14, <2 x i16> addrspace(1)* %tmp15, align 4
   ret void
 }
+
+; GCN-LABEL: {{^}}trunc_v2i64_arg_to_v2i16:
+; GCN: v_lshlrev_b32_e32 v1, 16, v2
+
+; SI-NEXT: v_and_b32_e32 v0, 0xffff, v0
+; SI-NEXT: v_or_b32_e32 v0, v0, v1
+; SI-NEXT: v_lshrrev_b32_e32 v1, 16, v0
+
+; VI-NEXT: v_or_b32_sdwa v0, v0, v1 dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_0 src1_sel:DWORD
+
+; GCN-NEXT: s_setpc_b64
+define <2 x i16> @trunc_v2i64_arg_to_v2i16(<2 x i64> %arg0) #0 {
+  %trunc = trunc <2 x i64> %arg0 to <2 x i16>
+  ret <2 x i16> %trunc
+}