[DAG] Fold (mul(abs(x),abs(x))) -> (mul(x,x)) (PR39476)

[RKSimon]

Summary

This patch adds support for discarding integer absolutes from self-multiplications.

Alive2: http://volta.cs.utah.edu:8080/z/rwcc8W

As part of this I've pulled out various absolute pattern matching code snippets we have and moved them into a MatchABS helper, which returns the source operand on success. We can add other ABS patterns here as required (e.g. PR45691).

NOTE: We shouldn't move the select(0>X,0-X,X) pattern as this is always canonicalized to a branchless pattern in visitVSELECT even if ISD::ABS isn't legal.

[lebedev.ri]

Echoing what @craig.topper asked, do we somehow end up with that pattern after instcombine had a chance to do this fold?

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

5059

Precommit NFC cleanup?

[lebedev.ri]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

3527–3530

This LG

[RKSimon]

In D79304#2016813, @lebedev.ri wrote:

Echoing what @craig.topper asked, do we somehow end up with that pattern after instcombine had a chance to do this fold?

Where did he ask that? https://godbolt.org/z/4vIPwf (from PR39476) shows that this IR currently comes through to the backend.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

5059

np - if people are happy with the patch I'll pull this out as a NFC precommit

[lebedev.ri]

In D79304#2016839, @RKSimon wrote:

In D79304#2016813, @lebedev.ri wrote:

Echoing what @craig.topper asked, do we somehow end up with that pattern after instcombine had a chance to do this fold?

Where did he ask that?

https://lists.llvm.org/pipermail/llvm-commits/Week-of-Mon-20200427/776590.html

https://godbolt.org/z/4vIPwf (from PR39476) shows that this IR currently comes through to the backend.

[lebedev.ri]

Now that we handle it in instcombine, do we still want dagcombine fold?

[RKSimon]

I'm going to abandon this (and I will remove the test cases from combine-mul.ll) - if we find a case where we need this in DAG then its trivial to reimplement.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue MatchRotatePosNeg(SDValue Shifted, SDValue Pos, SDValue Neg,
                               const SDLoc &DL);
     SDValue MatchFunnelPosNeg(SDValue N0, SDValue N1, SDValue Pos, SDValue Neg,
                               SDValue InnerPos, SDValue InnerNeg,
                               unsigned PosOpcode, unsigned NegOpcode,
                               const SDLoc &DL);
     SDValue MatchRotate(SDValue LHS, SDValue RHS, const SDLoc &DL);
     SDValue MatchLoadCombine(SDNode *N);
     SDValue MatchStoreCombine(StoreSDNode *N);
+    SDValue MatchABS(SDNode *N);
     SDValue ReduceLoadWidth(SDNode *N);
     SDValue ReduceLoadOpStoreWidth(SDNode *N);
     SDValue splitMergedValStore(StoreSDNode *ST);
     SDValue TransformFPLoadStorePair(SDNode *N);
     SDValue convertBuildVecZextToZext(SDNode *N);
     SDValue reduceBuildVecExtToExtBuildVec(SDNode *N);
     SDValue reduceBuildVecTruncToBitCast(SDNode *N);
     SDValue reduceBuildVecToShuffle(SDNode *N);
if (N1.getOpcode() == ISD::ZERO_EXTEND &&
       N1.getOperand(0).getScalarValueSizeInBits() == 1 &&
       TLI.getBooleanContents(VT) ==
           TargetLowering::ZeroOrNegativeOneBooleanContent) {
     SDValue SExt = DAG.getNode(ISD::SIGN_EXTEND, DL, VT, N1.getOperand(0));
     return DAG.getNode(ISD::ADD, DL, VT, N0, SExt);
   }
 
   // fold Y = sra (X, size(X)-1); sub (xor (X, Y), Y) -> (abs X)
-  if (TLI.isOperationLegalOrCustom(ISD::ABS, VT)) {
-    if (N0.getOpcode() == ISD::XOR && N1.getOpcode() == ISD::SRA) {
-      SDValue X0 = N0.getOperand(0), X1 = N0.getOperand(1);
-      SDValue S0 = N1.getOperand(0);
-      if ((X0 == S0 && X1 == N1) || (X0 == N1 && X1 == S0)) {
-        unsigned OpSizeInBits = VT.getScalarSizeInBits();
-        if (ConstantSDNode *C = isConstOrConstSplat(N1.getOperand(1)))
-          if (C->getAPIntValue() == (OpSizeInBits - 1))
-            return DAG.getNode(ISD::ABS, SDLoc(N), VT, S0);
-      }
-    }
-  }
+  if (TLI.isOperationLegalOrCustom(ISD::ABS, VT))
+    if (SDValue AbsSrc = MatchABS(N))
+      return DAG.getNode(ISD::ABS, DL, VT, AbsSrc);
 
   // If the relocation model supports it, consider symbol offsets.
   if (GlobalAddressSDNode *GA = dyn_cast<GlobalAddressSDNode>(N0))
     if (!LegalOperations && TLI.isOffsetFoldingLegal(GA)) {
       // fold (sub Sym, c) -> Sym-c
       if (N1C && GA->getOpcode() == ISD::GlobalAddress)
         return DAG.getGlobalAddress(GA->getGlobal(), SDLoc(N1C), VT,
                                     GA->getOffset() -
if (N1IsConst && !N1IsOpaqueConst && (-ConstValue1).isPowerOf2()) {
     // single-use add), we should put the negate there.
     return DAG.getNode(ISD::SUB, DL, VT,
                        DAG.getConstant(0, DL, VT),
                        DAG.getNode(ISD::SHL, DL, VT, N0,
                             DAG.getConstant(Log2Val, DL,
                                       getShiftAmountTy(N0.getValueType()))));
   }
 
+  // fold mul(abs(x),abs(x)) -> mul(x,x)
+  if (N0 == N1)
+    if (SDValue AbsSrc = MatchABS(N0.getNode()))
+      return DAG.getNode(ISD::MUL, SDLoc(N), VT, AbsSrc, AbsSrc);
+
   // Try to transform multiply-by-(power-of-2 +/- 1) into shift and add/sub.
   // mul x, (2^N + 1) --> add (shl x, N), x
   // mul x, (2^N - 1) --> sub (shl x, N), x
   // Examples: x * 33 --> (x << 5) + x
   //           x * 15 --> (x << 4) - x
   //           x * -33 --> -((x << 5) + x)
   //           x * -15 --> -((x << 4) - x) ; this reduces --> x - (x << 4)
   if (N1IsConst && TLI.decomposeMulByConstant(*DAG.getContext(), VT, N1)) {
for (auto *Load : Loads) {
       CombineTo(Load, NewLoad, NewLoad.getValue(1));
     }
     DAG.ReplaceAllUsesWith(N, N->getOperand(0).getNode());
     return true;
   }
   return false;
 }
 
+// Attempts to match an integer absolute pattern and returns the source operand
+// on success.
+SDValue DAGCombiner::MatchABS(SDNode *N) {
+  if (N->getOpcode() == ISD::ABS)
+    return N->getOperand(0);
+
+  // Y = sra (X, size(X)-1); sub (xor (X, Y), Y) -> (abs X)
+  if (N->getOpcode() == ISD::SUB) {
+    SDValue N0 = N->getOperand(0);
+    SDValue N1 = N->getOperand(1);
+    if (N0.getOpcode() == ISD::XOR && N1.getOpcode() == ISD::SRA) {
+      SDValue X0 = N0.getOperand(0), X1 = N0.getOperand(1);
+      SDValue S0 = N1.getOperand(0);
+      if ((X0 == S0 && X1 == N1) || (X0 == N1 && X1 == S0)) {
+        unsigned OpSizeInBits = N->getValueType(0).getScalarSizeInBits();
+        if (ConstantSDNode *C = isConstOrConstSplat(N1.getOperand(1)))
+          if (C->getAPIntValue() == (OpSizeInBits - 1))
+            return S0;
+      }
+    }
+  }
+
+  // Y = sra (X, size(X)-1); xor (add (X, Y), Y) -> (abs X)
+  if (N->getOpcode() == ISD::XOR) {
+    SDValue N0 = N->getOperand(0);
+    SDValue N1 = N->getOperand(1);
+    unsigned N0Opcode = N0.getOpcode();
+    SDValue A = N0Opcode == ISD::ADD ? N0 : N1;
+    SDValue S = N0Opcode == ISD::SRA ? N0 : N1;
+    if (A.getOpcode() == ISD::ADD && S.getOpcode() == ISD::SRA) {
+      SDValue A0 = A.getOperand(0), A1 = A.getOperand(1);
+      SDValue S0 = S.getOperand(0);
+      if ((A0 == S && A1 == S0) || (A1 == S && A0 == S0)) {
+        unsigned OpSizeInBits = N->getValueType(0).getScalarSizeInBits();
+        if (ConstantSDNode *C = isConstOrConstSplat(S.getOperand(1)))
+          if (C->getAPIntValue() == (OpSizeInBits - 1))
+            return S0;
+      }
+    }
+  }
+
+  return SDValue();
+}
+
 // Unfold
 //    x &  (-1 'logical shift' y)
 // To
 //    (x 'opposite logical shift' y) 'logical shift' y
 // if it is better for performance.
 SDValue DAGCombiner::unfoldExtremeBitClearingToShifts(SDNode *N) {
   assert(N->getOpcode() == ISD::AND);
 
if (XorC && ShiftC) {
           SDValue Not = DAG.getNOT(DL, N0.getOperand(0), VT);
           return DAG.getNode(N0Opcode, DL, VT, Not, N0.getOperand(1));
         }
       }
     }
   }
 
   // fold Y = sra (X, size(X)-1); xor (add (X, Y), Y) -> (abs X)
