[DAG][AArch64] Limit preferIncOfAddToSubOfNot until after legalization if the node has wrap flags

[dmgreen]

Summary

If the add node has wrap flags then they will be destroyed by converting to sub/not. The flags can be useful in converting to rhadd, for example, but that can happen late if the node types need to be legalized. This limits the preferIncOfAddToSubOfNot fold until after legalize DAG if the node have flags to allow more folding.

[samtebbs]

LGTM

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (N0.getOpcode() == ISD::ADD) {
         return DAG.getNode(ISD::SUB, DL, VT, A, Xor.getOperand(0));
     }
 
     // Look for:
     //   add (add x, y), 1
     // And if the target does not like this form then turn into:
     //   sub y, (xor x, -1)
     if (!TLI.preferIncOfAddToSubOfNot(VT) && N0.getOpcode() == ISD::ADD &&
-        N0.hasOneUse()) {
+        N0.hasOneUse() &&
+        // Limit this to after legalization if the add has wrap flags
+        (Level >= AfterLegalizeDAG || (!N->getFlags().hasNoUnsignedWrap() &&
+                                       !N->getFlags().hasNoSignedWrap()))) {
       SDValue Not = DAG.getNode(ISD::XOR, DL, VT, N0.getOperand(0),
                                 DAG.getAllOnesConstant(DL, VT));
       return DAG.getNode(ISD::SUB, DL, VT, N0.getOperand(1), Not);
     }
   }
 
   // (x - y) + -1  ->  add (xor y, -1), x
   if (N0.getOpcode() == ISD::SUB && N0.hasOneUse() &&
SDValue DAGCombiner::visitADDLikeCommutative(SDValue N0, SDValue N1,
   if (SDValue V = foldAddSubMasked1(true, N0, N1, DAG, DL))
     return V;
 
   // Look for:
   //   add (add x, 1), y
   // And if the target does not like this form then turn into:
   //   sub y, (xor x, -1)
   if (!TLI.preferIncOfAddToSubOfNot(VT) && N0.getOpcode() == ISD::ADD &&
-      N0.hasOneUse() && isOneOrOneSplat(N0.getOperand(1))) {
+      N0.hasOneUse() && isOneOrOneSplat(N0.getOperand(1)) &&
+      // Limit this to after legalization if the add has wrap flags
+      (Level >= AfterLegalizeDAG || (!N0->getFlags().hasNoUnsignedWrap() &&
+                                     !N0->getFlags().hasNoSignedWrap()))) {
     SDValue Not = DAG.getNode(ISD::XOR, DL, VT, N0.getOperand(0),
                               DAG.getAllOnesConstant(DL, VT));
     return DAG.getNode(ISD::SUB, DL, VT, N1, Not);
   }
 
   if (N0.getOpcode() == ISD::SUB && N0.hasOneUse()) {
     // Hoist one-use subtraction by non-opaque constant:
     //   (x - C) + y  ->  (x + y) - C

llvm/test/CodeGen/AArch64/arm64-vhadd.ll

; CHECK-NEXT:    ret
   ret <4 x i16> %resulti8
 }
 
 define <4 x i16> @rhadd8_sext_lsr(<4 x i8> %src1, <4 x i8> %src2) {
 ; CHECK-LABEL: rhadd8_sext_lsr:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    shl.4h v0, v0, #8
 ; CHECK-NEXT:    shl.4h v1, v1, #8
+; CHECK-NEXT:    movi.4h v2, #1
 ; CHECK-NEXT:    sshr.4h v0, v0, #8
-; CHECK-NEXT:    sshr.4h v1, v1, #8
-; CHECK-NEXT:    mvn.8b v0, v0
-; CHECK-NEXT:    sub.4h v0, v1, v0
+; CHECK-NEXT:    ssra.4h v0, v1, #8
+; CHECK-NEXT:    add.4h v0, v0, v2
 ; CHECK-NEXT:    ushr.4h v0, v0, #1
 ; CHECK-NEXT:    ret
   %zextsrc1 = sext <4 x i8> %src1 to <4 x i16>
   %zextsrc2 = sext <4 x i8> %src2 to <4 x i16>
   %add = add nsw <4 x i16> %zextsrc1, %zextsrc2
   %add2 = add nsw <4 x i16> %add, <i16 1, i16 1, i16 1, i16 1>
   %resulti8 = lshr <4 x i16> %add2, <i16 1, i16 1, i16 1, i16 1>
   ret <4 x i16> %resulti8

 
 define <2 x i16> @rhadd8x2_sext_asr(<2 x i8> %src1, <2 x i8> %src2) {
 ; CHECK-LABEL: rhadd8x2_sext_asr:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    shl.2s v0, v0, #24
 ; CHECK-NEXT:    shl.2s v1, v1, #24
 ; CHECK-NEXT:    sshr.2s v0, v0, #24
 ; CHECK-NEXT:    sshr.2s v1, v1, #24
-; CHECK-NEXT:    mvn.8b v0, v0
-; CHECK-NEXT:    sub.2s v0, v1, v0
-; CHECK-NEXT:    sshr.2s v0, v0, #1
+; CHECK-NEXT:    srhadd.2s v0, v0, v1
 ; CHECK-NEXT:    ret
   %zextsrc1 = sext <2 x i8> %src1 to <2 x i16>
   %zextsrc2 = sext <2 x i8> %src2 to <2 x i16>
   %add = add nsw <2 x i16> %zextsrc1, %zextsrc2
   %add2 = add nsw <2 x i16> %add, <i16 1, i16 1>
   %resulti8 = ashr <2 x i16> %add2, <i16 1, i16 1>
   ret <2 x i16> %resulti8
 }
 
 define <2 x i16> @rhadd8x2_zext_asr(<2 x i8> %src1, <2 x i8> %src2) {
 ; CHECK-LABEL: rhadd8x2_zext_asr:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi d2, #0x0000ff000000ff
 ; CHECK-NEXT:    and.8b v0, v0, v2
 ; CHECK-NEXT:    and.8b v1, v1, v2
-; CHECK-NEXT:    mvn.8b v0, v0
-; CHECK-NEXT:    sub.2s v0, v1, v0
-; CHECK-NEXT:    ushr.2s v0, v0, #1
+; CHECK-NEXT:    urhadd.2s v0, v0, v1
 ; CHECK-NEXT:    ret
   %zextsrc1 = zext <2 x i8> %src1 to <2 x i16>
   %zextsrc2 = zext <2 x i8> %src2 to <2 x i16>
   %add = add nuw nsw <2 x i16> %zextsrc1, %zextsrc2
   %add2 = add nuw nsw <2 x i16> %add, <i16 1, i16 1>
   %resulti8 = lshr <2 x i16> %add2, <i16 1, i16 1>
   ret <2 x i16> %resulti8
 }
 
 define <2 x i16> @rhadd8x2_sext_lsr(<2 x i8> %src1, <2 x i8> %src2) {
 ; CHECK-LABEL: rhadd8x2_sext_lsr:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    shl.2s v0, v0, #24
 ; CHECK-NEXT:    shl.2s v1, v1, #24
-; CHECK-NEXT:    movi d2, #0x00ffff0000ffff
+; CHECK-NEXT:    movi.2s v2, #1
 ; CHECK-NEXT:    sshr.2s v0, v0, #24
-; CHECK-NEXT:    sshr.2s v1, v1, #24
-; CHECK-NEXT:    mvn.8b v0, v0
-; CHECK-NEXT:    sub.2s v0, v1, v0
-; CHECK-NEXT:    and.8b v0, v0, v2
+; CHECK-NEXT:    ssra.2s v0, v1, #24
+; CHECK-NEXT:    movi d1, #0x00ffff0000ffff
+; CHECK-NEXT:    add.2s v0, v0, v2
+; CHECK-NEXT:    and.8b v0, v0, v1
 ; CHECK-NEXT:    ushr.2s v0, v0, #1
 ; CHECK-NEXT:    ret
   %zextsrc1 = sext <2 x i8> %src1 to <2 x i16>
   %zextsrc2 = sext <2 x i8> %src2 to <2 x i16>
   %add = add nsw <2 x i16> %zextsrc1, %zextsrc2
   %add2 = add nsw <2 x i16> %add, <i16 1, i16 1>
   %resulti8 = lshr <2 x i16> %add2, <i16 1, i16 1>
   ret <2 x i16> %resulti8
 }
 
 define <2 x i16> @rhadd8x2_zext_lsr(<2 x i8> %src1, <2 x i8> %src2) {
 ; CHECK-LABEL: rhadd8x2_zext_lsr:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi d2, #0x0000ff000000ff
 ; CHECK-NEXT:    and.8b v0, v0, v2
 ; CHECK-NEXT:    and.8b v1, v1, v2
-; CHECK-NEXT:    mvn.8b v0, v0
-; CHECK-NEXT:    sub.2s v0, v1, v0
-; CHECK-NEXT:    ushr.2s v0, v0, #1
+; CHECK-NEXT:    urhadd.2s v0, v0, v1
 ; CHECK-NEXT:    ret
   %zextsrc1 = zext <2 x i8> %src1 to <2 x i16>
   %zextsrc2 = zext <2 x i8> %src2 to <2 x i16>
   %add = add nuw nsw <2 x i16> %zextsrc1, %zextsrc2
   %add2 = add nuw nsw <2 x i16> %add, <i16 1, i16 1>
   %resulti8 = lshr <2 x i16> %add2, <i16 1, i16 1>
   ret <2 x i16> %resulti8
 }

llvm/test/CodeGen/AArch64/sve-hadd.ll

entry:
   %add = add <vscale x 2 x i32> %s0s, shufflevector (<vscale x 2 x i32> insertelement (<vscale x 2 x i32> poison, i32 1, i32 0), <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer)
   %add2 = add <vscale x 2 x i32> %add, %s1s
   %s = lshr <vscale x 2 x i32> %add2, shufflevector (<vscale x 2 x i32> insertelement (<vscale x 2 x i32> poison, i32 1, i32 0), <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer)
   %result = trunc <vscale x 2 x i32> %s to <vscale x 2 x i16>
   ret <vscale x 2 x i16> %result
 }
 
 define <vscale x 2 x i16> @rhaddu_v2i16(<vscale x 2 x i16> %s0, <vscale x 2 x i16> %s1) {
-; CHECK-LABEL: rhaddu_v2i16:
-; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    mov z2.d, #-1 // =0xffffffffffffffff
-; CHECK-NEXT:    and z0.d, z0.d, #0xffff
-; CHECK-NEXT:    and z1.d, z1.d, #0xffff
-; CHECK-NEXT:    eor z0.d, z0.d, z2.d
-; CHECK-NEXT:    sub z0.d, z1.d, z0.d
-; CHECK-NEXT:    lsr z0.d, z0.d, #1
-; CHECK-NEXT:    ret
+; SVE-LABEL: rhaddu_v2i16:
+; SVE:       // %bb.0: // %entry
+; SVE-NEXT:    mov z2.d, #-1 // =0xffffffffffffffff
+; SVE-NEXT:    and z0.d, z0.d, #0xffff
+; SVE-NEXT:    and z1.d, z1.d, #0xffff
+; SVE-NEXT:    eor z0.d, z0.d, z2.d
+; SVE-NEXT:    sub z0.d, z1.d, z0.d
+; SVE-NEXT:    lsr z0.d, z0.d, #1
+; SVE-NEXT:    ret
+;
+; SVE2-LABEL: rhaddu_v2i16:
+; SVE2:       // %bb.0: // %entry
+; SVE2-NEXT:    ptrue p0.d
+; SVE2-NEXT:    and z0.d, z0.d, #0xffff
+; SVE2-NEXT:    and z1.d, z1.d, #0xffff
+; SVE2-NEXT:    urhadd z0.d, p0/m, z0.d, z1.d
+; SVE2-NEXT:    ret
 entry:
   %s0s = zext <vscale x 2 x i16> %s0 to <vscale x 2 x i32>
   %s1s = zext <vscale x 2 x i16> %s1 to <vscale x 2 x i32>
   %add = add nuw nsw <vscale x 2 x i32> %s0s, shufflevector (<vscale x 2 x i32> insertelement (<vscale x 2 x i32> poison, i32 1, i32 0), <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer)
   %add2 = add nuw nsw <vscale x 2 x i32> %add, %s1s
   %s = lshr <vscale x 2 x i32> %add2, shufflevector (<vscale x 2 x i32> insertelement (<vscale x 2 x i32> poison, i32 1, i32 0), <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer)
   %result = trunc <vscale x 2 x i32> %s to <vscale x 2 x i16>
   ret <vscale x 2 x i16> %result
entry:
   %add = add <vscale x 4 x i16> %s0s, shufflevector (<vscale x 4 x i16> insertelement (<vscale x 4 x i16> poison, i16 1, i32 0), <vscale x 4 x i16> poison, <vscale x 4 x i32> zeroinitializer)
   %add2 = add <vscale x 4 x i16> %add, %s1s
   %s = lshr <vscale x 4 x i16> %add2, shufflevector (<vscale x 4 x i16> insertelement (<vscale x 4 x i16> poison, i16 1, i32 0), <vscale x 4 x i16> poison, <vscale x 4 x i32> zeroinitializer)
   %result = trunc <vscale x 4 x i16> %s to <vscale x 4 x i8>
   ret <vscale x 4 x i8> %result
 }
 
 define <vscale x 4 x i8> @rhaddu_v4i8(<vscale x 4 x i8> %s0, <vscale x 4 x i8> %s1) {
-; CHECK-LABEL: rhaddu_v4i8:
-; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    mov z2.s, #-1 // =0xffffffffffffffff
-; CHECK-NEXT:    and z0.s, z0.s, #0xff
-; CHECK-NEXT:    and z1.s, z1.s, #0xff
-; CHECK-NEXT:    eor z0.d, z0.d, z2.d
-; CHECK-NEXT:    sub z0.s, z1.s, z0.s
-; CHECK-NEXT:    lsr z0.s, z0.s, #1
-; CHECK-NEXT:    ret
+; SVE-LABEL: rhaddu_v4i8:
+; SVE:       // %bb.0: // %entry
+; SVE-NEXT:    mov z2.s, #-1 // =0xffffffffffffffff
+; SVE-NEXT:    and z0.s, z0.s, #0xff
+; SVE-NEXT:    and z1.s, z1.s, #0xff
+; SVE-NEXT:    eor z0.d, z0.d, z2.d
+; SVE-NEXT:    sub z0.s, z1.s, z0.s
+; SVE-NEXT:    lsr z0.s, z0.s, #1
+; SVE-NEXT:    ret
+;
+; SVE2-LABEL: rhaddu_v4i8:
+; SVE2:       // %bb.0: // %entry
+; SVE2-NEXT:    ptrue p0.s
+; SVE2-NEXT:    and z0.s, z0.s, #0xff
+; SVE2-NEXT:    and z1.s, z1.s, #0xff
+; SVE2-NEXT:    urhadd z0.s, p0/m, z0.s, z1.s
+; SVE2-NEXT:    ret
 entry:
   %s0s = zext <vscale x 4 x i8> %s0 to <vscale x 4 x i16>
   %s1s = zext <vscale x 4 x i8> %s1 to <vscale x 4 x i16>
   %add = add nuw nsw <vscale x 4 x i16> %s0s, shufflevector (<vscale x 4 x i16> insertelement (<vscale x 4 x i16> poison, i16 1, i32 0), <vscale x 4 x i16> poison, <vscale x 4 x i32> zeroinitializer)
   %add2 = add nuw nsw <vscale x 4 x i16> %add, %s1s
   %s = lshr <vscale x 4 x i16> %add2, shufflevector (<vscale x 4 x i16> insertelement (<vscale x 4 x i16> poison, i16 1, i32 0), <vscale x 4 x i16> poison, <vscale x 4 x i32> zeroinitializer)
   %result = trunc <vscale x 4 x i16> %s to <vscale x 4 x i8>
   ret <vscale x 4 x i8> %result