[SDAG] narrow truncated sign_extend_inreg

[spatel]

Summary

trunc (sign_ext_inreg X, iM) to iN --> sign_ext_inreg (trunc X to iN), iM

There are improvements on existing tests from this, and there are a pair of large regressions in D127115 for Thumb2 caused by not folding this pattern.

[spatel]

Patch updated:
On closer inspection, we can generalize the fold to only match trunc-of-sext (don't need to match the preceding extend), and that shows improvements on existing tests for x86 and (I think) AMDGPU.

[arsenm]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

13138

Shouldn't really require legal types, but the DAG is lacking in adequate legality checks for sext_inreg

llvm/test/CodeGen/AMDGPU/mul_int24.ll

184–186

This is a pretty big improvement

[spatel]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

13138

There are a couple of other sext_inreg folds in DAGCombiner that use something like this:

if (!LegalOperations ||
    TLI.getOperationAction(ISD::SIGN_EXTEND_INREG, ExtVT) ==
    TargetLowering::Legal)

I tried that and don't see any differences vs. this check, so I could use that too, but I think the current predicate would be more conservative just in case there's some unexpected type conversions during legalization.

llvm/test/CodeGen/AMDGPU/mul_int24.ll

184–186

Thanks for confirming!

[arsenm]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

13138

This is one of the places where we really need 2 type checks, but I think checking the ExtVT legality is probably enough here

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

     // if the source is larger than the dest, than we just need the truncate.
     if (N0.getOperand(0).getValueType().bitsGT(VT))
       return DAG.getNode(ISD::TRUNCATE, SDLoc(N), VT, N0.getOperand(0));
     // if the source and dest are the same type, we can drop both the extend
     // and the truncate.
     return N0.getOperand(0);
   }
 
+  // Try to narrow a truncate-of-sext_in_reg to the destination type:
+  // trunc (sign_ext_inreg X, iM) to iN --> sign_ext_inreg (trunc X to iN), iM
+  if (!LegalTypes && N0.getOpcode() == ISD::SIGN_EXTEND_INREG &&
+      N0.hasOneUse()) {
+    SDValue X = N0.getOperand(0);
+    SDValue ExtVal = N0.getOperand(1);
+    EVT ExtVT = cast<VTSDNode>(ExtVal)->getVT();
+    if (ExtVT.bitsLT(VT)) {
+      SDValue TrX = DAG.getNode(ISD::TRUNCATE, SDLoc(N), VT, X);
+      return DAG.getNode(ISD::SIGN_EXTEND_INREG, SDLoc(N), VT, TrX, ExtVal);
+    }
+  }
+
   // If this is anyext(trunc), don't fold it, allow ourselves to be folded.
   if (N->hasOneUse() && (N->use_begin()->getOpcode() == ISD::ANY_EXTEND))
     return SDValue();
 
   // Fold extract-and-trunc into a narrow extract. For example:
   //   i64 x = EXTRACT_VECTOR_ELT(v2i64 val, i32 1)
   //   i32 y = TRUNCATE(i64 x)
   //        -- becomes --

llvm/test/CodeGen/AMDGPU/mul_int24.ll

