[ISel] Port AArch64 SABD and UABD to DAGCombine

[dmgreen]

Summary

This ports the AArch64 SABD and USBD over to DAG Combine, where they can be used by more backends (notably MVE in a follow-up patch). The matching code has changed very little, just to handle legal operations and types differently. It selects from (ABS (SUB (EXTEND a), (EXTEND b))), producing a ubds/abdu which is zexted to the original type.

[RKSimon]

What other archs have equivalent instructions?

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

8882

Rename to VT1 /VT2? I don't think there's anything that technically requires them to be vectors?

llvm/lib/CodeGen/TargetLoweringBase.cpp

776

// Absolute differences

[RKSimon]

In D91937#2410809, @RKSimon wrote:

What other archs have equivalent instructions?

@nemanjai Does PPC have equivalent instructions?

[dmgreen]

Thanks for taking a look.

What other archs have equivalent instructions?

It was MVE I was hoping to share this with, over in the ARM backend. I don't know if any other backend would find this useful too - it would be great if they did. I find it hard enough to know what ARM instructions do some of the time let along other architectures! I believe the X86 instruction (PSADBW) works quite differently, as it includes a reduction step?

[RKSimon]

Yeah - sum-of-abs-diff is a bit more complex - technically we could use ABD as a stepping stone in pattern matching it

llvm/include/llvm/CodeGen/ISDOpcodes.h

584

Add comments showing both ADDS and ADDU patterns

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

8864

naming consistency - combineAbsToAbd or combineABSToABD

[dmgreen]

Updated patterns in comments and renamed to combineABSToABD.

[RKSimon]

llvm/include/llvm/Target/TargetSelectionDAG.td

362

Can't these be commutative ?

[RKSimon]

llvm/include/llvm/Target/TargetSelectionDAG.td

362

alive2 says yes: https://alive2.llvm.org/ce/z/8giJFY

[dmgreen]

llvm/include/llvm/Target/TargetSelectionDAG.td

362

Aye, I agree. Sorry I have not gotten back to this yet, and didn't have time over the weekend to update it. I'll try and do so soon.

I did happen to talk to someone how implemented the same thing in gcc and apparently they represent it as max(a, b) - min(a, b). Which I thought was nifty but does show how the whole thing can be commutative.

[dmgreen]

Now with Commutative. Thanks!

[RKSimon]

LGTM

[RKSimon]

@dmgreen What happened to this?

[dmgreen]

I still need to do something with D91921 before it can be used with MVE. I was working on that, but it was a bit slow going and other things had come up in the meantime. I was looking at it recently to try and improve trunc/extend lowering, but need to get some time to sort through that properly. There is also some stuff to do with the way that MVE wants to do lane interleaving that I thought might be an issue, but I'm pretty sure that will be fine.

This part can go in if we need it I think. But the MVE side won't be working well yet, and so there won't be another user other than AArch64.

[dmgreen]

Rebase to current trunk.

llvm/include/llvm/CodeGen/ISDOpcodes.h

enum NodeType {
   /// i[N+1], halving the result by shifting it one bit right.
   HADDS,
   HADDU,
   /// RHADDS/RHADDU - Rounding having add - Add two integers using an integer of
   /// type i[N+2], add 1 and halve the result by shifting it one bit right.
   RHADDS,
   RHADDU,
 
+  // ABDS/ABDU - Absolute difference - Return the absolute difference between
+  // two numbers interpreted as signed/unsigned.
+  // i.e trunc(abs(sext(Op0) - sext(Op1))) becomes abds(Op0, Op1)
+  //  or trunc(abs(zext(Op0) - zext(Op1))) becomes abdu(Op0, Op1)
+  ABDS,
+  ABDU,
+
   /// [US]{MIN/MAX} - Binary minimum or maximum or signed or unsigned
   /// integers.
   SMIN,
   SMAX,
   UMIN,
   UMAX,
 
   /// Bitwise operators - logical and, logical or, logical xor.

llvm/include/llvm/Target/TargetSelectionDAG.td

