[AArch64] Improve SAD pattern

[jaykang10]

Summary

Given a vecreduce_add node, detect the below pattern and convert it to the node sequence with UABDL, UADB and UADDLP.

i32 vecreduce_add(
 v16i32 abs(
   v16i32 sub(
     v16i32 zero_extend(v16i8 a), v16i32 zero_extend(v16i8 b))))

i32 vecreduce_add(
  v4i32 UADDLP(
    v8i16 add(
      v8i16 zext(
        v8i8 UABD low8:v16i8 a, low8:v16i8 b
      v8i16 zext(
        v8i8 UABD high8:v16i8 a, high8:v16i8 b

For example, the updated pattern improves the assembly output as below.
The source llvm IR

define i32 @test_sad_v16i8(i8* nocapture readonly %a, i8* nocapture readonly %b) {
entry:
  %0 = bitcast i8* %a to <16 x i8>*
  %1 = load <16 x i8>, <16 x i8>* %0
  %2 = zext <16 x i8> %1 to <16 x i32>
  %3 = bitcast i8* %b to <16 x i8>*
  %4 = load <16 x i8>, <16 x i8>* %3
  %5 = zext <16 x i8> %4 to <16 x i32>
  %6 = sub nsw <16 x i32> %5, %2
  %7 = call <16 x i32> @llvm.abs.v16i32(<16 x i32> %6, i1 true)
  %8 = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> %7) 
  ret i32 %8
}

The assembly output from original pattern.

	ldr	q0, [x0]
	ldr	q1, [x1]
	uabd	v0.16b, v1.16b, v0.16b
	ushll2	v1.8h, v0.16b, #0
	ushll	v0.8h, v0.8b, #0
	uaddl2	v2.4s, v0.8h, v1.8h
	uaddl	v0.4s, v0.4h, v1.4h
	add	v0.4s, v0.4s, v2.4s
	addv	s0, v0.4s
	fmov	w0, s0
	ret

The assembly output from updated pattern.

	ldr	q0, [x0]
	ldr	q1, [x1]
	uabdl	v2.8h, v1.8b, v0.8b
	uabal2	v2.8h, v1.16b, v0.16b
	uaddlp	v0.4s, v2.8h
	addv	s0, v0.4s
	fmov	w0, s0
	ret

[dmgreen]

Nice patch. It's a big combine, but we certainly have bigger elsewhere. Sometimes it's possible to do things like this as smaller pieces, that add up to the same thing. Speaking of which, would it be possible to combine (in a separate patch, maybe as a tablegen pattern) addv (uaddlp x) -> uaddlv x ?

Would it make sense to add signed handling too? If I'm not mistaken, if it has sext's as opposed to zext's, it can turn into addv(uaddlp(sabd,sabd2). The sadb's produce an unsigned value from signed inputs, so the remaining extensions are zext's not sexts and it doesn't need a SADDLP node. That may not be correct though, so, um, make sure you check it.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12004

This doesn't appear to use UADALP. Is that meant as a shorthand for UADA + UADDLP?

12011

This is mostly a style comment, so feel free to ignore if you like. They Opcode's are only used once in the checks, so could be done inline without the variable. Having the "ABSOp0" and "SUB" nodes could just be one variable called "SUB" too, as we are checking the opcode on them straight away.

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

7

It can be good to pre-commit the tests, to just show the differences in the review. It makes it easier to see what the patch does.

[jaykang10]

In D104042#2811212, @dmgreen wrote:

Nice patch. It's a big combine, but we certainly have bigger elsewhere. Sometimes it's possible to do things like this as smaller pieces, that add up to the same thing.

Thanks for your kind comments @dmgreen
I agree with you. I tried to do this combine with smaller pieces but the v16i32 type caused more nodes. Type legalizer split the v16i32 type operations into multiple v4i32 and v4i16 operations because the AArch64 target does not support the native register class for the v16i32 type. In order to avoid the type legalization, I combined vecreduce_add.

Speaking of which, would it be possible to combine (in a separate patch, maybe as a tablegen pattern) addv (uaddlp x) -> uaddlv x ?

It is good point! I have tried below pattern following your suggestion. It seems to work. If you are ok, let me add below pattern in this patch.

let AddedComplexity = 10 in { 
def : Pat<(i32 (extractelt
                 (v4i32 (AArch64uaddv (v4i32 (AArch64uaddlp (v8i16 V128:$op))))),
                 (i64 0))),
          (UADDLVv8i16v V128:$op)>;
}

Would it make sense to add signed handling too? If I'm not mistaken, if it has sext's as opposed to zext's, it can turn into addv(uaddlp(sabd,sabd2). The sadb's produce an unsigned value from signed inputs, so the remaining extensions are zext's not sexts and it doesn't need a SADDLP node. That may not be correct though, so, um, make sure you check it.

Yep, I agree with you. Let me update code to support signed one.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12004

Sorry for typo. Let me update the name.

12011

Yep, let me update it.

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

7

Yep, let me pre-commit this test.

[jaykang10]

Following comments from @dmgreen, updated code.

[jaykang10]

Fixed typo

[dmgreen]

It is good point! I have tried below pattern following your suggestion. It seems to work. If you are ok, let me add below pattern in this patch.

let AddedComplexity = 10 in { 
def : Pat<(i32 (extractelt
                 (v4i32 (AArch64uaddv (v4i32 (AArch64uaddlp (v8i16 V128:$op))))),
                 (i64 0))),
          (UADDLVv8i16v V128:$op)>;
}

Do you mind doing this as a new patch? As it does feel logically separable. If we can test them, it would be good to add the various other sizes too. And, I'm not sure about this, but maybe it doesn't need to start from the extract, and can produce a INSERT_SUBREG like some of the other patterns do (like the ones from SIMDAcrossLanesIntrinsic). That might remove the need for the added complexity, and the INSERT_SUBREG / EXRACT_SUBREG should all get cleared up later in the pipeline.

[jaykang10]

In D104042#2813609, @dmgreen wrote:

It is good point! I have tried below pattern following your suggestion. It seems to work. If you are ok, let me add below pattern in this patch.

let AddedComplexity = 10 in { 
def : Pat<(i32 (extractelt
                 (v4i32 (AArch64uaddv (v4i32 (AArch64uaddlp (v8i16 V128:$op))))),
                 (i64 0))),
          (UADDLVv8i16v V128:$op)>;
}

Do you mind doing this as a new patch? As it does feel logically separable. If we can test them, it would be good to add the various other sizes too. And, I'm not sure about this, but maybe it doesn't need to start from the extract, and can produce a INSERT_SUBREG like some of the other patterns do (like the ones from SIMDAcrossLanesIntrinsic). That might remove the need for the added complexity, and the INSERT_SUBREG / EXRACT_SUBREG should all get cleared up later in the pipeline.

Yep, let me create a new patch for it.

[jaykang10]

Following comments from @dmgreen, updated code.

[jaykang10]

@dmgreen If possible, can we push this patch please? After that, I will create a new patch for the TableGen pattern.

[dmgreen]

Thanks. LGTM

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12073–12074

Just use UABAL directly? It would seem simpler that way. Same for the other SDValues in this function, which seem to be copied more than they need to be.

[jaykang10]

Following comments from @dmgreen, updated code.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12073–12074

Yep, let me update it.

llvm/lib/Target/AArch64/AArch64ISelLowering.h

enum NodeType : unsigned {
   // Vector rounding halving addition
   SRHADD,
   URHADD,
 
   // Absolute difference
   UABD,
   SABD,
 
+  // Unsigned Add Long Pairwise
+  UADDLP,
+
   // udot/sdot instructions
   UDOT,
   SDOT,
 
   // Vector across-lanes min/max
   // Only the lower result lane is defined.
   SMINV,
   UMINV,

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

case AArch64ISD::FIRST_NUMBER:
     MAKE_CASE(AArch64ISD::BITREVERSE_MERGE_PASSTHRU)
     MAKE_CASE(AArch64ISD::BSWAP_MERGE_PASSTHRU)
     MAKE_CASE(AArch64ISD::CTLZ_MERGE_PASSTHRU)
     MAKE_CASE(AArch64ISD::CTPOP_MERGE_PASSTHRU)
     MAKE_CASE(AArch64ISD::DUP_MERGE_PASSTHRU)
     MAKE_CASE(AArch64ISD::INDEX_VECTOR)
     MAKE_CASE(AArch64ISD::UABD)
     MAKE_CASE(AArch64ISD::SABD)
+    MAKE_CASE(AArch64ISD::UADDLP)
     MAKE_CASE(AArch64ISD::CALL_RVMARKER)
   }
 #undef MAKE_CASE
   return nullptr;
 }
 
 MachineBasicBlock *
 AArch64TargetLowering::EmitF128CSEL(MachineInstr &MI,
SDValue AArch64TargetLowering::LowerINTRINSIC_WO_CHAIN(SDValue Op,
   }
   case Intrinsic::aarch64_neon_sabd:
   case Intrinsic::aarch64_neon_uabd: {
     unsigned Opcode = IntNo == Intrinsic::aarch64_neon_uabd ? AArch64ISD::UABD
                                                             : AArch64ISD::SABD;
     return DAG.getNode(Opcode, dl, Op.getValueType(), Op.getOperand(1),
                        Op.getOperand(2));
   }
+  case Intrinsic::aarch64_neon_uaddlp: {
+    unsigned Opcode = AArch64ISD::UADDLP;
+    return DAG.getNode(Opcode, dl, Op.getValueType(), Op.getOperand(1));
+  }
   case Intrinsic::aarch64_neon_sdot:
   case Intrinsic::aarch64_neon_udot:
   case Intrinsic::aarch64_sve_sdot:
   case Intrinsic::aarch64_sve_udot: {
     unsigned Opcode = (IntNo == Intrinsic::aarch64_neon_udot ||
                        IntNo == Intrinsic::aarch64_sve_udot)
                           ? AArch64ISD::UDOT
                           : AArch64ISD::SDOT;
static SDValue foldVectorXorShiftIntoCmp(SDNode *N, SelectionDAG &DAG,
   auto *ShiftAmt = dyn_cast<ConstantSDNode>(Shift.getOperand(1));
   EVT ShiftEltTy = Shift.getValueType().getVectorElementType();
   if (!ShiftAmt || ShiftAmt->getZExtValue() != ShiftEltTy.getSizeInBits() - 1)
     return SDValue();
 
   return DAG.getNode(AArch64ISD::CMGEz, SDLoc(N), VT, Shift.getOperand(0));
 }
 
+// Given a vecreduce_add node, detect the below pattern and convert it to the
+// node sequence with UABDL, UADB and UADDLP.
+//
+// i32 vecreduce_add(
+//  v16i32 abs(
+//    v16i32 sub(
+//      v16i32 zero_extend(v16i8 a), v16i32 zero_extend(v16i8 b))))
+// =================>
+// i32 vecreduce_add(
+//   v4i32 UADDLP(
+//     v8i16 add(
+//       v8i16 zext(
+//         v8i8 UABD low8:v16i8 a, low8:v16i8 b
+//       v8i16 zext(
+//         v8i8 UABD high8:v16i8 a, high8:v16i8 b
+static SDValue performVecReduceAddCombineWithUADALP(SDNode *N,
+                                                    SelectionDAG &DAG) {
+  // Assumed i32 vecreduce_add
+  if (N->getValueType(0) != MVT::i32)
+    return SDValue();
+
+  SDValue VecReduceOp0 = N->getOperand(0);
+  unsigned Opcode = VecReduceOp0.getOpcode();
+  // Assumed v16i32 abs
+  if (Opcode != ISD::ABS || VecReduceOp0->getValueType(0) != MVT::v16i32)
+    return SDValue();
+
+  SDValue ABS = VecReduceOp0;
+  SDValue ABSOp0 = ABS->getOperand(0);
+  Opcode = ABSOp0.getOpcode();
+  // Assumed v16i32 sub
+  if (Opcode != ISD::SUB || ABSOp0->getValueType(0) != MVT::v16i32)
+    return SDValue();
+
+  SDValue SUB = ABSOp0;
+  SDValue SUBOp0 = SUB->getOperand(0);
+  SDValue SUBOp1 = SUB->getOperand(1);
+  unsigned Opcode0 = SUBOp0.getOpcode();
+  unsigned Opcode1 = SUBOp1.getOpcode();
+  // Assumed v16i32 zext
+  if (Opcode0 != ISD::ZERO_EXTEND || Opcode1 != ISD::ZERO_EXTEND ||
+      SUBOp0->getValueType(0) != MVT::v16i32 ||
+      SUBOp1->getValueType(0) != MVT::v16i32)
+    return SDValue();
+
+  SDValue ZEXT0 = SUBOp0;
+  SDValue ZEXT1 = SUBOp1;
+  SDValue ZEXT0Op0 = ZEXT0->getOperand(0);
+  SDValue ZEXT1Op0 = ZEXT1->getOperand(0);
+  // Assumed zext's operand has v16i8 type
+  if (ZEXT0Op0->getValueType(0) != MVT::v16i8 ||
+      ZEXT1Op0->getValueType(0) != MVT::v16i8)
+    return SDValue();
+
+  // Pattern is dectected. Let's convert it to sequence of nodes.
+  SDLoc DL(N);
+
+  // First, create the node pattern of UABDL.
+  SDValue UABDHigh8Op0 =
+      DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, MVT::v8i8, ZEXT0Op0,
+                  DAG.getConstant(8, DL, MVT::i64));
+  SDValue UABDHigh8Op1 =
+      DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, MVT::v8i8, ZEXT1Op0,
+                  DAG.getConstant(8, DL, MVT::i64));
+  SDValue UABDHigh8 =
+      DAG.getNode(AArch64ISD::UABD, DL, MVT::v8i8, UABDHigh8Op0, UABDHigh8Op1);
+  SDValue UABDL = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::v8i16, UABDHigh8);
+
+  // Second, create the node pattern of UABAL.
+  SDValue UABDLo8Op0 = DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, MVT::v8i8,
+                                   ZEXT0Op0, DAG.getConstant(0, DL, MVT::i64));
+  SDValue UABDLo8Op1 = DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, MVT::v8i8,
+                                   ZEXT1Op0, DAG.getConstant(0, DL, MVT::i64));
+  SDValue UABDLo8 =
+      DAG.getNode(AArch64ISD::UABD, DL, MVT::v8i8, UABDLo8Op0, UABDLo8Op1);
+  SDValue ZExtUABD = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::v8i16, UABDLo8);
+  SDValue UABAL = DAG.getNode(ISD::ADD, DL, MVT::v8i16, UABDL, ZExtUABD);
+
+  // Third, create the node of UADDLP.
+  SDValue UADDLPOp0 = UABAL;
+  SDValue UADDLP = DAG.getNode(AArch64ISD::UADDLP, DL, MVT::v4i32, UADDLPOp0);
+
+  // Fourth, create the node of VECREDUCE_ADD.
+  return DAG.getNode(ISD::VECREDUCE_ADD, DL, MVT::i32, UADDLP);
+}
+
 // Turn a v8i8/v16i8 extended vecreduce into a udot/sdot and vecreduce
 //   vecreduce.add(ext(A)) to vecreduce.add(DOT(zero, A, one))
 //   vecreduce.add(mul(ext(A), ext(B))) to vecreduce.add(DOT(zero, A, B))
 static SDValue performVecReduceAddCombine(SDNode *N, SelectionDAG &DAG,
                                           const AArch64Subtarget *ST) {
+  if (!ST->hasDotProd())
+    return performVecReduceAddCombineWithUADALP(N, DAG);
+
   SDValue Op0 = N->getOperand(0);
-  if (!ST->hasDotProd() || N->getValueType(0) != MVT::i32 ||
+  if (N->getValueType(0) != MVT::i32 ||
       Op0.getValueType().getVectorElementType() != MVT::i32)
     return SDValue();
 
   unsigned ExtOpcode = Op0.getOpcode();
   SDValue A = Op0;
   SDValue B;
   if (ExtOpcode == ISD::MUL) {
     A = Op0.getOperand(0);

llvm/lib/Target/AArch64/AArch64InstrInfo.td

 def SDT_AArch64TCRET : SDTypeProfile<0, 2, [SDTCisPtrTy<0>]>;
 def SDT_AArch64PREFETCH : SDTypeProfile<0, 2, [SDTCisVT<0, i32>, SDTCisPtrTy<1>]>;
 
 def SDT_AArch64ITOF  : SDTypeProfile<1, 1, [SDTCisFP<0>, SDTCisSameAs<0,1>]>;
 
 def SDT_AArch64TLSDescCall : SDTypeProfile<0, -2, [SDTCisPtrTy<0>,
                                                  SDTCisPtrTy<1>]>;
 
+def SDT_AArch64uaddlp : SDTypeProfile<1, 1, [SDTCisVec<0>, SDTCisVec<1>]>;
+
 def SDT_AArch64ldp : SDTypeProfile<2, 1, [SDTCisVT<0, i64>, SDTCisSameAs<0, 1>, SDTCisPtrTy<2>]>;
 def SDT_AArch64stp : SDTypeProfile<0, 3, [SDTCisVT<0, i64>, SDTCisSameAs<0, 1>, SDTCisPtrTy<2>]>;
 def SDT_AArch64stnp : SDTypeProfile<0, 3, [SDTCisVT<0, v4i32>, SDTCisSameAs<0, 1>, SDTCisPtrTy<2>]>;
 
 // Generates the general dynamic sequences, i.e.
 //  adrp  x0, :tlsdesc:var
 //  ldr   x1, [x0, #:tlsdesc_lo12:var]
 //  add   x0, x0, #:tlsdesc_lo12:var

 
 def AArch64uabd     : PatFrags<(ops node:$lhs, node:$rhs),
                                [(AArch64uabd_n 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),
                                 (int_aarch64_neon_sabd node:$lhs, node:$rhs)]>;
 
+def AArch64uaddlp_n : SDNode<"AArch64ISD::UADDLP", SDT_AArch64uaddlp>;
+def AArch64uaddlp   : PatFrags<(ops node:$src),
+                               [(AArch64uaddlp_n node:$src),
+                                (int_aarch64_neon_uaddlp node:$src)]>;
+
 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]>;
 
 def SDT_AArch64unpk : SDTypeProfile<1, 1, [
     SDTCisInt<0>, SDTCisInt<1>, SDTCisOpSmallerThanOp<1, 0>

 defm SCVTF  : SIMDTwoVectorIntToFP<0, 0, 0b11101, "scvtf", sint_to_fp>;
 defm SHLL   : SIMDVectorLShiftLongBySizeBHS;
 defm SQABS  : SIMDTwoVectorBHSD<0, 0b00111, "sqabs", int_aarch64_neon_sqabs>;
 defm SQNEG  : SIMDTwoVectorBHSD<1, 0b00111, "sqneg", int_aarch64_neon_sqneg>;
 defm SQXTN  : SIMDMixedTwoVector<0, 0b10100, "sqxtn", int_aarch64_neon_sqxtn>;
 defm SQXTUN : SIMDMixedTwoVector<1, 0b10010, "sqxtun", int_aarch64_neon_sqxtun>;
 defm SUQADD : SIMDTwoVectorBHSDTied<0, 0b00011, "suqadd",int_aarch64_neon_suqadd>;
 defm UADALP : SIMDLongTwoVectorTied<1, 0b00110, "uadalp",
-       BinOpFrag<(add node:$LHS, (int_aarch64_neon_uaddlp node:$RHS))> >;
-defm UADDLP : SIMDLongTwoVector<1, 0b00010, "uaddlp",
-                    int_aarch64_neon_uaddlp>;
+       BinOpFrag<(add node:$LHS, (AArch64uaddlp node:$RHS))> >;
+defm UADDLP : SIMDLongTwoVector<1, 0b00010, "uaddlp", AArch64uaddlp>;
 defm UCVTF  : SIMDTwoVectorIntToFP<1, 0, 0b11101, "ucvtf", uint_to_fp>;
 defm UQXTN  : SIMDMixedTwoVector<1, 0b10100, "uqxtn", int_aarch64_neon_uqxtn>;
 defm URECPE : SIMDTwoVectorS<0, 1, 0b11100, "urecpe", int_aarch64_neon_urecpe>;
 defm URSQRTE: SIMDTwoVectorS<1, 1, 0b11100, "ursqrte", int_aarch64_neon_ursqrte>;
 defm USQADD : SIMDTwoVectorBHSDTied<1, 0b00011, "usqadd",int_aarch64_neon_usqadd>;
 defm XTN    : SIMDMixedTwoVector<0, 0b10010, "xtn", trunc>;
 
 def : Pat<(v4f16  (AArch64rev32 V64:$Rn)),  (REV32v4i16 V64:$Rn)>;

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

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple aarch64-none-linux-gnu < %s | FileCheck %s
+
+declare <16 x i32> @llvm.abs.v16i32(<16 x i32>, i1 immarg)
+declare i32 @llvm.vector.reduce.add.v16i32(<16 x i32>)
+
+define i32 @test_sad_v16i8(i8* nocapture readonly %a, i8* nocapture readonly %b) {
+; CHECK-LABEL: test_sad_v16i8:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ldr q0, [x0]
+; CHECK-NEXT:    ldr q1, [x1]
+; CHECK-NEXT:    uabdl v2.8h, v1.8b, v0.8b
+; CHECK-NEXT:    uabal2 v2.8h, v1.16b, v0.16b
+; CHECK-NEXT:    uaddlp v0.4s, v2.8h
+; CHECK-NEXT:    addv s0, v0.4s
+; CHECK-NEXT:    fmov w0, s0
+; CHECK-NEXT:    ret
+entry:
+  %0 = bitcast i8* %a to <16 x i8>*
+  %1 = load <16 x i8>, <16 x i8>* %0
+  %2 = zext <16 x i8> %1 to <16 x i32>
+  %3 = bitcast i8* %b to <16 x i8>*
+  %4 = load <16 x i8>, <16 x i8>* %3
+  %5 = zext <16 x i8> %4 to <16 x i32>
+  %6 = sub nsw <16 x i32> %5, %2
+  %7 = call <16 x i32> @llvm.abs.v16i32(<16 x i32> %6, i1 true)
+  %8 = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> %7)
+  ret i32 %8
+}