-  if (TLI.isOperationLegalOrCustom(ISD::ABS, VT)) {
-    SDValue A = N0Opcode == ISD::ADD ? N0 : N1;
-    SDValue S = N0Opcode == ISD::SRA ? N0 : N1;
-    if (A.getOpcode() == ISD::ADD && S.getOpcode() == ISD::SRA) {
-      SDValue A0 = A.getOperand(0), A1 = A.getOperand(1);
-      SDValue S0 = S.getOperand(0);
-      if ((A0 == S && A1 == S0) || (A1 == S && A0 == S0)) {
-        unsigned OpSizeInBits = VT.getScalarSizeInBits();
-        if (ConstantSDNode *C = isConstOrConstSplat(S.getOperand(1)))
-          if (C->getAPIntValue() == (OpSizeInBits - 1))
-            return DAG.getNode(ISD::ABS, DL, VT, S0);
-      }
-    }
-  }
+  if (TLI.isOperationLegalOrCustom(ISD::ABS, VT))
+    if (SDValue AbsSrc = MatchABS(N))
+      return DAG.getNode(ISD::ABS, DL, VT, AbsSrc);
 
   // fold (xor x, x) -> 0
   if (N0 == N1)
     return tryFoldToZero(DL, TLI, VT, DAG, LegalOperations);
 
   // fold (xor (shl 1, x), -1) -> (rotl ~1, x)
   // Here is a concrete example of this equivalence:
   // i16   x ==  14

llvm/test/CodeGen/X86/combine-mul.ll

 ; AVX-NEXT:    vpmulld {{.*}}(%rip), %xmm0, %xmm0
 ; AVX-NEXT:    vpaddd {{.*}}(%rip), %xmm0, %xmm0
 ; AVX-NEXT:    retq
   %1 = add <4 x i32> %x, <i32 1, i32 2, i32 8, i32 16>
   %2 = mul <4 x i32> %1, <i32 4, i32 6, i32 2, i32 0>
   ret <4 x i32> %2
 }
 
