[AArch64] Fix abs(sub nsw) -> absd

[rjj]

Summary

This partially reverts a regression introduced in 8f25e382c5b1 for AArch64 targets. In particular, we restore the logic of (abs (sub nsw x, y)) -> abds(x, y) for all targets except X86, which keeps the logic introduced in 8f25e382c5b1. See also https://reviews.llvm.org/D142288.

Depends on D144399

[dmgreen]

Can you make sure you upload with full context (-U999999), as per https://llvm.org/docs/DeveloperPolicy.html#making-and-submitting-a-patch. It can help to make the patch easier to review.

llvm/include/llvm/CodeGen/TargetLowering.h

801

Perhaps preferABDSToABSWithNSW? It would be good to add the "nsw" part, as there are other ways to produce ABDS

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

1 ↗

(On Diff #498794)

Nothing in this file seems to change. We already have abd tests for the nsw cases in neon-abd.ll.

If you would like to add saba tests too, perhaps split this off into a separate commit? It can be good to do that anyway to see the differences in generated code in the patch.

[SjoerdMeijer]

Thanks for looking into this.
I will let @RKSimon comment on the usage of the TTI hook.

I had a couple of nits about the tests, see inlined.
You can consider to "precommit" the tests: put them up for review separately, then rebase this patch on top of that. The codegen differences are easy to spot here, so not necessary, but up to you if you want to get that out of the way first.

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

5 ↗

(On Diff #498794)

If this is NEON, then we probably don't want to be passing +sve2 here.
And +neon is passed twice.

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

5

sve is implied by sve2, and I think neon is implied too, so just sve2 will suffice.

I don't have a strong opinion on this, but you can also use llc -mattr=+sve2 so you don't need these attributes at all.

[rjj]

• Split tests to D144399
• Rename preferABDSToABS -> preferABDSToABSWithNSW

[RKSimon]

llvm/include/llvm/CodeGen/TargetLowering.h

802

just return true by default - isOperationLegalOrCustom can cause problems after custom lowering if it re-expands the ABDS pattern

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10523–10524

It would be better to use hasOperation() && TLI.preferABDSToABSWithNSW and remove the Legal/Custom checks from the callback - otherwise you can potentially end up in loops where the custom lowering re-expands the abds pattern and its gets matched over and over again....

llvm/lib/Target/X86/X86ISelLowering.cpp

57121

just return false - x86 doesn't have any ABDS instructions

[rjj]

Remove Legal/Custom checks from the callback as suggested by @RKSimon.

[RKSimon]

LGTM - cheers

llvm/include/llvm/CodeGen/TargetLowering.h

public:
   ///   add (add %x, 1), %y
   /// The variant with two add's is IR-canonical.
   /// Some targets may prefer one to the other.
   virtual bool preferIncOfAddToSubOfNot(EVT VT) const {
     // By default, let's assume that everyone prefers the form with two add's.
     return true;
   }
 
+  // By default prefer folding (abs (sub nsw x, y)) -> abds(x, y). Some targets
+  // may want to avoid this to prevent loss of sub_nsw pattern.
+  virtual bool preferABDSToABSWithNSW(EVT VT) const {
+    return true;
+  }
+
   // Return true if the target wants to transform Op(Splat(X)) -> Splat(Op(X))
   virtual bool preferScalarizeSplat(unsigned Opc) const { return true; }
 
   /// Return true if the target wants to use the optimization that
   /// turns ext(promotableInst1(...(promotableInstN(load)))) into
   /// promotedInst1(...(promotedInstN(ext(load)))).
   bool enableExtLdPromotion() const { return EnableExtLdPromotion; }
 

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::foldABSToABD(SDNode *N) {
   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)) {
     // fold (abs (sub nsw x, y)) -> abds(x, y)
-    // Limit this to legal ops to prevent loss of sub_nsw pattern.
-    if (AbsOp1->getFlags().hasNoSignedWrap() &&
-        TLI.isOperationLegal(ISD::ABDS, VT)) {
+    if (AbsOp1->getFlags().hasNoSignedWrap() && hasOperation(ISD::ABDS, VT) &&
+        TLI.preferABDSToABSWithNSW(VT)) {
       SDValue ABD = DAG.getNode(ISD::ABDS, DL, VT, Op0, Op1);
       return DAG.getZExtOrTrunc(ABD, DL, SrcVT);
     }
     return SDValue();
   }
 
   EVT VT1 = Op0.getOperand(0).getValueType();
   EVT VT2 = Op1.getOperand(0).getValueType();

llvm/lib/Target/X86/X86ISelLowering.h

public:
     /// Replace the results of node with an illegal result
     /// type with new values built out of custom code.
     ///
     void ReplaceNodeResults(SDNode *N, SmallVectorImpl<SDValue>&Results,
                             SelectionDAG &DAG) const override;
 
     SDValue PerformDAGCombine(SDNode *N, DAGCombinerInfo &DCI) const override;
 
+    bool preferABDSToABSWithNSW(EVT VT) const override;
+
     /// Return true if the target has native support for
     /// the specified value type and it is 'desirable' to use the type for the
     /// given node type. e.g. On x86 i16 is legal, but undesirable since i16
     /// instruction encodings are longer and some i16 instructions are slow.
     bool isTypeDesirableForOp(unsigned Opc, EVT VT) const override;
 
     /// Return true if the target has native support for the
     /// specified value type and it is 'desirable' to use the type. e.g. On x86

llvm/lib/Target/X86/X86ISelLowering.cpp

SDValue X86TargetLowering::PerformDAGCombine(SDNode *N,
   case X86ISD::SUBV_BROADCAST_LOAD: return combineBROADCAST_LOAD(N, DAG, DCI);
   case X86ISD::MOVDQ2Q:     return combineMOVDQ2Q(N, DAG);
   case X86ISD::PDEP:        return combinePDEP(N, DAG, DCI);
   }
 
   return SDValue();
 }
 
+bool X86TargetLowering::preferABDSToABSWithNSW(EVT VT) const {
+  return false;
+}
+
 bool X86TargetLowering::isTypeDesirableForOp(unsigned Opc, EVT VT) const {
   if (!isTypeLegal(VT))
     return false;
 
   // There are no vXi8 shifts.
   if (Opc == ISD::SHL && VT.isVector() && VT.getVectorElementType() == MVT::i8)
     return false;
 

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

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple aarch64-unknown-linux-gnu -mattr=+sve2 < %s | FileCheck %s
 
 ; SABA from ADD(ABS(SUB NSW))
 
 define <vscale x 2 x i64> @saba_abs_d(<vscale x 2 x i64> %a, <vscale x 2 x i64> %b, <vscale x 2 x i64> %c) #0 {
 ; CHECK-LABEL: saba_abs_d:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptrue p0.d
-; CHECK-NEXT:    sub z1.d, z1.d, z2.d
-; CHECK-NEXT:    abs z1.d, p0/m, z1.d
-; CHECK-NEXT:    add z0.d, z0.d, z1.d
+; CHECK-NEXT:    saba z0.d, z1.d, z2.d
 ; CHECK-NEXT:    ret
   %sub = sub nsw <vscale x 2 x i64> %b, %c
   %abs = call <vscale x 2 x i64> @llvm.abs.nxv2i64(<vscale x 2 x i64> %sub, i1 true)
   %add = add <vscale x 2 x i64> %a, %abs
   ret <vscale x 2 x i64> %add
 }
 
 define <vscale x 4 x i32> @saba_abs_s(<vscale x 4 x i32> %a, <vscale x 4 x i32> %b, <vscale x 4 x i32> %c) #0 {
 ; CHECK-LABEL: saba_abs_s:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptrue p0.s
-; CHECK-NEXT:    sub z1.s, z1.s, z2.s
-; CHECK-NEXT:    abs z1.s, p0/m, z1.s
-; CHECK-NEXT:    add z0.s, z0.s, z1.s
+; CHECK-NEXT:    saba z0.s, z1.s, z2.s
 ; CHECK-NEXT:    ret
   %sub = sub nsw <vscale x 4 x i32> %b, %c
   %abs = call <vscale x 4 x i32> @llvm.abs.nxv4i32(<vscale x 4 x i32> %sub, i1 true)
   %add = add <vscale x 4 x i32> %a, %abs
   ret <vscale x 4 x i32> %add
 }
 
 define <vscale x 8 x i16> @saba_abs_h(<vscale x 8 x i16> %a, <vscale x 8 x i16> %b, <vscale x 8 x i16> %c) #0 {
 ; CHECK-LABEL: saba_abs_h:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptrue p0.h
