[DAGCombine][AArch64][ARM] Combine abd(sub(x, y)) to abd if the sub is nsw

[dmgreen]

Summary

This implements the fold (abs (sub nsw x, y)) -> abds(x, y). Providing the sub is nsw this appears to be valid without the extensions that are usually used for abds. https://alive2.llvm.org/ce/z/XHVaB3. The equivalent abdu combine seems to not be valid.

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10168

Worth moving this inside combineABSToABD?

[dmgreen]

Moved into combineABSToABD.

[RKSimon]

LGTM - cheers

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10120

You can probably hoist this before the sext/zext test?

[dmgreen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10120

The idea is that if we have a sext/zext pair then is better to use them in the abd, creating the smaller type. We use the wrapping flags if that fails.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitSHLSAT(SDNode *N) {
   return SDValue();
 }
 
 // Given a ABS node, detect the following pattern:
 // (ABS (SUB (EXTEND a), (EXTEND b))).
 // Generates UABD/SABD instruction.
 static SDValue combineABSToABD(SDNode *N, SelectionDAG &DAG,
                                const TargetLowering &TLI) {
+  EVT VT = N->getValueType(0);
   SDValue AbsOp1 = N->getOperand(0);
   SDValue Op0, Op1;
 
   if (AbsOp1.getOpcode() != ISD::SUB)
     return SDValue();
 
   Op0 = AbsOp1.getOperand(0);
   Op1 = AbsOp1.getOperand(1);
 
   unsigned Opc0 = Op0.getOpcode();
   // Check if the operands of the sub are (zero|sign)-extended.
   if (Opc0 != Op1.getOpcode() ||
-      (Opc0 != ISD::ZERO_EXTEND && Opc0 != ISD::SIGN_EXTEND))
+      (Opc0 != ISD::ZERO_EXTEND && Opc0 != ISD::SIGN_EXTEND)) {
+    // fold (abs (sub nsw x, y)) -> abds(x, y)
+    if (AbsOp1->getFlags().hasNoSignedWrap() &&
+        TLI.isOperationLegalOrCustom(ISD::ABDS, VT))
+      return DAG.getNode(ISD::ABDS, SDLoc(N), VT, Op0, Op1);
     return SDValue();
+  }
 
-  EVT VT = N->getValueType(0);
   EVT VT1 = Op0.getOperand(0).getValueType();
   EVT VT2 = Op1.getOperand(0).getValueType();
   unsigned ABDOpcode = (Opc0 == ISD::SIGN_EXTEND) ? ISD::ABDS : ISD::ABDU;
 
   // fold abs(sext(x) - sext(y)) -> zext(abds(x, y))
   // fold abs(zext(x) - zext(y)) -> zext(abdu(x, y))
   // NOTE: Extensions must be equivalent.
   if (VT1 == VT2 && TLI.isOperationLegalOrCustom(ABDOpcode, VT1)) {
SDValue DAGCombiner::visitABS(SDNode *N) {
   if (SDValue ABD = combineABSToABD(N, DAG, TLI))
     return ABD;
 
   // fold (abs (sign_extend_inreg x)) -> (zero_extend (abs (truncate x)))
   // iff zero_extend/truncate are free.
   if (N0.getOpcode() == ISD::SIGN_EXTEND_INREG) {
     EVT ExtVT = cast<VTSDNode>(N0.getOperand(1))->getVT();
     if (TLI.isTruncateFree(VT, ExtVT) && TLI.isZExtFree(ExtVT, VT) &&
         TLI.isTypeDesirableForOp(ISD::ABS, ExtVT) &&
         hasOperation(ISD::ABS, ExtVT)) {
       SDLoc DL(N);
       return DAG.getNode(
           ISD::ZERO_EXTEND, DL, VT,
           DAG.getNode(ISD::ABS, DL, ExtVT,
                       DAG.getNode(ISD::TRUNCATE, DL, ExtVT, N0.getOperand(0))));
     }
   }

llvm/test/CodeGen/AArch64/neon-abd.ll

; CHECK-NEXT:    ret
   %sub = sub nuw <2 x i64> %a, %b
   %abs = call <2 x i64> @llvm.abs.v2i64(<2 x i64> %sub, i1 true)
   ret <2 x i64> %abs
 }
 
 define <16 x i8> @sabd_v16i8_nsw(<16 x i8> %a, <16 x i8> %b) #0 {
 ; CHECK-LABEL: sabd_v16i8_nsw:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    sub v0.16b, v0.16b, v1.16b
-; CHECK-NEXT:    abs v0.16b, v0.16b
+; CHECK-NEXT:    sabd v0.16b, v0.16b, v1.16b
 ; CHECK-NEXT:    ret
   %sub = sub nsw <16 x i8> %a, %b
   %abs = call <16 x i8> @llvm.abs.v16i8(<16 x i8> %sub, i1 true)
   ret <16 x i8> %abs
 }
 
 define <8 x i16> @sabd_v8i16_nsw(<8 x i16> %a, <8 x i16> %b) #0 {
 ; CHECK-LABEL: sabd_v8i16_nsw:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    sub v0.8h, v0.8h, v1.8h
-; CHECK-NEXT:    abs v0.8h, v0.8h
+; CHECK-NEXT:    sabd v0.8h, v0.8h, v1.8h
 ; CHECK-NEXT:    ret
   %sub = sub nsw <8 x i16> %a, %b
   %abs = call <8 x i16> @llvm.abs.v8i16(<8 x i16> %sub, i1 true)
   ret <8 x i16> %abs
 }
 
 define <4 x i32> @sabd_v4i32_nsw(<4 x i32> %a, <4 x i32> %b) #0 {
 ; CHECK-LABEL: sabd_v4i32_nsw:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    sub v0.4s, v0.4s, v1.4s
-; CHECK-NEXT:    abs v0.4s, v0.4s
+; CHECK-NEXT:    sabd v0.4s, v0.4s, v1.4s
 ; CHECK-NEXT:    ret
   %sub = sub nsw <4 x i32> %a, %b
   %abs = call <4 x i32> @llvm.abs.v4i32(<4 x i32> %sub, i1 true)
   ret <4 x i32> %abs
 }
 
 define <2 x i64> @sabd_v2i64_nsw(<2 x i64> %a, <2 x i64> %b) #0 {
 ; CHECK-LABEL: sabd_v2i64_nsw: