[DAGCombiner] Add support for mulhi const folding in DAGCombiner

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5088

APInt::extractBits ?

5094

APInt::extractBits ?

[RKSimon]

Test cases would be useful as well of course

[foad]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5086

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

4456

Use DL variable that already exists

4508

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 ↗

(On Diff #356448)

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

210 ↗

(On Diff #356448)

Fix missing newline

[dstuttard]

Pre-committed test and rebased on top

See D105424

[dstuttard]

llvm/test/CodeGen/AMDGPU/udiv.ll

206 ↗

(On Diff #356448)

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)