AMDGPU: Extend extract_vector_elt fneg combine to fabs

[arsenm]

Summary

Fixes a regression in a future commit.

[cfang]

LGTM,

[arsenm]

r330981

lib/Target/AMDGPU/SIISelLowering.cpp

SDValue SITargetLowering::performCvtPkRTZCombine(SDNode *N,
   return SDValue();
 }
 
 SDValue SITargetLowering::performExtractVectorEltCombine(
   SDNode *N, DAGCombinerInfo &DCI) const {
   SDValue Vec = N->getOperand(0);
 
   SelectionDAG &DAG = DCI.DAG;
-  if (Vec.getOpcode() == ISD::FNEG && allUsesHaveSourceMods(N)) {
+  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(ISD::FNEG, SL, EltVT, Elt);
+    return DAG.getNode(Vec.getOpcode(), SL, EltVT, Elt);
   }
 
   return SDValue();
 }
 
 static bool convertBuildVectorCastElt(SelectionDAG &DAG,
                                       SDValue &Lo, SDValue &Hi) {
   if (Hi.getOpcode() == ISD::BITCAST &&

test/CodeGen/AMDGPU/fabs.f16.ll

define amdgpu_kernel void @v_fabs_fold_v2f16(<2 x half> addrspace(1)* %out, <2 x half> addrspace(1)* %in) #0 {
   %gep = getelementptr <2 x half>, <2 x half> addrspace(1)* %in, i32 %tid
   %val = load <2 x half>, <2 x half> addrspace(1)* %gep
   %fabs = call <2 x half> @llvm.fabs.v2f16(<2 x half> %val)
   %fmul = fmul <2 x half> %fabs, %val
   store <2 x half> %fmul, <2 x half> addrspace(1)* %out
   ret void
 }
 
+; GCN-LABEL: {{^}}v_extract_fabs_fold_v2f16:
+; GCN-DAG: {{flat|global}}_load_dword [[VAL:v[0-9]+]]
+; CI-DAG: v_mul_f32_e32 v{{[0-9]+}}, 4.0, v{{[0-9]+}}
+; CI-DAG: v_add_f32_e32 v{{[0-9]+}}, 2.0, v{{[0-9]+}}
+
+; GFX89-DAG: v_mul_f16_e32 v{{[0-9]+}}, -4.0, [[VAL]]
+; GFX89-DAG: v_mov_b32_e32 [[CONST2:v[0-9]+]], 0x4000
+; GFX89-DAG: v_sub_f16_sdwa v{{[0-9]+}}, [[CONST2]], [[VAL]] dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:DWORD src1_sel:WORD_1
+define amdgpu_kernel void @v_extract_fabs_fold_v2f16(<2 x half> addrspace(1)* %in) #0 {
+  %tid = call i32 @llvm.amdgcn.workitem.id.x()
+  %gep.in = getelementptr inbounds <2 x half>, <2 x half> addrspace(1)* %in, i32 %tid
+  %val = load <2 x half>, <2 x half> addrspace(1)* %gep.in
+  %fabs = call <2 x half> @llvm.fabs.v2f16(<2 x half> %val)
+  %elt0 = extractelement <2 x half> %fabs, i32 0
+  %elt1 = extractelement <2 x half> %fabs, i32 1
+
+  %fmul0 = fmul half %elt0, 4.0
+  %fadd1 = fadd half %elt1, 2.0
+  store volatile half %fmul0, half addrspace(1)* undef
+  store volatile half %fadd1, half addrspace(1)* undef
+  ret void
+}
+
+; GCN-LABEL: {{^}}v_extract_fabs_no_fold_v2f16:
+; GCN: {{flat|global}}_load_dword [[VAL:v[0-9]+]]
+
+; FIXME: Extra bfe on VI
+; GFX9-NOT: v_bfe_u32
+; VI: v_bfe_u32
+; GCN: v_and_b32_e32 [[AND:v[0-9]+]], 0x7fff7fff, [[VAL]]
+; GFX9: global_store_short_d16_hi v{{\[[0-9]+:[0-9]+\]}}, [[AND]], off
+define amdgpu_kernel void @v_extract_fabs_no_fold_v2f16(<2 x half> addrspace(1)* %in) #0 {
+  %tid = call i32 @llvm.amdgcn.workitem.id.x()
+  %gep.in = getelementptr inbounds <2 x half>, <2 x half> addrspace(1)* %in, i32 %tid
+  %val = load <2 x half>, <2 x half> addrspace(1)* %gep.in
+  %fabs = call <2 x half> @llvm.fabs.v2f16(<2 x half> %val)
+  %elt0 = extractelement <2 x half> %fabs, i32 0
+  %elt1 = extractelement <2 x half> %fabs, i32 1
+  store volatile half %elt0, half addrspace(1)* undef
+  store volatile half %elt1, half addrspace(1)* undef
+  ret void
+}
+
 declare half @llvm.fabs.f16(half) #1
 declare <2 x half> @llvm.fabs.v2f16(<2 x half>) #1
 declare <4 x half> @llvm.fabs.v4f16(<4 x half>) #1
 declare i32 @llvm.amdgcn.workitem.id.x() #1
 
 attributes #0 = { nounwind }
 attributes #1 = { nounwind readnone }

