[llvm] [AArch64] Transform add(x, abs(y)) -> saba(x, y, 0) (PR #156615)

[David Sherwood via llvm-commits]

================
@@ -21913,6 +21914,56 @@ static SDValue performExtBinopLoadFold(SDNode *N, SelectionDAG &DAG) {
   return DAG.getNode(N->getOpcode(), DL, VT, Ext0, NShift);
 }
 
+// Transform the following:
+// - add(x, abs(y)) -> saba(x, y, 0)
+// - add(x, zext(abs(y))) -> sabal(x, y, 0)
+static SDValue performAddSABACombine(SDNode *N,
+                                     TargetLowering::DAGCombinerInfo &DCI) {
+  if (N->getOpcode() != ISD::ADD)
+    return SDValue();
+
+  EVT VT = N->getValueType(0);
+  if (!VT.isFixedLengthVector())
+    return SDValue();
+
+  SDValue N0 = N->getOperand(0);
+  SDValue N1 = N->getOperand(1);
+
+  auto MatchAbsOrZExtAbs = [](SDValue V0, SDValue V1, SDValue &AbsOp,
+                              SDValue &Other, bool &IsZExt) {
+    Other = V1;
+    if (sd_match(V0, m_Abs(SDPatternMatch::m_Value(AbsOp)))) {
+      IsZExt = false;
+      return true;
+    }
+    if (sd_match(V0, SDPatternMatch::m_ZExt(
+                         m_Abs(SDPatternMatch::m_Value(AbsOp))))) {
+      IsZExt = true;
+      return true;
+    }
+
+    return false;
+  };
+
+  SDValue AbsOp;
+  SDValue Other;
+  bool IsZExt;
+  if (!MatchAbsOrZExtAbs(N0, N1, AbsOp, Other, IsZExt) &&
+      !MatchAbsOrZExtAbs(N1, N0, AbsOp, Other, IsZExt))
+    return SDValue();
+
+  // Don't perform this on abs(sub), as this will become an ABD/ABA anyway.
+  if (AbsOp.getOpcode() == ISD::SUB)
+    return SDValue();
+
+  SDLoc DL(N);
+  SDValue Zero = DCI.DAG.getConstant(0, DL, MVT::i64);
+  SDValue Zeros = DCI.DAG.getSplatVector(AbsOp.getValueType(), DL, Zero);
+
+  unsigned Opcode = IsZExt ? AArch64ISD::SABAL : AArch64ISD::SABA;
+  return DCI.DAG.getNode(Opcode, DL, VT, Other, AbsOp, Zeros);
----------------
david-arm wrote:

Shouldn't we also be checking if the operation is legal for the instruction? If we allow this DAG combine to run pre-legalisation then we could be zero-extending 64-bit to 128-bit and I don't think we support that.

https://github.com/llvm/llvm-project/pull/156615