[AArch64] Remove vector shift instrinsic with shift amount zero

[jaykang10]

Summary

It looks gcc folds vector shift intrinsic with zero shift amount from below example.

#include <arm_neon.h>

inline void foo(int64x2_t a, int64x2_t b, int64_t *dst, int df) {
    int64x2_t df_s64 = vdupq_n_s64(df);
    a = vpaddq_s64(a, b); 
    a = vshlq_s64(a, df_s64);
    vst1q_s64(dst, a); 
} 

void bar(int64x2_t a, int64x2_t b, int64_t *dst) {
    foo(a, b, dst, 0); 
}

gcc output
bar:
	addp	v0.2d, v0.2d, v1.2d
	str	q0, [x0]
	ret

llvm output
bar:
	addp	v0.2d, v0.2d, v1.2d
	shl	v0.2d, v0.2d, #0
	str	q0, [x0]
	ret

It looks llvm's AArch64 target lowers the intrinsic to target custom node in SelectionDAG and is missing to fold the custom node with zero shift amount.
With this patch, the llvm output is as below.

bar:
	addp	v0.2d, v0.2d, v1.2d
	str	q0, [x0]
	ret

[dmgreen]

Thanks - looks mostly good to me. Like I said before, there might be some advantage to doing this in instcombine in order to doing the transform earlier, but this will be useful in DAG too.

llvm/test/CodeGen/AArch64/arm64-vshift.ll

3491

Apparently this one is not correct, as the sqshlu will round the input even with a zero shift. The others look OK.

[jaykang10]

llvm/test/CodeGen/AArch64/arm64-vshift.ll

3491

Ah... I did not know that.
If possible, can you let me know where I can find that the sqshlu will round the input even with a zero shift please?

[jaykang10]

Following @dmgreen's comment, updated code.

[dmgreen]

