DAG: Combine extract_vector_elt of concat_vectors

[arsenm]

Summary

Fixes extra canonicalize regressions when legalizing
vector fminnum/fmaxnum.

[RKSimon]

Handling INSERT_SUBVECTOR might make sense as well?

[arsenm]

In D51737#1226081, @RKSimon wrote:

Handling INSERT_SUBVECTOR might make sense as well?

Probably. I think that's a separate patch, since now both indexes need to be considered. I think this only works if the indexes match

[RKSimon]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15493 ↗

(On Diff #164234)

Do we need a TLI.isOperationLegal(ISD::EXTRACT_VECTOR_ELT, ConcatVT) check?

[RKSimon]

@arsenm What's happening with this patch?

[arsenm]

Rebase

[lebedev.ri]

lg

[RKSimon]

LGTM

[arsenm]

70726cec5b16cfa28b18084e276f9b64f70e6550

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (VecOp.getOpcode() == ISD::BITCAST) {
 
       VecOp = VecOp.getOperand(0);
     }
     if (ISD::isNormalLoad(VecOp.getNode())) {
       LN0 = cast<LoadSDNode>(VecOp);
       Elt = (Idx < (int)NumElts) ? Idx : Idx - (int)NumElts;
       Index = DAG.getConstant(Elt, DL, Index.getValueType());
     }
+  } else if (VecOp.getOpcode() == ISD::CONCAT_VECTORS &&
+             !BCNumEltsChanged && VecVT.getVectorElementType() == ScalarVT) {
+    // extract_vector_elt (concat_vectors v2i16:a, v2i16:b), 0
+    //      -> extract_vector_elt a, 0
+    // extract_vector_elt (concat_vectors v2i16:a, v2i16:b), 1
+    //      -> extract_vector_elt a, 1
+    // extract_vector_elt (concat_vectors v2i16:a, v2i16:b), 2
+    //      -> extract_vector_elt b, 0
+    // extract_vector_elt (concat_vectors v2i16:a, v2i16:b), 3
+    //      -> extract_vector_elt b, 1
+    SDLoc SL(N);
+    EVT ConcatVT = VecOp.getOperand(0).getValueType();
+    unsigned ConcatNumElts = ConcatVT.getVectorNumElements();
+    SDValue NewIdx = DAG.getConstant(Elt % ConcatNumElts, SL,
+                                     Index.getValueType());
+
+    SDValue ConcatOp = VecOp.getOperand(Elt / ConcatNumElts);
+    SDValue Elt = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL,
+                              ConcatVT.getVectorElementType(),
+                              ConcatOp, NewIdx);
+    return DAG.getNode(ISD::BITCAST, SL, ScalarVT, Elt);
   }
 
   // Make sure we found a non-volatile load and the extractelement is
   // the only use.
   if (!LN0 || !LN0->hasNUsesOfValue(1,0) || !LN0->isSimple())
     return SDValue();
 
   // If Idx was -1 above, Elt is going to be -1, so just return undef.

llvm/test/CodeGen/AMDGPU/reduction.ll

entry:
   ret i16 %res
 }
 
 ; GCN-LABEL: {{^}}reduction_maxnum_v4f16:
 ; GFX9: s_waitcnt
 ; GFX9-NEXT: v_pk_max_f16 [[CANON1:v[0-9]+]], v1, v1
 ; GFX9-NEXT: v_pk_max_f16 [[CANON0:v[0-9]+]], v0, v0
 ; GFX9-NEXT: v_pk_max_f16 [[MAX:v[0-9]+]], [[CANON0]], [[CANON1]]{{$}}
+; GFX9-NEXT: v_max_f16_sdwa v{{[0-9]+}}, [[MAX]], [[MAX]] dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:DWORD src1_sel:WORD_1
 
-; FIXME: Extra canonicalize leftover
-; GFX9-NEXT: v_max_f16_sdwa [[TMP:v[0-9]+]], [[MAX]], [[MAX]] dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_1 src1_sel:WORD_1
-; GFX9-NEXT: v_max_f16_e32 v0, [[MAX]], [[TMP]]
 
 ; VI-DAG: v_max_f16_sdwa [[CANON1:v[0-9]+]], v0, v0 dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_1 src1_sel:WORD_1
 ; VI-DAG: v_max_f16_sdwa [[CANON3:v[0-9]+]], v1, v1 dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_1 src1_sel:WORD_1
 ; VI-DAG: v_max_f16_e32 [[CANON0:v[0-9]+]], v0, v0
 ; VI-DAG: v_max_f16_e32 [[CANON2:v[0-9]+]], v1, v1
 
 ; VI-DAG: v_max_f16_e32 [[MAX0:v[0-9]+]], [[CANON1]], [[CANON3]]
 ; VI-DAG: v_max_f16_e32 [[MAX1:v[0-9]+]], [[CANON0]], [[CANON2]]
 ; VI: v_max_f16_e32 v0, [[MAX1]], [[MAX0]]
 define half @reduction_maxnum_v4f16(<4 x half> %vec4) {
 entry:
   %rdx.shuf = shufflevector <4 x half> %vec4, <4 x half> undef, <4 x i32> <i32 2, i32 3, i32 undef, i32 undef>
   %rdx.minmax = call <4 x half> @llvm.maxnum.v4f16(<4 x half> %vec4, <4 x half> %rdx.shuf)
   %rdx.shuf1 = shufflevector <4 x half> %rdx.minmax, <4 x half> undef, <4 x i32> <i32 1, i32 undef, i32 undef, i32 undef>
   %rdx.minmax3 = call <4 x half> @llvm.maxnum.v4f16(<4 x half> %rdx.minmax, <4 x half> %rdx.shuf1)
   %res = extractelement <4 x half> %rdx.minmax3, i32 0
   ret half %res
 }
 
 ; GCN-LABEL: {{^}}reduction_minnum_v4f16:
 ; GFX9: s_waitcnt
 ; GFX9-NEXT: v_pk_max_f16 [[CANON1:v[0-9]+]], v1, v1
 ; GFX9-NEXT: v_pk_max_f16 [[CANON0:v[0-9]+]], v0, v0
 ; GFX9-NEXT: v_pk_min_f16 [[MIN:v[0-9]+]], [[CANON0]], [[CANON1]]{{$}}
