[SDAG] enable binop identity constant folds for sub

[LuoYuanke]

Summary

This patch extract the sub folding from D119654 and leave only add
folding in that patch.

[RKSimon]

LGTM

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (ConstantFPSDNode *C = isConstOrConstSplatFP(V)) {
         return C->isZero() && C->isNegative();
       case ISD::FSUB: // X - 0.0 --> X
         return C->isZero() && !C->isNegative();
       case ISD::FMUL: // X * 1.0 --> X
       case ISD::FDIV: // X / 1.0 --> X
         return C->isExactlyValue(1.0);
       }
     }
+    if (ConstantSDNode *C = isConstOrConstSplat(V)) {
+      switch (Opcode) {
+      case ISD::SUB: // X - 0 --> X
+        return C->isZero();
+      }
+    }
     return false;
   };
 
   // This transform increases uses of N0, so freeze it to be safe.
   // binop N0, (vselect Cond, IDC, FVal) --> vselect Cond, N0, (binop N0, FVal)
   if (isIdentityConstantForOpcode(Opcode, TVal)) {
     SDValue F0 = DAG.getFreeze(N0);
     SDValue NewBO = DAG.getNode(Opcode, SDLoc(N), VT, F0, FVal, N->getFlags());

 }
 
 SDValue DAGCombiner::visitSUB(SDNode *N) {
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
   EVT VT = N0.getValueType();
   SDLoc DL(N);
 
+  auto PeekThroughFreeze = [](SDValue N) {
+    if (N->getOpcode() == ISD::FREEZE && N.hasOneUse())
+      return N->getOperand(0);
+    return N;
+  };
+
   // fold (sub x, x) -> 0
   // FIXME: Refactor this and xor and other similar operations together.
-  if (N0 == N1)
+  if (PeekThroughFreeze(N0) == PeekThroughFreeze(N1))
     return tryFoldToZero(DL, TLI, VT, DAG, LegalOperations);
 
   // fold (sub c1, c2) -> c3
   if (SDValue C = DAG.FoldConstantArithmetic(ISD::SUB, DL, VT, {N0, N1}))
     return C;
 
   // fold vector ops
   if (VT.isVector()) {

llvm/test/CodeGen/X86/vector-bo-select.ll

 ; AVX512F-NEXT:    vpsubd %xmm0, %xmm1, %xmm0
 ; AVX512F-NEXT:    vzeroupper
 ; AVX512F-NEXT:    retq
 ;
 ; AVX512VL-LABEL: sub_v4i32:
 ; AVX512VL:       # %bb.0:
 ; AVX512VL-NEXT:    vpslld $31, %xmm0, %xmm0
 ; AVX512VL-NEXT:    vptestmd %xmm0, %xmm0, %k1
-; AVX512VL-NEXT:    vmovdqa32 %xmm2, %xmm0 {%k1} {z}
-; AVX512VL-NEXT:    vpsubd %xmm0, %xmm1, %xmm0
+; AVX512VL-NEXT:    vpsubd %xmm2, %xmm1, %xmm1 {%k1}
+; AVX512VL-NEXT:    vmovdqa %xmm1, %xmm0
 ; AVX512VL-NEXT:    retq
   %s = select <4 x i1> %b, <4 x i32> %y, <4 x i32> zeroinitializer
   %r = sub <4 x i32> %x, %s
   ret <4 x i32> %r
 }
 
 ; negative test - sub is not commutative; there is no identity constant for operand 0
 

 ; AVX2-NEXT:    vpsubd %ymm0, %ymm1, %ymm0
 ; AVX2-NEXT:    vpsubd %ymm4, %ymm2, %ymm1
 ; AVX2-NEXT:    retq
 ;
 ; AVX512-LABEL: sub_v16i32_swap:
 ; AVX512:       # %bb.0:
 ; AVX512-NEXT:    vpmovsxbd %xmm0, %zmm0
 ; AVX512-NEXT:    vpslld $31, %zmm0, %zmm0
-; AVX512-NEXT:    vptestnmd %zmm0, %zmm0, %k1
-; AVX512-NEXT:    vmovdqa32 %zmm2, %zmm0 {%k1} {z}
-; AVX512-NEXT:    vpsubd %zmm0, %zmm1, %zmm0
+; AVX512-NEXT:    vptestmd %zmm0, %zmm0, %k1
+; AVX512-NEXT:    vpsubd %zmm2, %zmm1, %zmm0
+; AVX512-NEXT:    vmovdqa32 %zmm1, %zmm0 {%k1}
 ; AVX512-NEXT:    retq
   %s = select <16 x i1> %b, <16 x i32> zeroinitializer, <16 x i32> %y
   %r = sub <16 x i32> %x, %s
   ret <16 x i32> %r
 }
 
 ; negative test - sub is not commutative; there is no identity constant for operand 0
 

 ; AVX512F-NEXT:    kmovw %edi, %k1
 ; AVX512F-NEXT:    vmovdqa32 %zmm1, %zmm1 {%k1} {z}
 ; AVX512F-NEXT:    vpsubd %ymm1, %ymm0, %ymm0
 ; AVX512F-NEXT:    retq
 ;
 ; AVX512VL-LABEL: sub_v8i32_cast_cond:
 ; AVX512VL:       # %bb.0:
 ; AVX512VL-NEXT:    kmovw %edi, %k1
-; AVX512VL-NEXT:    vmovdqa32 %ymm1, %ymm1 {%k1} {z}
-; AVX512VL-NEXT:    vpsubd %ymm1, %ymm0, %ymm0
+; AVX512VL-NEXT:    vpsubd %ymm1, %ymm0, %ymm0 {%k1}
 ; AVX512VL-NEXT:    retq
   %b = bitcast i8 %pb to <8 x i1>
   %s = select <8 x i1> %b, <8 x i32> %y, <8 x i32> zeroinitializer
   %r = sub <8 x i32> %x, %s
   ret <8 x i32> %r
 }
 
 define <8 x i64> @sub_v8i64_cast_cond(i8 noundef zeroext %pb, <8 x i64> noundef %x, <8 x i64> noundef %y) {

 ; AVX2-NEXT:    vpand %ymm2, %ymm4, %ymm2
 ; AVX2-NEXT:    vpsubq %ymm2, %ymm0, %ymm0
 ; AVX2-NEXT:    vpsubq %ymm3, %ymm1, %ymm1
 ; AVX2-NEXT:    retq
 ;
 ; AVX512-LABEL: sub_v8i64_cast_cond:
 ; AVX512:       # %bb.0:
 ; AVX512-NEXT:    kmovw %edi, %k1
-; AVX512-NEXT:    vmovdqa64 %zmm1, %zmm1 {%k1} {z}
-; AVX512-NEXT:    vpsubq %zmm1, %zmm0, %zmm0
+; AVX512-NEXT:    vpsubq %zmm1, %zmm0, %zmm0 {%k1}
 ; AVX512-NEXT:    retq
   %b = bitcast i8 %pb to <8 x i1>
   %s = select <8 x i1> %b, <8 x i64> %y, <8 x i64> zeroinitializer
   %r = sub <8 x i64> %x, %s
   ret <8 x i64> %r
 }