[AArch64] Remove copy instruction between uaddlv and urshr

[jaykang10]

Summary

gcc generates less number of instructions from below example than llvm.

#include <arm_neon.h>

uint8x8_t foo(uint8x8_t a) {
    return vdup_n_u8(vrshrd_n_u64(vaddlv_u8(a), 3));
}

gcc output
foo:
        uaddlv  h0, v0.8b
        urshr   d0, d0, 3
        dup     v0.8b, v0.b[0]
        ret

llvm output
foo:
        uaddlv  h0, v0.8b
        fmov    w8, s0
        fmov    d0, x8
        urshr   d0, d0, #3
        dup     v0.8b, v0.b[0]
        ret

There are copy instructions between gpr and fpr. We could remove them as below pattern.

def : Pat<(v1i64 (scalar_to_vector (i64 (zext (i32 (int_aarch64_neon_uaddlv (v8i8 V64:$Rn))))))),
          (INSERT_SUBREG (v1i64 (IMPLICIT_DEF)), (UADDLVv8i8v V64:$Rn), hsub)>;

With above pattern, llvm generates below output.

foo:                                    // @foo
        uaddlv  h0, v0.8b
        urshr   d0, d0, #3
        dup     v0.8b, v0.b[0]
        ret

The pattern could be too specific for this example. If you have other idea to generalize this case, please let me know.

[dmgreen]

Eli reminded us recently that there is a v16i8 variant of uaddlv that should be handled as well. Maybe the i16->i32 variants of uaddlv too, for matching the i32->i64 zext.

You may find that if this is able to be done in DAG (as in D159267), that the pattern is unneeded though. The transform might become simplified in the DAG before it gets to selection.

[jaykang10]

Updated pattern using UADDLV SDNode

[efriedma]

Do we need to be concerned at all about big-endian here? (Actually, also for D159267.) This is basically bitcasting from <2 x i32> to <1 x i64>.

We might want to consider teaching DAGCombine to optimize this sequence to an ISD::BITCAST instead of pattern-matching it late. Might unblock other optimizations? Maybe there aren't really any other optimizations we can do on a uaddlv, though.

[jaykang10]

Thanks for comment.

Do we need to be concerned at all about big-endian here? (Actually, also for D159267.) This is basically bitcasting from <2 x i32> to <1 x i64>.

I am not sure... It would be fine because compiler adds the rev instructions where they are needed for big-endian... but it could be wrong...
If you are concerned about something for big-endian, please let me know.

We might want to consider teaching DAGCombine to optimize this sequence to an ISD::BITCAST instead of pattern-matching it late. Might unblock other optimizations? Maybe there aren't really any other optimizations we can do on a uaddlv, though.

Let me try to detect the sequence and generate BITCAST with DAGCombine.

[jaykang10]

Following @efriedma's comment, folded sdnode sequence into bitcast in DAGCombine.

[efriedma]

LGTM

In D159265#4640515, @jaykang10 wrote:

Thanks for comment.

Do we need to be concerned at all about big-endian here? (Actually, also for D159267.) This is basically bitcasting from <2 x i32> to <1 x i64>.

I am not sure... It would be fine because compiler adds the rev instructions where they are needed for big-endian... but it could be wrong...
If you are concerned about something for big-endian, please let me know.

BITCAST itself should be fine; I meant, if we use a substitute sequence, we still need a REV, but there isn't any code to generate it.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

23159

I think we should be able to generalize this to other operations that only produce a result in the low element, but I guess we can leave that for a followup.

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

23169

I'm pretty sure this needs to be an AArch64ISD::NVCAST, not a BITCAST. The BITCAST will swap the 0th and 1st lanes into the i64, we need to keep them inorder.

[jaykang10]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

23169

Ah, I did not know we do not need rev instruction here.
Let me change BITCAST to NVCAST.
Thanks for letting me know.

[jaykang10]

Following @dmgreen's comment, changed BITCAST to NVCAST.

[jaykang10]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

23159

Yep, let's check it with other patch.

[dmgreen]

Thanks. It would be good to add other ops too if we can. Otherwise LGTM.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

#undef LCALLNAME5
   setTargetDAGCombine(ISD::GlobalAddress);
 
   setTargetDAGCombine(ISD::CTLZ);
 
   setTargetDAGCombine(ISD::VECREDUCE_AND);
   setTargetDAGCombine(ISD::VECREDUCE_OR);
   setTargetDAGCombine(ISD::VECREDUCE_XOR);
 
+  setTargetDAGCombine(ISD::SCALAR_TO_VECTOR);
+
   // In case of strict alignment, avoid an excessive number of byte wide stores.
   MaxStoresPerMemsetOptSize = 8;
   MaxStoresPerMemset =
       Subtarget->requiresStrictAlign() ? MaxStoresPerMemsetOptSize : 32;
 
   MaxGluedStoresPerMemcpy = 4;
   MaxStoresPerMemcpyOptSize = 4;
   MaxStoresPerMemcpy =
if (SDValue Val =
     return Val;
 
   if (SDValue Val = tryCombineMULLWithUZP1(N, DCI, DAG))
     return Val;
 
   return SDValue();
 }
 