 def mul        : SDNode<"ISD::MUL"       , SDTIntBinOp,
                         [SDNPCommutative, SDNPAssociative]>;
 def mulhs      : SDNode<"ISD::MULHS"     , SDTIntBinOp, [SDNPCommutative]>;
 def mulhu      : SDNode<"ISD::MULHU"     , SDTIntBinOp, [SDNPCommutative]>;
 def hadds      : SDNode<"ISD::HADDS"     , SDTIntBinOp, [SDNPCommutative]>;
 def haddu      : SDNode<"ISD::HADDU"     , SDTIntBinOp, [SDNPCommutative]>;
 def rhadds     : SDNode<"ISD::RHADDS"    , SDTIntBinOp, [SDNPCommutative]>;
 def rhaddu     : SDNode<"ISD::RHADDU"    , SDTIntBinOp, [SDNPCommutative]>;
+def abds       : SDNode<"ISD::ABDS"      , SDTIntBinOp, [SDNPCommutative]>;
+def abdu       : SDNode<"ISD::ABDU"      , SDTIntBinOp, [SDNPCommutative]>;
 def smullohi   : SDNode<"ISD::SMUL_LOHI" , SDTIntBinHiLoOp, [SDNPCommutative]>;
 def umullohi   : SDNode<"ISD::UMUL_LOHI" , SDTIntBinHiLoOp, [SDNPCommutative]>;
 def sdiv       : SDNode<"ISD::SDIV"      , SDTIntBinOp>;
 def udiv       : SDNode<"ISD::UDIV"      , SDTIntBinOp>;
 def srem       : SDNode<"ISD::SREM"      , SDTIntBinOp>;
 def urem       : SDNode<"ISD::UREM"      , SDTIntBinOp>;
 def sdivrem    : SDNode<"ISD::SDIVREM"   , SDTIntBinHiLoOp>;
 def udivrem    : SDNode<"ISD::UDIVREM"   , SDTIntBinHiLoOp>;

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitFunnelShift(SDNode *N) {
 
   // Simplify, based on bits shifted out of N0/N1.
   if (SimplifyDemandedBits(SDValue(N, 0)))
     return SDValue(N, 0);
 
   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) {
+  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))
+    return SDValue();
+
+  EVT VT1 = Op0.getOperand(0).getValueType();
+  EVT VT2 = Op1.getOperand(0).getValueType();
+  // Check if the operands are of same type and valid size.
+  unsigned ABDOpcode = (Opc0 == ISD::SIGN_EXTEND) ? ISD::ABDS : ISD::ABDU;
+  if (VT1 != VT2 || !TLI.isOperationLegalOrCustom(ABDOpcode, VT1))
+    return SDValue();
+
+  Op0 = Op0.getOperand(0);
+  Op1 = Op1.getOperand(0);
+  SDValue ABD =
+      DAG.getNode(ABDOpcode, SDLoc(N), Op0->getValueType(0), Op0, Op1);
+  return DAG.getNode(ISD::ZERO_EXTEND, SDLoc(N), N->getValueType(0), ABD);
+}
+
 SDValue DAGCombiner::visitABS(SDNode *N) {
   SDValue N0 = N->getOperand(0);
   EVT VT = N->getValueType(0);
 
   // fold (abs c1) -> c2
   if (DAG.isConstantIntBuildVectorOrConstantInt(N0))
     return DAG.getNode(ISD::ABS, SDLoc(N), VT, N0);
   // fold (abs (abs x)) -> (abs x)
   if (N0.getOpcode() == ISD::ABS)
     return N0;
   // fold (abs x) -> x iff not-negative
   if (DAG.SignBitIsZero(N0))
     return N0;
+
+  if (SDValue ABD = combineABSToABD(N, DAG, TLI))
+    return ABD;
+
   return SDValue();
 }
 
 SDValue DAGCombiner::visitBSWAP(SDNode *N) {
   SDValue N0 = N->getOperand(0);
   EVT VT = N->getValueType(0);
 
   // fold (bswap c1) -> c2

llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp

#endif
   case ISD::SUB:                        return "sub";
   case ISD::MUL:                        return "mul";
   case ISD::MULHU:                      return "mulhu";
   case ISD::MULHS:                      return "mulhs";
   case ISD::HADDU:                      return "haddu";
   case ISD::HADDS:                      return "hadds";
   case ISD::RHADDU:                     return "rhaddu";
   case ISD::RHADDS:                     return "rhadds";
+  case ISD::ABDS:                       return "abds";
+  case ISD::ABDU:                       return "abdu";
   case ISD::SDIV:                       return "sdiv";
   case ISD::UDIV:                       return "udiv";
   case ISD::SREM:                       return "srem";
   case ISD::UREM:                       return "urem";
   case ISD::SMUL_LOHI:                  return "smul_lohi";
   case ISD::UMUL_LOHI:                  return "umul_lohi";
   case ISD::SDIVREM:                    return "sdivrem";
   case ISD::UDIVREM:                    return "udivrem";

llvm/lib/CodeGen/TargetLoweringBase.cpp

for (MVT VT : MVT::all_valuetypes()) {
     setOperationAction(ISD::SUBC, VT, Expand);
     setOperationAction(ISD::SUBE, VT, Expand);
 
     // Halving adds
     setOperationAction(ISD::HADDS, VT, Expand);
     setOperationAction(ISD::HADDU, VT, Expand);
     setOperationAction(ISD::RHADDS, VT, Expand);
     setOperationAction(ISD::RHADDU, VT, Expand);
+    // Absolute difference
+    setOperationAction(ISD::ABDS, VT, Expand);
+    setOperationAction(ISD::ABDU, VT, Expand);
 
     // These default to Expand so they will be expanded to CTLZ/CTTZ by default.
     setOperationAction(ISD::CTLZ_ZERO_UNDEF, VT, Expand);
     setOperationAction(ISD::CTTZ_ZERO_UNDEF, VT, Expand);
 
     setOperationAction(ISD::BITREVERSE, VT, Expand);
     setOperationAction(ISD::PARITY, VT, Expand);
 

llvm/lib/Target/AArch64/AArch64ISelLowering.h

enum NodeType : unsigned {
   FCMLEz,
   FCMLTz,
 
   // Vector across-lanes addition
   // Only the lower result lane is defined.
   SADDV,
   UADDV,
 
-  // Absolute difference
-  UABD,
-  SABD,
-
   // Vector across-lanes min/max
   // Only the lower result lane is defined.
   SMINV,
   UMINV,
   SMAXV,
   UMAXV,
 
   SADDV_PRED,

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (Subtarget->hasNEON()) {
     }
 
     for (MVT VT : {MVT::v8i8, MVT::v4i16, MVT::v2i32, MVT::v16i8, MVT::v8i16,
                    MVT::v4i32}) {
       setOperationAction(ISD::HADDS, VT, Legal);
       setOperationAction(ISD::HADDU, VT, Legal);
       setOperationAction(ISD::RHADDS, VT, Legal);
       setOperationAction(ISD::RHADDU, VT, Legal);
+      setOperationAction(ISD::ABDS, VT, Legal);
+      setOperationAction(ISD::ABDU, VT, Legal);
     }
 
     // Vector reductions
     for (MVT VT : { MVT::v4f16, MVT::v2f32,
                     MVT::v8f16, MVT::v4f32, MVT::v2f64 }) {
       setOperationAction(ISD::VECREDUCE_FMAX, VT, Custom);
       setOperationAction(ISD::VECREDUCE_FMIN, VT, Custom);
     }
case AArch64ISD::FIRST_NUMBER:
     MAKE_CASE(AArch64ISD::SST1_IMM_PRED)
     MAKE_CASE(AArch64ISD::SSTNT1_PRED)
     MAKE_CASE(AArch64ISD::SSTNT1_INDEX_PRED)
     MAKE_CASE(AArch64ISD::LDP)
     MAKE_CASE(AArch64ISD::STP)
     MAKE_CASE(AArch64ISD::STNP)
     MAKE_CASE(AArch64ISD::DUP_MERGE_PASSTHRU)
     MAKE_CASE(AArch64ISD::INDEX_VECTOR)
-    MAKE_CASE(AArch64ISD::UABD)
-    MAKE_CASE(AArch64ISD::SABD)
   }
 #undef MAKE_CASE
   return nullptr;
 }
 
 MachineBasicBlock *
 AArch64TargetLowering::EmitF128CSEL(MachineInstr &MI,
                                     MachineBasicBlock *MBB) const {
bool IsRoundingAdd = (IntNo == Intrinsic::aarch64_neon_srhadd ||
                           IntNo == Intrinsic::aarch64_neon_urhadd);
     unsigned Opcode = IsSignedAdd ? (IsRoundingAdd ? ISD::RHADDS : ISD::HADDS)
                                   : (IsRoundingAdd ? ISD::RHADDU : ISD::HADDU);
     return DAG.getNode(Opcode, dl, Op.getValueType(), Op.getOperand(1),
                        Op.getOperand(2));
   }
 
   case Intrinsic::aarch64_neon_uabd: {
-    return DAG.getNode(AArch64ISD::UABD, dl, Op.getValueType(),
+    return DAG.getNode(ISD::ABDU, dl, Op.getValueType(),
                        Op.getOperand(1), Op.getOperand(2));
   }
   case Intrinsic::aarch64_neon_sabd: {
-    return DAG.getNode(AArch64ISD::SABD, dl, Op.getValueType(),
+    return DAG.getNode(ISD::ABDS, dl, Op.getValueType(),
                        Op.getOperand(1), Op.getOperand(2));
   }
   }
 }
 
 bool AArch64TargetLowering::shouldRemoveExtendFromGSIndex(EVT VT) const {
   if (VT.getVectorElementType() == MVT::i32 &&
       VT.getVectorElementCount().getKnownMinValue() >= 4)
auto DotIntrisic = (ExtOpcode == ISD::ZERO_EXTEND)
                          ? Intrinsic::aarch64_neon_udot
                          : Intrinsic::aarch64_neon_sdot;
   SDValue Dot = DAG.getNode(ISD::INTRINSIC_WO_CHAIN, DL, Zeros.getValueType(),
                             DAG.getConstant(DotIntrisic, DL, MVT::i32), Zeros,
                             Ones, Op0.getOperand(0));
   return DAG.getNode(ISD::VECREDUCE_ADD, DL, N->getValueType(0), Dot);
 }
 
-// Given a ABS node, detect the following pattern:
-// (ABS (SUB (EXTEND a), (EXTEND b))).
-// Generates UABD/SABD instruction.
-static SDValue performABSCombine(SDNode *N, SelectionDAG &DAG,
-                                 TargetLowering::DAGCombinerInfo &DCI,
-                                 const AArch64Subtarget *Subtarget) {
-  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))
-    return SDValue();
-
-  EVT VectorT1 = Op0.getOperand(0).getValueType();
-  EVT VectorT2 = Op1.getOperand(0).getValueType();
-  // Check if vectors are of same type and valid size.
-  uint64_t Size = VectorT1.getFixedSizeInBits();
-  if (VectorT1 != VectorT2 || (Size != 64 && Size != 128))
-    return SDValue();
-
-  // Check if vector element types are valid.
-  EVT VT1 = VectorT1.getVectorElementType();
-  if (VT1 != MVT::i8 && VT1 != MVT::i16 && VT1 != MVT::i32)
-    return SDValue();
-
-  Op0 = Op0.getOperand(0);
-  Op1 = Op1.getOperand(0);
-  unsigned ABDOpcode =
-      (Opc0 == ISD::SIGN_EXTEND) ? AArch64ISD::SABD : AArch64ISD::UABD;
-  SDValue ABD =
-      DAG.getNode(ABDOpcode, SDLoc(N), Op0->getValueType(0), Op0, Op1);
-  return DAG.getNode(ISD::ZERO_EXTEND, SDLoc(N), N->getValueType(0), ABD);
-}
-
 static SDValue performXorCombine(SDNode *N, SelectionDAG &DAG,
                                  TargetLowering::DAGCombinerInfo &DCI,
                                  const AArch64Subtarget *Subtarget) {
   if (DCI.isBeforeLegalizeOps())
     return SDValue();
 
   if (SDValue Cmp = foldVectorXorShiftIntoCmp(N, DAG, Subtarget))
     return Cmp;

 static SDValue performExtendCombine(SDNode *N,
                                     TargetLowering::DAGCombinerInfo &DCI,
                                     SelectionDAG &DAG) {
   // If we see something like (zext (sabd (extract_high ...), (DUP ...))) then
   // we can convert that DUP into another extract_high (of a bigger DUP), which
   // helps the backend to decide that an sabdl2 would be useful, saving a real
   // extract_high operation.
   if (!DCI.isBeforeLegalizeOps() && N->getOpcode() == ISD::ZERO_EXTEND &&
-      (N->getOperand(0).getOpcode() == AArch64ISD::UABD ||
-       N->getOperand(0).getOpcode() == AArch64ISD::SABD)) {
+      (N->getOperand(0).getOpcode() == ISD::ABDU ||
+       N->getOperand(0).getOpcode() == ISD::ABDS)) {
     SDNode *ABDNode = N->getOperand(0).getNode();
     SDValue NewABD =
         tryCombineLongOpWithDup(Intrinsic::not_intrinsic, ABDNode, DCI, DAG);
     if (!NewABD.getNode())
       return SDValue();
 
     return DAG.getNode(ISD::ZERO_EXTEND, SDLoc(N), N->getValueType(0), NewABD);
   }

 
 SDValue AArch64TargetLowering::PerformDAGCombine(SDNode *N,
                                                  DAGCombinerInfo &DCI) const {
   SelectionDAG &DAG = DCI.DAG;
   switch (N->getOpcode()) {
   default:
     LLVM_DEBUG(dbgs() << "Custom combining: skipping\n");
     break;
-  case ISD::ABS:
-    return performABSCombine(N, DAG, DCI, Subtarget);
   case ISD::ADD:
   case ISD::SUB:
     return performAddSubCombine(N, DCI, DAG);
   case ISD::XOR:
     return performXorCombine(N, DAG, DCI, Subtarget);
   case ISD::MUL:
     return performMulCombine(N, DAG, DCI, Subtarget);
   case ISD::SINT_TO_FP:

llvm/lib/Target/AArch64/AArch64InstrInfo.td

 
 def AArch64saddv    : SDNode<"AArch64ISD::SADDV", SDT_AArch64UnaryVec>;
 def AArch64uaddv    : SDNode<"AArch64ISD::UADDV", SDT_AArch64UnaryVec>;
 def AArch64sminv    : SDNode<"AArch64ISD::SMINV", SDT_AArch64UnaryVec>;
 def AArch64uminv    : SDNode<"AArch64ISD::UMINV", SDT_AArch64UnaryVec>;
 def AArch64smaxv    : SDNode<"AArch64ISD::SMAXV", SDT_AArch64UnaryVec>;
 def AArch64umaxv    : SDNode<"AArch64ISD::UMAXV", SDT_AArch64UnaryVec>;
 
-def AArch64uabd_n   : SDNode<"AArch64ISD::UABD", SDT_AArch64binvec>;
-def AArch64sabd_n   : SDNode<"AArch64ISD::SABD", SDT_AArch64binvec>;
-
 def AArch64uabd     : PatFrags<(ops node:$lhs, node:$rhs),
-                               [(AArch64uabd_n node:$lhs, node:$rhs),
+                               [(abdu node:$lhs, node:$rhs),
                                 (int_aarch64_neon_uabd node:$lhs, node:$rhs)]>;
 def AArch64sabd     : PatFrags<(ops node:$lhs, node:$rhs),
-                               [(AArch64sabd_n node:$lhs, node:$rhs),
+                               [(abds node:$lhs, node:$rhs),
                                 (int_aarch64_neon_sabd node:$lhs, node:$rhs)]>;
 
 def SDT_AArch64SETTAG : SDTypeProfile<0, 2, [SDTCisPtrTy<0>, SDTCisPtrTy<1>]>;
 def AArch64stg : SDNode<"AArch64ISD::STG", SDT_AArch64SETTAG, [SDNPHasChain, SDNPMayStore, SDNPMemOperand]>;
 def AArch64stzg : SDNode<"AArch64ISD::STZG", SDT_AArch64SETTAG, [SDNPHasChain, SDNPMayStore, SDNPMemOperand]>;
 def AArch64st2g : SDNode<"AArch64ISD::ST2G", SDT_AArch64SETTAG, [SDNPHasChain, SDNPMayStore, SDNPMemOperand]>;
 def AArch64stz2g : SDNode<"AArch64ISD::STZ2G", SDT_AArch64SETTAG, [SDNPHasChain, SDNPMayStore, SDNPMemOperand]>;