[llvm] 412d59f - [DAG] combineShiftToMULH - handle zext nneg as sext

[Simon Pilgrim via llvm-commits]

Author: Simon Pilgrim
Date: 2024-09-30T12:12:32+01:00
New Revision: 412d59f0a510a05c08ed45545943dfd2f901bc5d

URL: https://github.com/llvm/llvm-project/commit/412d59f0a510a05c08ed45545943dfd2f901bc5d
DIFF: https://github.com/llvm/llvm-project/commit/412d59f0a510a05c08ed45545943dfd2f901bc5d.diff

LOG: [DAG] combineShiftToMULH - handle zext nneg as sext

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

Added: 
    

Modified: 
    llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
    llvm/test/CodeGen/X86/pmulh.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index c6f6fc25080541..65a620b70d8f0c 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/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();

diff  --git a/llvm/test/CodeGen/X86/pmulh.ll b/llvm/test/CodeGen/X86/pmulh.ll
index 300da68d9a3b34..4e9ce012aff219 100644
--- a/llvm/test/CodeGen/X86/pmulh.ll
+++ b/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>