-
-; FIXME: Extra canonicalize leftover
-; GFX9-NEXT: v_max_f16_sdwa [[TMP:v[0-9]+]], [[MIN]], [[MIN]] dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_1 src1_sel:WORD_1
-; GFX9-NEXT: v_min_f16_e32 v0, [[MIN]], [[TMP]]
-
+; GFX9-NEXT: v_min_f16_sdwa v{{[0-9]+}}, [[MIN]], [[MIN]] dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:DWORD src1_sel:WORD_1
 
 ; VI-DAG: v_max_f16_sdwa [[CANON1:v[0-9]+]], v0, v0 dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_1 src1_sel:WORD_1
 ; VI-DAG: v_max_f16_sdwa [[CANON3:v[0-9]+]], v1, v1 dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_1 src1_sel:WORD_1
 ; VI-DAG: v_max_f16_e32 [[CANON0:v[0-9]+]], v0, v0
 ; VI-DAG: v_max_f16_e32 [[CANON2:v[0-9]+]], v1, v1
 
 ; VI-DAG: v_min_f16_e32 [[MAX0:v[0-9]+]], [[CANON1]], [[CANON3]]
 ; VI-DAG: v_min_f16_e32 [[MAX1:v[0-9]+]], [[CANON0]], [[CANON2]]

 ; XVI-NEXT: v_max_f16_e32 v0, v0, v1
 ; XVI-NEXT: v_max_f16_e32 v0, v0, v2
 ; XVI-NEXT: s_setpc_b64
 
 ; GFX9: s_waitcnt
 ; GFX9-NEXT: v_pk_max_f16 [[CANON1:v[0-9]+]], v1, v1
 ; GFX9-NEXT: v_pk_max_f16 [[CANON0:v[0-9]+]], v0, v0
 ; GFX9-NEXT: v_pk_max_f16 [[MAX:v[0-9]+]], [[CANON0]], [[CANON1]]{{$}}
-
-; FIXME: Extra canonicalize leftover
-; GFX9-NEXT: v_max_f16_sdwa [[TMP:v[0-9]+]], [[MAX]], [[MAX]] dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_1 src1_sel:WORD_1
-; GFX9-NEXT: v_max_f16_e32 v0, [[MAX]], [[TMP]]
+; GFX9-NEXT: v_max_f16_sdwa v{{[0-9]+}}, [[MAX]], [[MAX]] dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:DWORD src1_sel:WORD_1
 
 ; VI-DAG: v_max_f16_sdwa [[CANON1:v[0-9]+]], v0, v0 dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_1 src1_sel:WORD_1
 ; VI-DAG: v_max_f16_sdwa [[CANON3:v[0-9]+]], v1, v1 dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_1 src1_sel:WORD_1
 ; VI-DAG: v_max_f16_e32 [[CANON0:v[0-9]+]], v0, v0
 ; VI-DAG: v_max_f16_e32 [[CANON2:v[0-9]+]], v1, v1
 
 ; VI-DAG: v_max_f16_e32 [[MAX0:v[0-9]+]], [[CANON1]], [[CANON3]]
 ; VI-DAG: v_max_f16_e32 [[MAX1:v[0-9]+]], [[CANON0]], [[CANON2]]

 ; XVI-NEXT: v_min_f16_e32 v0, v0, v1
 ; XVI-NEXT: v_min_f16_e32 v0, v0, v2
 ; XVI-NEXT: s_setpc_b64
 
 ; GFX9: s_waitcnt
 ; GFX9-NEXT: v_pk_max_f16 [[CANON1:v[0-9]+]], v1, v1
 ; GFX9-NEXT: v_pk_max_f16 [[CANON0:v[0-9]+]], v0, v0
 ; GFX9-NEXT: v_pk_min_f16 [[MIN:v[0-9]+]], [[CANON0]], [[CANON1]]{{$}}
-
-; FIXME: Extra canonicalize leftover
-; GFX9-NEXT: v_max_f16_sdwa [[TMP:v[0-9]+]], [[MIN]], [[MIN]] dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_1 src1_sel:WORD_1
-; GFX9-NEXT: v_min_f16_e32 v0, [[MIN]], [[TMP]]
-
+; GFX9-NEXT: v_min_f16_sdwa v{{[0-9]+}}, [[MIN]], [[MIN]] dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:DWORD src1_sel:WORD_1
 
 ; VI-DAG: v_max_f16_sdwa [[CANON1:v[0-9]+]], v0, v0 dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_1 src1_sel:WORD_1
 ; VI-DAG: v_max_f16_sdwa [[CANON3:v[0-9]+]], v1, v1 dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_1 src1_sel:WORD_1
 ; VI-DAG: v_max_f16_e32 [[CANON0:v[0-9]+]], v0, v0
 ; VI-DAG: v_max_f16_e32 [[CANON2:v[0-9]+]], v1, v1
 
 ; VI-DAG: v_min_f16_e32 [[MAX0:v[0-9]+]], [[CANON1]], [[CANON3]]
 ; VI-DAG: v_min_f16_e32 [[MAX1:v[0-9]+]], [[CANON0]], [[CANON2]]