[DAGCombiner] Add support for mulhi const folding in DAGCombiner

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5086

APInt::extractBits ?

5092

APInt::extractBits ?

[RKSimon]

Test cases would be useful as well of course

[foad]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5084

You don't need "OrTrunc" here and below.

[dstuttard]

Thanks for reviews.
Made suggested changes.

I'll add some test cases as well - but I'm off for a few days. I'll do it on return.

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

4467

Use DL variable that already exists

4519

Use DL variable that already exists

[dstuttard]

Updating for review comments.

Tests still to be added...

[RKSimon]

In D103323#2823907, @dstuttard wrote:

Updating for review comments.

Tests still to be added...

Do you have any in mind? I can add SSE2 vector examples if that would help? But scalars are trickier

[dstuttard]

Add test

[foad]

You can simplify the test case to:

define amdgpu_cs i64 @main(i64 %arg) {
entry:
  %d = udiv i64 %arg, 100000
  ret i64 %d
}

and it still shows the effect. Surely there are already some tests for i64 divide-by-constant that you could tweak, rather than adding a whole new file.

llvm/test/CodeGen/AMDGPU/dagcombine-mulhs-const.ll

5 ↗

(On Diff #352960)

Obviously folding the mul_hi is good, but the s_add that you check for looks like this:

s_mov_b32 s0, 0x346d900
...
s_add_u32 s0, 0x4237, s0

so it should also be folded to a constant!

[RKSimon]

@dstuttard Please can you rebase? rGcc38f8939da4aec85e7d0ef4de412e30d4de5a14 should give you vector coverage

[dstuttard]

Updated existing test based on feedback

Rebase and adjust X86 test that now folds

[dstuttard]

In D103323#2856828, @RKSimon wrote:

@dstuttard Please can you rebase? rGcc38f8939da4aec85e7d0ef4de412e30d4de5a14 should give you vector coverage

Thanks - have rebased and updated the test

[dstuttard]

llvm/test/CodeGen/AMDGPU/dagcombine-mulhs-const.ll

5 ↗

(On Diff #352960)

Yes - that could be another one to do - then fix up this test (or not worry about it at all given that there's now an X86 test that tests this, thanks to Simon)

[RKSimon]

Cheers!

llvm/test/CodeGen/AMDGPU/udiv.ll

206

Possibly pre-commit this to show current codegen? Do we need a GFX1030-NOT v_mul_hi_u32 check of some kind?

212

Fix missing newline

[dstuttard]

Pre-committed test and rebased on top

See D105424

[dstuttard]

llvm/test/CodeGen/AMDGPU/udiv.ll

206

Good idea. See latest change.

[RKSimon]

LGTM - cheers!

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitMULHS(SDNode *N) {
   if (VT.isVector()) {
     // fold (mulhs x, 0) -> 0
     // do not return N0/N1, because undef node may exist.
     if (ISD::isConstantSplatVectorAllZeros(N0.getNode()) ||
         ISD::isConstantSplatVectorAllZeros(N1.getNode()))
       return DAG.getConstant(0, DL, VT);
   }
 
+  // fold (mulhs c1, c2)
+  if (SDValue C = DAG.FoldConstantArithmetic(ISD::MULHS, DL, VT, {N0, N1}))
+    return C;
+
   // fold (mulhs x, 0) -> 0
   if (isNullConstant(N1))
     return N1;
   // fold (mulhs x, 1) -> (sra x, size(x)-1)
   if (isOneConstant(N1))
     return DAG.getNode(ISD::SRA, DL, N0.getValueType(), N0,
                        DAG.getConstant(N0.getScalarValueSizeInBits() - 1, DL,
                                        getShiftAmountTy(N0.getValueType())));
SDValue DAGCombiner::visitMULHU(SDNode *N) {
   if (VT.isVector()) {
     // fold (mulhu x, 0) -> 0
     // do not return N0/N1, because undef node may exist.
     if (ISD::isConstantSplatVectorAllZeros(N0.getNode()) ||
         ISD::isConstantSplatVectorAllZeros(N1.getNode()))
       return DAG.getConstant(0, DL, VT);
   }
 
+  // fold (mulhu c1, c2)
+  if (SDValue C = DAG.FoldConstantArithmetic(ISD::MULHU, DL, VT, {N0, N1}))
+    return C;
+
   // fold (mulhu x, 0) -> 0
   if (isNullConstant(N1))
     return N1;
   // fold (mulhu x, 1) -> 0
   if (isOneConstant(N1))
     return DAG.getConstant(0, DL, N0.getValueType());
   // fold (mulhu x, undef) -> 0
   if (N0.isUndef() || N1.isUndef())

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

static llvm::Optional<APInt> FoldValue(unsigned Opcode, const APInt &C1,
   case ISD::SDIV:
     if (!C2.getBoolValue())
       break;
     return C1.sdiv(C2);
   case ISD::SREM:
     if (!C2.getBoolValue())
       break;
     return C1.srem(C2);
+  case ISD::MULHS: {
+    unsigned FullWidth = C1.getBitWidth() * 2;
+    APInt C1Ext = C1.sext(FullWidth);
+    APInt C2Ext = C2.sext(FullWidth);
+    return (C1Ext * C2Ext).extractBits(C1.getBitWidth(), C1.getBitWidth());
+  }
+  case ISD::MULHU: {
+    unsigned FullWidth = C1.getBitWidth() * 2;
+    APInt C1Ext = C1.zext(FullWidth);
+    APInt C2Ext = C2.zext(FullWidth);
+    return (C1Ext * C2Ext).extractBits(C1.getBitWidth(), C1.getBitWidth());
+  }
   }
   return llvm::None;
 }
 
 SDValue SelectionDAG::FoldSymbolOffset(unsigned Opcode, EVT VT,
                                        const GlobalAddressSDNode *GA,
                                        const SDNode *N2) {
   if (GA->getOpcode() != ISD::GlobalAddress)

llvm/test/CodeGen/AMDGPU/udiv.ll

bb:
   %tmp5 = sdiv i32 %tmp1, %tmp4
   %tmp6 = trunc i32 %tmp5 to i8
   store i8 %tmp6, i8 addrspace(1)* null, align 1
   ret void
 }
 
 define i64 @v_test_udiv64_mulhi_fold(i64 %arg) {
 ; GCN-LABEL: v_test_udiv64_mulhi_fold
-; GFX1030: s_mov_b32 [[VAL1:s[0-9]+]], 0xa9000000
-; GFX1030: s_brev_b32 [[VAL2:s[0-9]+]], 6
-; GFX1030: s_movk_i32 [[VAL3:s[0-9]+]], 0x500
-; GFX1030: s_mul_hi_u32 s7, [[VAL1]], [[VAL2]]
-; GFX1030: s_mov_b32 [[VAL4:s[0-9]+]], 0xa7c5
-; GFX1030: s_mul_hi_u32 s8, [[VAL1]], [[VAL3]]
-; GFX1030: s_mul_hi_u32 s5, [[VAL4]], [[VAL2]]
-; GFX1030: s_mul_hi_u32 s6, [[VAL4]], [[VAL3]]
-; GFX1030: v_add_co_u32 v{{[0-9]+}}, s{{[0-9]+}}, s{{[0-9]+}}, s{{[0-9]+}}
+; GFX1030: s_add_u32 [[VAL:s[0-9]+]], 0x4237, s{{[0-9]+}}
+; GFX1030-NOT: s_mul_hi_u32
+; GFX1030: v_add_co_u32 v{{[0-9]+}}, [[VAL]], 0xa9000000, [[VAL]]
   %d = udiv i64 %arg, 100000
   ret i64 %d
 }

llvm/test/CodeGen/X86/pmulh.ll

; AVX512-NEXT:    retq
   %c = mul <8 x i64> %a1, %b1
   %d = ashr <8 x i64> %c, <i64 16, i64 16, i64 16, i64 16, i64 16, i64 16, i64 16, i64 16>
   ret <8 x i64> %d
 }
 
 define <8 x i16> @sse2_pmulh_w_const(<8 x i16> %a0, <8 x i16> %a1) {
 ; SSE-LABEL: sse2_pmulh_w_const:
 ; SSE:       # %bb.0:
-; SSE-NEXT:    movdqa {{.*#+}} xmm0 = [65535,65534,65533,65532,65531,65530,65529,0]
-; SSE-NEXT:    pmulhw {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0
+; SSE-NEXT:    movaps {{.*#+}} xmm0 = [0,65535,65535,65535,65535,65535,65535,0]
 ; SSE-NEXT:    retq
 ;
 ; AVX-LABEL: sse2_pmulh_w_const:
 ; AVX:       # %bb.0:
-; AVX-NEXT:    vmovdqa {{.*#+}} xmm0 = [65535,65534,65533,65532,65531,65530,65529,0]
-; AVX-NEXT:    vpmulhw {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
+; AVX-NEXT:    vmovaps {{.*#+}} xmm0 = [0,65535,65535,65535,65535,65535,65535,0]
 ; AVX-NEXT:    retq
   %res = call <8 x i16> @llvm.x86.sse2.pmulh.w(<8 x i16> <i16 -1, i16 -2, i16 -3, i16 -4, i16 -5, i16 -6, i16 -7, i16 0>, <8 x i16> <i16 0, i16 1, i16 2, i16 3, i16 4, i16 5, i16 6, i16 7>)
   ret <8 x i16> %res
 }
 declare <8 x i16> @llvm.x86.sse2.pmulh.w(<8 x i16>, <8 x i16>)
 
 define <8 x i16> @sse2_pmulhu_w_const(<8 x i16> %a0, <8 x i16> %a1) {
 ; SSE-LABEL: sse2_pmulhu_w_const:
 ; SSE:       # %bb.0:
-; SSE-NEXT:    movdqa {{.*#+}} xmm0 = [65535,65534,65533,65532,65531,65530,65529,0]
-; SSE-NEXT:    pmulhuw {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0
+; SSE-NEXT:    movaps {{.*#+}} xmm0 = [0,0,1,2,3,4,5,0]
 ; SSE-NEXT:    retq
 ;
 ; AVX-LABEL: sse2_pmulhu_w_const:
 ; AVX:       # %bb.0:
-; AVX-NEXT:    vmovdqa {{.*#+}} xmm0 = [65535,65534,65533,65532,65531,65530,65529,0]
-; AVX-NEXT:    vpmulhuw {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
+; AVX-NEXT:    vmovaps {{.*#+}} xmm0 = [0,0,1,2,3,4,5,0]
 ; AVX-NEXT:    retq
   %res = call <8 x i16> @llvm.x86.sse2.pmulhu.w(<8 x i16> <i16 -1, i16 -2, i16 -3, i16 -4, i16 -5, i16 -6, i16 -7, i16 0>, <8 x i16> <i16 0, i16 1, i16 2, i16 3, i16 4, i16 5, i16 6, i16 7>)
   ret <8 x i16> %res
 }
 declare <8 x i16> @llvm.x86.sse2.pmulhu.w(<8 x i16>, <8 x i16>)