+static SDValue
+performScalarToVectorCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI,
+                             SelectionDAG &DAG) {
+  // Let's do below transform.
+  //
+  //         t34: v4i32 = AArch64ISD::UADDLV t2
+  //       t35: i32 = extract_vector_elt t34, Constant:i64<0>
+  //     t7: i64 = zero_extend t35
+  //   t20: v1i64 = scalar_to_vector t7
+  // ==>
+  //      t34: v4i32 = AArch64ISD::UADDLV t2
+  //    t39: v2i32 = extract_subvector t34, Constant:i64<0>
+  //  t40: v1i64 = AArch64ISD::NVCAST t39
+  if (DCI.isBeforeLegalizeOps())
+    return SDValue();
+
+  EVT VT = N->getValueType(0);
+  if (VT != MVT::v1i64)
+    return SDValue();
+
+  SDValue ZEXT = N->getOperand(0);
+  if (ZEXT.getOpcode() != ISD::ZERO_EXTEND || ZEXT.getValueType() != MVT::i64)
+    return SDValue();
+
+  SDValue EXTRACT_VEC_ELT = ZEXT.getOperand(0);
+  if (EXTRACT_VEC_ELT.getOpcode() != ISD::EXTRACT_VECTOR_ELT ||
+      EXTRACT_VEC_ELT.getValueType() != MVT::i32)
+    return SDValue();
+
+  if (!isNullConstant(EXTRACT_VEC_ELT.getOperand(1)))
+    return SDValue();
+
+  SDValue UADDLV = EXTRACT_VEC_ELT.getOperand(0);
+  if (UADDLV.getOpcode() != AArch64ISD::UADDLV ||
+      UADDLV.getValueType() != MVT::v4i32 ||
+      UADDLV.getOperand(0).getValueType() != MVT::v8i8)
+    return SDValue();
+
+  // Let's generate new sequence with AArch64ISD::NVCAST.
+  SDLoc DL(N);
+  SDValue EXTRACT_SUBVEC =
+      DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, MVT::v2i32, UADDLV,
+                  DAG.getConstant(0, DL, MVT::i64));
+  SDValue NVCAST =
+      DAG.getNode(AArch64ISD::NVCAST, DL, MVT::v1i64, EXTRACT_SUBVEC);
+
+  return NVCAST;
+}
+
 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::VECREDUCE_AND:
case ISD::INTRINSIC_W_CHAIN:
     default:
       break;
     }
     break;
   case ISD::GlobalAddress:
     return performGlobalAddressCombine(N, DAG, Subtarget, getTargetMachine());
   case ISD::CTLZ:
     return performCTLZCombine(N, DAG, Subtarget);
+  case ISD::SCALAR_TO_VECTOR:
+    return performScalarToVectorCombine(N, DCI, DAG);
   }
   return SDValue();
 }
 
 // Check if the return value is used as only a return value, as otherwise
 // we can't perform a tail-call. In particular, we need to check for
 // target ISD nodes that are returns and any other "odd" constructs
 // that the generic analysis code won't necessarily catch.

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

 ; CHECK-NEXT:    fmov w0, s0
 ; CHECK-NEXT:    ret
 entry:
   %vaddlv.i = tail call i32 @llvm.aarch64.neon.uaddlv.i32.v16i8(<16 x i8> %a)
   %0 = and i32 %vaddlv.i, 65535
   ret i32 %0
 }
 
-define dso_local <8 x i8> @bar(<8 x i8> noundef %a) local_unnamed_addr #0 {
-; CHECK-LABEL: bar:
+define dso_local <8 x i8> @uaddlv_v8i8_dup(<8 x i8> %a) {
+; CHECK-LABEL: uaddlv_v8i8_dup:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    uaddlv h0, v0.8b
 ; CHECK-NEXT:    dup v0.8h, v0.h[0]
 ; CHECK-NEXT:    rshrn v0.8b, v0.8h, #3
 ; CHECK-NEXT:    ret
 entry:
   %vaddlv.i = tail call i32 @llvm.aarch64.neon.uaddlv.i32.v8i8(<8 x i8> %a)
   %0 = trunc i32 %vaddlv.i to i16
   %vecinit.i = insertelement <8 x i16> undef, i16 %0, i64 0
   %vecinit7.i = shufflevector <8 x i16> %vecinit.i, <8 x i16> poison, <8 x i32> zeroinitializer
   %vrshrn_n2 = tail call <8 x i8> @llvm.aarch64.neon.rshrn.v8i8(<8 x i16> %vecinit7.i, i32 3)
   ret <8 x i8> %vrshrn_n2
 }
 
 declare <8 x i8> @llvm.aarch64.neon.rshrn.v8i8(<8 x i16>, i32)
+
+declare i64 @llvm.aarch64.neon.urshl.i64(i64, i64)
+
+define <8 x i8> @uaddlv_v8i8_urshr(<8 x i8> %a) {
+; CHECK-LABEL: uaddlv_v8i8_urshr:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    uaddlv h0, v0.8b
+; CHECK-NEXT:    urshr d0, d0, #3
+; CHECK-NEXT:    dup v0.8b, v0.b[0]
+; CHECK-NEXT:    ret
+entry:
+  %vaddlv.i = tail call i32 @llvm.aarch64.neon.uaddlv.i32.v8i8(<8 x i8> %a)
+  %0 = and i32 %vaddlv.i, 65535
+  %conv = zext i32 %0 to i64
+  %vrshr_n = tail call i64 @llvm.aarch64.neon.urshl.i64(i64 %conv, i64 -3)
+  %conv1 = trunc i64 %vrshr_n to i8
+  %vecinit.i = insertelement <8 x i8> undef, i8 %conv1, i64 0
+  %vecinit7.i = shufflevector <8 x i8> %vecinit.i, <8 x i8> poison, <8 x i32> zeroinitializer
+  ret <8 x i8> %vecinit7.i
+}