[DAG] combineShiftToMULH - handle zext nneg as sext

Fixes poor codegen on AVX512 targets for a test case from #109790

Author: RKSimon

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -10287,8 +10287,10 @@ static SDValue combineShiftToMULH(SDNode *N, const SDLoc &DL, SelectionDAG &DAG,
   SDValue LeftOp = ShiftOperand.getOperand(0);
   SDValue RightOp = ShiftOperand.getOperand(1);
 
-  bool IsSignExt = LeftOp.getOpcode() == ISD::SIGN_EXTEND;
-  bool IsZeroExt = LeftOp.getOpcode() == ISD::ZERO_EXTEND;
+  // Treat zext nneg as sext - we might need to support handling these as zext
+  // as well in the future, but for now just prefer sext.
+  bool IsSignExt = sd_match(LeftOp, m_SExtLike(m_Value()));
+  bool IsZeroExt = sd_match(LeftOp, m_ZExt(m_Value()));
 
   if (!IsSignExt && !IsZeroExt)
     return SDValue();

llvm/test/CodeGen/X86/pmulh.ll

@@ -953,39 +953,15 @@ define void @PR109790(ptr sret([32 x i8]) %ret, ptr %a) {
 ; SSE-NEXT:    movdqa %xmm0, 16(%rdi)
 ; SSE-NEXT:    retq
 ;
-; AVX2-LABEL: PR109790:
-; AVX2:       # %bb.0:
-; AVX2-NEXT:    movq %rdi, %rax
-; AVX2-NEXT:    vmovdqa (%rsi), %ymm0
-; AVX2-NEXT:    vpand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %ymm0, %ymm0
-; AVX2-NEXT:    vpmulhw {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %ymm0, %ymm0 # [64536,64536,64536,64536,64536,64536,64536,64536,64536,64536,64536,64536,64536,64536,64536,64536]
-; AVX2-NEXT:    vmovdqa %ymm0, (%rdi)
-; AVX2-NEXT:    vzeroupper
-; AVX2-NEXT:    retq
-;
-; AVX512F-LABEL: PR109790:
-; AVX512F:       # %bb.0:
-; AVX512F-NEXT:    movq %rdi, %rax
-; AVX512F-NEXT:    vmovdqa (%rsi), %ymm0
-; AVX512F-NEXT:    vpand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %ymm0, %ymm0
-; AVX512F-NEXT:    vpmovzxwd {{.*#+}} zmm0 = ymm0[0],zero,ymm0[1],zero,ymm0[2],zero,ymm0[3],zero,ymm0[4],zero,ymm0[5],zero,ymm0[6],zero,ymm0[7],zero,ymm0[8],zero,ymm0[9],zero,ymm0[10],zero,ymm0[11],zero,ymm0[12],zero,ymm0[13],zero,ymm0[14],zero,ymm0[15],zero
-; AVX512F-NEXT:    vpmulld {{\.?LCPI[0-9]+_[0-9]+}}(%rip){1to16}, %zmm0, %zmm0
-; AVX512F-NEXT:    vpsrld $16, %zmm0, %zmm0
-; AVX512F-NEXT:    vpmovdw %zmm0, (%rdi)
-; AVX512F-NEXT:    vzeroupper
-; AVX512F-NEXT:    retq
-;
-; AVX512BW-LABEL: PR109790:
-; AVX512BW:       # %bb.0:
-; AVX512BW-NEXT:    movq %rdi, %rax
-; AVX512BW-NEXT:    vmovdqa (%rsi), %ymm0
-; AVX512BW-NEXT:    vpand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %ymm0, %ymm0
-; AVX512BW-NEXT:    vpmovzxwd {{.*#+}} zmm0 = ymm0[0],zero,ymm0[1],zero,ymm0[2],zero,ymm0[3],zero,ymm0[4],zero,ymm0[5],zero,ymm0[6],zero,ymm0[7],zero,ymm0[8],zero,ymm0[9],zero,ymm0[10],zero,ymm0[11],zero,ymm0[12],zero,ymm0[13],zero,ymm0[14],zero,ymm0[15],zero
-; AVX512BW-NEXT:    vpmaddwd {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %zmm0, %zmm0 # [64536,0,64536,0,64536,0,64536,0,64536,0,64536,0,64536,0,64536,0,64536,0,64536,0,64536,0,64536,0,64536,0,64536,0,64536,0,64536,0]
-; AVX512BW-NEXT:    vpsrld $16, %zmm0, %zmm0
-; AVX512BW-NEXT:    vpmovdw %zmm0, (%rdi)
-; AVX512BW-NEXT:    vzeroupper
-; AVX512BW-NEXT:    retq
+; AVX-LABEL: PR109790:
+; AVX:       # %bb.0:
+; AVX-NEXT:    movq %rdi, %rax
+; AVX-NEXT:    vmovdqa (%rsi), %ymm0
+; AVX-NEXT:    vpand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %ymm0, %ymm0
+; AVX-NEXT:    vpmulhw {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %ymm0, %ymm0 # [64536,64536,64536,64536,64536,64536,64536,64536,64536,64536,64536,64536,64536,64536,64536,64536]
+; AVX-NEXT:    vmovdqa %ymm0, (%rdi)
+; AVX-NEXT:    vzeroupper
+; AVX-NEXT:    retq
   %load = load <16 x i16>, ptr %a, align 32
   %and = and <16 x i16> %load, <i16 32767, i16 32767, i16 32767, i16 32767, i16 32767, i16 32767, i16 32767, i16 32767, i16 32767, i16 32767, i16 32767, i16 32767, i16 32767, i16 32767, i16 32767, i16 32767>
   %ext = zext nneg <16 x i16> %and to <16 x i32>