[AArch64] Transform add(x, abs(y)) -> saba(x, y, 0)

Author: hazzlim

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -50,6 +50,7 @@
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineMemOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/SDPatternMatch.h"
 #include "llvm/CodeGen/SelectionDAG.h"
 #include "llvm/CodeGen/SelectionDAGNodes.h"
 #include "llvm/CodeGen/TargetCallingConv.h"
@@ -21914,6 +21915,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);
+}
+
 static SDValue performAddSubCombine(SDNode *N,
                                     TargetLowering::DAGCombinerInfo &DCI) {
   // Try to change sum of two reductions.
@@ -21939,6 +21990,9 @@ static SDValue performAddSubCombine(SDNode *N,
   if (SDValue Val = performExtBinopLoadFold(N, DCI.DAG))
     return Val;
 
+  if (SDValue Val = performAddSABACombine(N, DCI))
+    return Val;
+
   return performAddSubLongCombine(N, DCI);
 }
 

llvm/lib/Target/AArch64/AArch64InstrInfo.td

@@ -1059,6 +1059,10 @@ def AArch64sdot     : SDNode<"AArch64ISD::SDOT", SDT_AArch64Dot>;
 def AArch64udot     : SDNode<"AArch64ISD::UDOT", SDT_AArch64Dot>;
 def AArch64usdot    : SDNode<"AArch64ISD::USDOT", SDT_AArch64Dot>;
 
+// saba/sabal
+def AArch64neonsaba     : SDNode<"AArch64ISD::SABA", SDT_AArch64trivec>;
+def AArch64neonsabal    : SDNode<"AArch64ISD::SABAL", SDT_AArch64Dot>;
+
 // Vector across-lanes addition
 // Only the lower result lane is defined.
 def AArch64saddv    : SDNode<"AArch64ISD::SADDV", SDT_AArch64UnaryVec>;
@@ -6121,6 +6125,19 @@ defm SQRDMLAH : SIMDThreeSameVectorSQRDMLxHTiedHS<1,0b10000,"sqrdmlah",
 defm SQRDMLSH : SIMDThreeSameVectorSQRDMLxHTiedHS<1,0b10001,"sqrdmlsh",
                                                     int_aarch64_neon_sqrdmlsh>;
 
+def : Pat<(AArch64neonsaba (v8i8 V64:$Rd), V64:$Rn, V64:$Rm),
+          (SABAv8i8 V64:$Rd, V64:$Rn, V64:$Rm)>;
+def : Pat<(AArch64neonsaba (v4i16 V64:$Rd), V64:$Rn, V64:$Rm),
+          (SABAv4i16 V64:$Rd, V64:$Rn, V64:$Rm)>;
+def : Pat<(AArch64neonsaba (v2i32 V64:$Rd), V64:$Rn, V64:$Rm),
+          (SABAv2i32 V64:$Rd, V64:$Rn, V64:$Rm)>;
+def : Pat<(AArch64neonsaba (v16i8 V128:$Rd), V128:$Rn, V128:$Rm),
+          (SABAv16i8 V128:$Rd, V128:$Rn, V128:$Rm)>;
+def : Pat<(AArch64neonsaba (v8i16 V128:$Rd), V128:$Rn, V128:$Rm),
+          (SABAv8i16 V128:$Rd, V128:$Rn, V128:$Rm)>;
+def : Pat<(AArch64neonsaba (v4i32 V128:$Rd), V128:$Rn, V128:$Rm),
+          (SABAv4i32 V128:$Rd, V128:$Rn, V128:$Rm)>;
+
 defm AND : SIMDLogicalThreeVector<0, 0b00, "and", and>;
 defm BIC : SIMDLogicalThreeVector<0, 0b01, "bic",
                                   BinOpFrag<(and node:$LHS, (vnot node:$RHS))> >;
@@ -7008,6 +7025,14 @@ defm : AddSubHNPatterns<ADDHNv2i64_v2i32, ADDHNv2i64_v4i32,
                         SUBHNv2i64_v2i32, SUBHNv2i64_v4i32,
                         v2i32, v2i64, 32>;
 
+// Patterns for SABAL
+def : Pat<(AArch64neonsabal (v8i16 V128:$Rd), (v8i8 V64:$Rn), (v8i8 V64:$Rm)),
+          (SABALv8i8_v8i16 V128:$Rd, V64:$Rn, V64:$Rm)>;
+def : Pat<(AArch64neonsabal (v4i32 V128:$Rd), (v4i16 V64:$Rn), (v4i16 V64:$Rm)),
+          (SABALv4i16_v4i32 V128:$Rd, V64:$Rn, V64:$Rm)>;
+def : Pat<(AArch64neonsabal (v2i64 V128:$Rd), (v2i32 V64:$Rn), (v2i32 V64:$Rm)),
+          (SABALv2i32_v2i64 V128:$Rd, V64:$Rn, V64:$Rm)>;
+
 //----------------------------------------------------------------------------
 // AdvSIMD bitwise extract from vector instruction.
 //----------------------------------------------------------------------------