-; CHECK-NEXT:    sub z1.h, z1.h, z2.h
-; CHECK-NEXT:    abs z1.h, p0/m, z1.h
-; CHECK-NEXT:    add z0.h, z0.h, z1.h
+; CHECK-NEXT:    saba z0.h, z1.h, z2.h
 ; CHECK-NEXT:    ret
   %sub = sub nsw <vscale x 8 x i16> %b, %c
   %abs = call <vscale x 8 x i16> @llvm.abs.nxv8i16(<vscale x 8 x i16> %sub, i1 true)
   %add = add <vscale x 8 x i16> %a, %abs
   ret <vscale x 8 x i16> %add
 }
 
 define <vscale x 16 x i8> @saba_abs_b(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b, <vscale x 16 x i8> %c) #0 {
 ; CHECK-LABEL: saba_abs_b:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptrue p0.b
-; CHECK-NEXT:    sub z1.b, z1.b, z2.b
-; CHECK-NEXT:    abs z1.b, p0/m, z1.b
-; CHECK-NEXT:    add z0.b, z0.b, z1.b
+; CHECK-NEXT:    saba z0.b, z1.b, z2.b
 ; CHECK-NEXT:    ret
   %sub = sub nsw <vscale x 16 x i8> %b, %c
   %abs = call <vscale x 16 x i8> @llvm.abs.nxv16i8(<vscale x 16 x i8> %sub, i1 true)
   %add = add <vscale x 16 x i8> %a, %abs
   ret <vscale x 16 x i8> %add
 }
 
 ; SABA from ADD(SABD)