Thanks. LGTM

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (!BVN->isConstantSplat(SplatValue, SplatUndef, SplatBitSize,
       return SDValue();
 
     ShiftAmount = SplatValue.getSExtValue();
   } else if (ConstantSDNode *CVN = dyn_cast<ConstantSDNode>(N->getOperand(2))) {
     ShiftAmount = CVN->getSExtValue();
   } else
     return SDValue();
 
+  // If the shift amount is zero, remove the shift intrinsic.
+  if (ShiftAmount == 0 && IID != Intrinsic::aarch64_neon_sqshlu)
+    return N->getOperand(1);
+
   unsigned Opcode;
   bool IsRightShift;
   switch (IID) {
   default:
     llvm_unreachable("Unknown shift intrinsic");
   case Intrinsic::aarch64_neon_sqshl:
     Opcode = AArch64ISD::SQSHL_I;
     IsRightShift = false;

llvm/test/CodeGen/AArch64/arm64-vshift.ll

 ; CHECK-LABEL: ashr_v1i64:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    neg d1, d1
 ; CHECK-NEXT:    sshl d0, d0, d1
 ; CHECK-NEXT:    ret
   %c = ashr <1 x i64> %a, %b
   ret <1 x i64> %c
 }
+
+define void @sqshl_zero_shift_amount(<2 x i64> %a, <2 x i64> %b, ptr %dst) {
+; CHECK-LABEL: sqshl_zero_shift_amount:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    addp.2d v0, v0, v1
+; CHECK-NEXT:    str q0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  %vpaddq_v2.i.i = tail call <2 x i64> @llvm.aarch64.neon.addp.v2i64(<2 x i64> %a, <2 x i64> %b)
+  %vshlq_v2.i.i = tail call <2 x i64> @llvm.aarch64.neon.sqshl.v2i64(<2 x i64> %vpaddq_v2.i.i, <2 x i64> zeroinitializer)
+  store <2 x i64> %vshlq_v2.i.i, ptr %dst, align 8
+  ret void
+}
+
+define void @uqshl_zero_shift_amount(<2 x i64> %a, <2 x i64> %b, ptr %dst) {
+; CHECK-LABEL: uqshl_zero_shift_amount:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    addp.2d v0, v0, v1
+; CHECK-NEXT:    str q0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  %vpaddq_v2.i.i = tail call <2 x i64> @llvm.aarch64.neon.addp.v2i64(<2 x i64> %a, <2 x i64> %b)
+  %vshlq_v2.i.i = tail call <2 x i64> @llvm.aarch64.neon.uqshl.v2i64(<2 x i64> %vpaddq_v2.i.i, <2 x i64> zeroinitializer)
+  store <2 x i64> %vshlq_v2.i.i, ptr %dst, align 8
+  ret void
+}
+
+define void @srshl_zero_shift_amount(<2 x i64> %a, <2 x i64> %b, ptr %dst) {
+; CHECK-LABEL: srshl_zero_shift_amount:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    addp.2d v0, v0, v1
+; CHECK-NEXT:    str q0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  %vpaddq_v2.i.i = tail call <2 x i64> @llvm.aarch64.neon.addp.v2i64(<2 x i64> %a, <2 x i64> %b)
+  %vshlq_v2.i.i = tail call <2 x i64> @llvm.aarch64.neon.srshl.v2i64(<2 x i64> %vpaddq_v2.i.i, <2 x i64> zeroinitializer)
+  store <2 x i64> %vshlq_v2.i.i, ptr %dst, align 8
+  ret void
+}
+
+define void @urshl_zero_shift_amount(<2 x i64> %a, <2 x i64> %b, ptr %dst) {
+; CHECK-LABEL: urshl_zero_shift_amount:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    addp.2d v0, v0, v1
+; CHECK-NEXT:    str q0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  %vpaddq_v2.i.i = tail call <2 x i64> @llvm.aarch64.neon.addp.v2i64(<2 x i64> %a, <2 x i64> %b)
+  %vshlq_v2.i.i = tail call <2 x i64> @llvm.aarch64.neon.urshl.v2i64(<2 x i64> %vpaddq_v2.i.i, <2 x i64> zeroinitializer)
+  store <2 x i64> %vshlq_v2.i.i, ptr %dst, align 8
+  ret void
+}
+
+define void @sqshlu_zero_shift_amount(<2 x i64> %a, <2 x i64> %b, ptr %dst) {
+; CHECK-LABEL: sqshlu_zero_shift_amount:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    addp.2d v0, v0, v1
+; CHECK-NEXT:    sqshlu.2d v0, v0, #0
+; CHECK-NEXT:    str q0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  %vpaddq_v2.i.i = tail call <2 x i64> @llvm.aarch64.neon.addp.v2i64(<2 x i64> %a, <2 x i64> %b)
+  %vshlq_v2.i.i = tail call <2 x i64> @llvm.aarch64.neon.sqshlu.v2i64(<2 x i64> %vpaddq_v2.i.i, <2 x i64> zeroinitializer)
+  store <2 x i64> %vshlq_v2.i.i, ptr %dst, align 8
+  ret void
+}
+
+define void @sshl_zero_shift_amount(<2 x i64> %a, <2 x i64> %b, ptr %dst) {
+; CHECK-LABEL: sshl_zero_shift_amount:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    addp.2d v0, v0, v1
+; CHECK-NEXT:    str q0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  %vpaddq_v2.i.i = tail call <2 x i64> @llvm.aarch64.neon.addp.v2i64(<2 x i64> %a, <2 x i64> %b)
+  %vshlq_v2.i.i = tail call <2 x i64> @llvm.aarch64.neon.sshl.v2i64(<2 x i64> %vpaddq_v2.i.i, <2 x i64> zeroinitializer)
+  store <2 x i64> %vshlq_v2.i.i, ptr %dst, align 8
+  ret void
+}
+
+define void @ushl_zero_shift_amount(<2 x i64> %a, <2 x i64> %b, ptr %dst) {
+; CHECK-LABEL: ushl_zero_shift_amount:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    addp.2d v0, v0, v1
+; CHECK-NEXT:    str q0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  %vpaddq_v2.i.i = tail call <2 x i64> @llvm.aarch64.neon.addp.v2i64(<2 x i64> %a, <2 x i64> %b)
+  %vshlq_v2.i.i = tail call <2 x i64> @llvm.aarch64.neon.ushl.v2i64(<2 x i64> %vpaddq_v2.i.i, <2 x i64> zeroinitializer)
+  store <2 x i64> %vshlq_v2.i.i, ptr %dst, align 8
+  ret void
+}
+
+declare <2 x i64> @llvm.aarch64.neon.addp.v2i64(<2 x i64>, <2 x i64>)