entry:
   store i32 %mul24hi, i32 addrspace(1)* %out
   ret void
 }
 
 define i64 @test_smul48_i64(i64 %lhs, i64 %rhs) {
 ; SI-LABEL: test_smul48_i64:
 ; SI:       ; %bb.0:
 ; SI-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
-; SI-NEXT:    v_lshlrev_b32_e32 v1, 8, v0
-; SI-NEXT:    v_lshlrev_b32_e32 v2, 8, v2
-; SI-NEXT:    v_ashr_i64 v[3:4], v[0:1], 40
-; SI-NEXT:    v_ashr_i64 v[1:2], v[1:2], 40
-; SI-NEXT:    v_mul_i32_i24_e32 v0, v3, v1
-; SI-NEXT:    v_mul_hi_i32_i24_e32 v1, v3, v1
+; SI-NEXT:    v_mul_i32_i24_e32 v3, v0, v2
+; SI-NEXT:    v_mul_hi_i32_i24_e32 v1, v0, v2
+; SI-NEXT:    v_mov_b32_e32 v0, v3
 ; SI-NEXT:    s_setpc_b64 s[30:31]
 ;
 ; VI-LABEL: test_smul48_i64:
 ; VI:       ; %bb.0:
 ; VI-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
-; VI-NEXT:    v_lshlrev_b32_e32 v1, 8, v0
-; VI-NEXT:    v_ashrrev_i64 v[3:4], 40, v[0:1]
-; VI-NEXT:    v_lshlrev_b32_e32 v1, 8, v2
-; VI-NEXT:    v_ashrrev_i64 v[1:2], 40, v[0:1]
-; VI-NEXT:    v_mul_i32_i24_e32 v0, v3, v1
-; VI-NEXT:    v_mul_hi_i32_i24_e32 v1, v3, v1
+; VI-NEXT:    v_mul_i32_i24_e32 v3, v0, v2
+; VI-NEXT:    v_mul_hi_i32_i24_e32 v1, v0, v2
+; VI-NEXT:    v_mov_b32_e32 v0, v3
 ; VI-NEXT:    s_setpc_b64 s[30:31]
 ;
 ; GFX9-LABEL: test_smul48_i64:
 ; GFX9:       ; %bb.0:
 ; GFX9-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
-; GFX9-NEXT:    v_lshlrev_b32_e32 v1, 8, v0
-; GFX9-NEXT:    v_ashrrev_i64 v[3:4], 40, v[0:1]
-; GFX9-NEXT:    v_lshlrev_b32_e32 v1, 8, v2
-; GFX9-NEXT:    v_ashrrev_i64 v[1:2], 40, v[0:1]
-; GFX9-NEXT:    v_mul_i32_i24_e32 v0, v3, v1
-; GFX9-NEXT:    v_mul_hi_i32_i24_e32 v1, v3, v1
+; GFX9-NEXT:    v_mul_i32_i24_e32 v3, v0, v2
+; GFX9-NEXT:    v_mul_hi_i32_i24_e32 v1, v0, v2
+; GFX9-NEXT:    v_mov_b32_e32 v0, v3
 ; GFX9-NEXT:    s_setpc_b64 s[30:31]
 ;
 ; EG-LABEL: test_smul48_i64:
 ; EG:       ; %bb.0:
 ; EG-NEXT:    CF_END
 ; EG-NEXT:    PAD
 ;
 ; CM-LABEL: test_smul48_i64:

llvm/test/CodeGen/X86/pmulh.ll

 ; AVX2-NEXT:    vextracti128 $1, %ymm0, %xmm2
 ; AVX2-NEXT:    vpackssdw %xmm2, %xmm0, %xmm0
 ; AVX2-NEXT:    vpmulhw %xmm1, %xmm0, %xmm0
 ; AVX2-NEXT:    vzeroupper
 ; AVX2-NEXT:    retq
 ;
 ; AVX512-LABEL: sextinreg_mulhw_v8i16:
 ; AVX512:       # %bb.0:
-; AVX512-NEXT:    vpslld $24, %ymm0, %ymm0
-; AVX512-NEXT:    vpsrad $24, %ymm0, %ymm0
-; AVX512-NEXT:    vpslld $25, %ymm1, %ymm1
-; AVX512-NEXT:    vpsrad $25, %ymm1, %ymm1
+; AVX512-NEXT:    # kill: def $ymm1 killed $ymm1 def $zmm1
+; AVX512-NEXT:    # kill: def $ymm0 killed $ymm0 def $zmm0
 ; AVX512-NEXT:    vpmovdw %zmm1, %ymm1
+; AVX512-NEXT:    vpsllw $9, %xmm1, %xmm1
+; AVX512-NEXT:    vpsraw $9, %xmm1, %xmm1
 ; AVX512-NEXT:    vpmovdw %zmm0, %ymm0
+; AVX512-NEXT:    vpsllw $8, %xmm0, %xmm0
+; AVX512-NEXT:    vpsraw $8, %xmm0, %xmm0
 ; AVX512-NEXT:    vpmulhw %xmm1, %xmm0, %xmm0
 ; AVX512-NEXT:    vzeroupper
 ; AVX512-NEXT:    retq
   %a1 = shl <8 x i32> %a, <i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24>
   %b1 = shl <8 x i32> %b, <i32 25, i32 25, i32 25, i32 25, i32 25, i32 25, i32 25, i32 25>
   %a2 = ashr <8 x i32> %a1, <i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24>
   %b2 = ashr <8 x i32> %b1, <i32 25, i32 25, i32 25, i32 25, i32 25, i32 25, i32 25, i32 25>
   %c = mul <8 x i32> %a2, %b2

llvm/test/CodeGen/X86/sar_fold64.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown | FileCheck %s --check-prefixes=CHECK,SSE
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+avx  | FileCheck %s --check-prefixes=CHECK,AVX1
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+avx2 | FileCheck %s --check-prefixes=CHECK,AVX2
 
 define i32 @shl48sar47(i64 %a) #0 {
 ; CHECK-LABEL: shl48sar47:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movswq %di, %rax
+; CHECK-NEXT:    movswl %di, %eax
 ; CHECK-NEXT:    addl %eax, %eax
-; CHECK-NEXT:    # kill: def $eax killed $eax killed $rax
 ; CHECK-NEXT:    retq
   %1 = shl i64 %a, 48
   %2 = ashr exact i64 %1, 47
   %3 = trunc i64 %2 to i32
   ret i32 %3
 }
 
 define i32 @shl48sar49(i64 %a) #0 {
 ; CHECK-LABEL: shl48sar49:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    movswq %di, %rax
 ; CHECK-NEXT:    shrq %rax
 ; CHECK-NEXT:    # kill: def $eax killed $eax killed $rax
 ; CHECK-NEXT:    retq
   %1 = shl i64 %a, 48
   %2 = ashr exact i64 %1, 49
   %3 = trunc i64 %2 to i32
   ret i32 %3
 }
 
 define i32 @shl56sar55(i64 %a) #0 {
 ; CHECK-LABEL: shl56sar55:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movsbq %dil, %rax
+; CHECK-NEXT:    movsbl %dil, %eax
 ; CHECK-NEXT:    addl %eax, %eax
-; CHECK-NEXT:    # kill: def $eax killed $eax killed $rax
 ; CHECK-NEXT:    retq
   %1 = shl i64 %a, 56
   %2 = ashr exact i64 %1, 55
   %3 = trunc i64 %2 to i32
   ret i32 %3
 }
 
 define i32 @shl56sar57(i64 %a) #0 {