[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 ↗

(On Diff #556143)

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 ↗

(On Diff #556143)

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 ↗

(On Diff #556143)

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 ↗

(On Diff #556143)

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/AArch64InstrInfo.td

def : Pat<(i32 (int_aarch64_neon_uaddlv (v8i16 (AArch64uaddlp (v16i8 V128:$op))))),
             ssub))>;
 
 def : Pat<(v4i32 (AArch64uaddlv (v8i8 V64:$Rn))),
           (v4i32 (SUBREG_TO_REG (i64 0), (UADDLVv8i8v V64:$Rn), hsub))>;
 
 def : Pat<(v4i32 (AArch64uaddlv (v16i8 V128:$Rn))),
           (v4i32 (SUBREG_TO_REG (i64 0), (UADDLVv16i8v V128:$Rn), hsub))>;
 
+def : Pat<(v1i64 (scalar_to_vector
+                   (i64 (zext (i32 (vector_extract (v4i32 (AArch64uaddlv (v8i8 V64:$Rn))),
+                                                   (i64 0))))))),
+          (v1i64 (SUBREG_TO_REG (i64 0), (UADDLVv8i8v V64:$Rn), hsub))>;
+
 // Patterns for across-vector intrinsics, that have a node equivalent, that
 // returns a vector (with only the low lane defined) instead of a scalar.
 // In effect, opNode is the same as (scalar_to_vector (IntNode)).
 multiclass SIMDAcrossLanesIntrinsic<string baseOpc,
                                     SDPatternOperator opNode> {
 // If a lane instruction caught the vector_extract around opNode, we can
 // directly match the latter to the instruction.
 def : Pat<(v8i8 (opNode V64:$Rn)),

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
+}