-; TODO fold mul(abs(x),abs(x)) -> mul(x,x)
+; fold mul(abs(x),abs(x)) -> mul(x,x)
 
 define i31 @combine_mul_abs_i31(i31 %0) {
 ; SSE-LABEL: combine_mul_abs_i31:
 ; SSE:       # %bb.0:
-; SSE-NEXT:    addl %edi, %edi
-; SSE-NEXT:    sarl %edi
 ; SSE-NEXT:    movl %edi, %eax
-; SSE-NEXT:    negl %eax
-; SSE-NEXT:    cmovll %edi, %eax
-; SSE-NEXT:    imull %eax, %eax
+; SSE-NEXT:    imull %edi, %eax
 ; SSE-NEXT:    retq
 ;
 ; AVX-LABEL: combine_mul_abs_i31:
 ; AVX:       # %bb.0:
-; AVX-NEXT:    addl %edi, %edi
-; AVX-NEXT:    sarl %edi
 ; AVX-NEXT:    movl %edi, %eax
-; AVX-NEXT:    negl %eax
-; AVX-NEXT:    cmovll %edi, %eax
-; AVX-NEXT:    imull %eax, %eax
+; AVX-NEXT:    imull %edi, %eax
 ; AVX-NEXT:    retq
   %c = icmp slt i31 %0, 0
   %s = sub nsw i31 0, %0
   %r = select i1 %c, i31 %s, i31 %0
   %m = mul i31 %r, %r
   ret i31 %m
 }
 
 define i32 @combine_mul_abs_i32(i32 %0) {
 ; SSE-LABEL: combine_mul_abs_i32:
 ; SSE:       # %bb.0:
 ; SSE-NEXT:    movl %edi, %eax
-; SSE-NEXT:    negl %eax
-; SSE-NEXT:    cmovll %edi, %eax
-; SSE-NEXT:    imull %eax, %eax
+; SSE-NEXT:    imull %edi, %eax
 ; SSE-NEXT:    retq
 ;
 ; AVX-LABEL: combine_mul_abs_i32:
 ; AVX:       # %bb.0:
 ; AVX-NEXT:    movl %edi, %eax
-; AVX-NEXT:    negl %eax
-; AVX-NEXT:    cmovll %edi, %eax
-; AVX-NEXT:    imull %eax, %eax
+; AVX-NEXT:    imull %edi, %eax
 ; AVX-NEXT:    retq
   %c = icmp slt i32 %0, 0
   %s = sub nsw i32 0, %0
   %r = select i1 %c, i32 %s, i32 %0
   %m = mul i32 %r, %r
   ret i32 %m
 }
 
 define <4 x i32> @combine_mul_abs_v4i32(<4 x i32> %0) {
 ; SSE-LABEL: combine_mul_abs_v4i32:
 ; SSE:       # %bb.0:
-; SSE-NEXT:    pabsd %xmm0, %xmm0
 ; SSE-NEXT:    pmulld %xmm0, %xmm0
 ; SSE-NEXT:    retq
 ;
 ; AVX-LABEL: combine_mul_abs_v4i32:
 ; AVX:       # %bb.0:
-; AVX-NEXT:    vpabsd %xmm0, %xmm0
 ; AVX-NEXT:    vpmulld %xmm0, %xmm0, %xmm0
 ; AVX-NEXT:    retq
   %c = icmp slt <4 x i32> %0, zeroinitializer
   %s = sub nsw <4 x i32> zeroinitializer, %0
   %r = select <4 x i1> %c, <4 x i32> %s, <4 x i32> %0
   %m = mul <4 x i32> %r, %r
   ret <4 x i32> %m
 }