test/CodeGen/AMDGPU/mad-mix.ll

define float @v_mad_mix_f32_preextractfneg_f16hi_f16lo_f16lo(i32 %src0.arg, half %src1, half %src2) #0 {
   %src0 = extractelement <2 x half> %fneg, i32 1
   %src0.ext = fpext half %src0 to float
   %src1.ext = fpext half %src1 to float
   %src2.ext = fpext half %src2 to float
   %result = tail call float @llvm.fmuladd.f32(float %src0.ext, float %src1.ext, float %src2.ext)
   ret float %result
 }
 
-; FIXME: Should be able to fold
 ; GCN-LABEL: {{^}}v_mad_mix_f32_preextractfabs_f16hi_f16lo_f16lo:
 ; GFX9: s_waitcnt
-; GFX9-NEXT: v_and_b32_e32 v0, 0x7fff0000, v0
-; GFX9-NEXT: v_mad_mix_f32 v0, v0, v1, v2 op_sel:[1,0,0] op_sel_hi:[1,1,1]
+; GFX9-NEXT: v_mad_mix_f32 v0, |v0|, v1, v2 op_sel:[1,0,0] op_sel_hi:[1,1,1]
 ; GFX9-NEXT: s_setpc_b64
 define float @v_mad_mix_f32_preextractfabs_f16hi_f16lo_f16lo(i32 %src0.arg, half %src1, half %src2) #0 {
   %src0.arg.bc = bitcast i32 %src0.arg to <2 x half>
   %fabs = call <2 x half> @llvm.fabs.v2f16(<2 x half> %src0.arg.bc)
   %src0 = extractelement <2 x half> %fabs, i32 1
   %src0.ext = fpext half %src0 to float
   %src1.ext = fpext half %src1 to float
   %src2.ext = fpext half %src2 to float
   %result = tail call float @llvm.fmuladd.f32(float %src0.ext, float %src1.ext, float %src2.ext)
   ret float %result
 }
 
-; FIXME: Should be able to fold
 ; GCN-LABEL: {{^}}v_mad_mix_f32_preextractfabsfneg_f16hi_f16lo_f16lo:
 ; GFX9: s_waitcnt
-; GFX9-NEXT: v_and_b32_e32 v0, 0x7fff0000, v0
-; GFX9-NEXT: v_mad_mix_f32 v0, -v0, v1, v2 op_sel:[1,0,0] op_sel_hi:[1,1,1]
+; GFX9-NEXT: v_mad_mix_f32 v0, -|v0|, v1, v2 op_sel:[1,0,0] op_sel_hi:[1,1,1]
 ; GFX9-NEXT: s_setpc_b64
 define float @v_mad_mix_f32_preextractfabsfneg_f16hi_f16lo_f16lo(i32 %src0.arg, half %src1, half %src2) #0 {
   %src0.arg.bc = bitcast i32 %src0.arg to <2 x half>
   %fabs = call <2 x half> @llvm.fabs.v2f16(<2 x half> %src0.arg.bc)
   %fneg.fabs = fsub <2 x half> <half -0.0, half -0.0>, %fabs
   %src0 = extractelement <2 x half> %fneg.fabs, i32 1
   %src0.ext = fpext half %src0 to float
   %src1.ext = fpext half %src1 to float