[AArch64] Change the type of i64 neon shifts to v1i64

This alters the lowering of shifts by a constant, so that the type is lowered
to a v1i64 instead of a i64. This helps communicate that the type will live in
a neon register, and can help clean up surrounding code. Note this is only
currently for the scalar shifts of a constant that go through the nodes in
tryCombineShiftImm.

ssra instructions are no longer being recognized in places, but that can be
cleaned up in a followup patch that combines the i64 add into a v1i64 add.

Differential Revision: https://reviews.llvm.org/D148309

2 files changed, 33 insertions, 26 deletions

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 56be99ee5ff9..4f0be5f98b0a 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -18384,14 +18384,28 @@ static SDValue tryCombineShiftImm(unsigned IID, SDNode *N, SelectionDAG &DAG) {
     break;
   }
 
+  EVT VT = N->getValueType(0);
+  SDValue Op = N->getOperand(1);
+  SDLoc dl(N);
+  if (VT == MVT::i64) {
+    Op = DAG.getNode(ISD::SCALAR_TO_VECTOR, dl, MVT::v1i64, Op);
+    VT = MVT::v1i64;
+  }
+
   if (IsRightShift && ShiftAmount <= -1 && ShiftAmount >= -(int)ElemBits) {
-    SDLoc dl(N);
-    return DAG.getNode(Opcode, dl, N->getValueType(0), N->getOperand(1),
-                       DAG.getConstant(-ShiftAmount, dl, MVT::i32));
+    Op = DAG.getNode(Opcode, dl, VT, Op,
+                     DAG.getConstant(-ShiftAmount, dl, MVT::i32));
+    if (N->getValueType(0) == MVT::i64)
+      Op = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, dl, MVT::i64, Op,
+                       DAG.getConstant(0, dl, MVT::i64));
+    return Op;
   } else if (!IsRightShift && ShiftAmount >= 0 && ShiftAmount < ElemBits) {
-    SDLoc dl(N);
-    return DAG.getNode(Opcode, dl, N->getValueType(0), N->getOperand(1),
-                       DAG.getConstant(ShiftAmount, dl, MVT::i32));
+    Op = DAG.getNode(Opcode, dl, VT, Op,
+                     DAG.getConstant(ShiftAmount, dl, MVT::i32));
+    if (N->getValueType(0) == MVT::i64)
+      Op = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, dl, MVT::i64, Op,
+                       DAG.getConstant(0, dl, MVT::i64));
+    return Op;
   }
 
   return SDValue();
diff --git a/llvm/test/CodeGen/AArch64/arm64-vshift.ll b/llvm/test/CodeGen/AArch64/arm64-vshift.ll
index dca690e1f693..082268a836ad 100644
--- a/llvm/test/CodeGen/AArch64/arm64-vshift.ll
+++ b/llvm/test/CodeGen/AArch64/arm64-vshift.ll
@@ -83,8 +83,7 @@ define i64 @sqshl_scalar(ptr %A, ptr %B) nounwind {
 define i64 @sqshl_scalar_constant(ptr %A) nounwind {
 ; CHECK-LABEL: sqshl_scalar_constant:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ldr x8, [x0]
-; CHECK-NEXT:    fmov d0, x8
+; CHECK-NEXT:    ldr d0, [x0]
 ; CHECK-NEXT:    sqshl d0, d0, #1
 ; CHECK-NEXT:    fmov x0, d0
 ; CHECK-NEXT:    ret
@@ -279,8 +278,7 @@ define i64 @uqshl_scalar(ptr %A, ptr %B) nounwind {
 define i64 @uqshl_scalar_constant(ptr %A) nounwind {
 ; CHECK-LABEL: uqshl_scalar_constant:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ldr x8, [x0]
-; CHECK-NEXT:    fmov d0, x8
+; CHECK-NEXT:    ldr d0, [x0]
 ; CHECK-NEXT:    uqshl d0, d0, #1
 ; CHECK-NEXT:    fmov x0, d0
 ; CHECK-NEXT:    ret
@@ -1039,8 +1037,7 @@ define <1 x i64> @urshr1d(ptr %A) nounwind {
 define i64 @urshr_scalar(ptr %A) nounwind {
 ; CHECK-LABEL: urshr_scalar:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ldr x8, [x0]
-; CHECK-NEXT:    fmov d0, x8
+; CHECK-NEXT:    ldr d0, [x0]
 ; CHECK-NEXT:    urshr d0, d0, #1
 ; CHECK-NEXT:    fmov x0, d0
 ; CHECK-NEXT:    ret
@@ -1140,8 +1137,7 @@ define <1 x i64> @srshr1d(ptr %A) nounwind {
 define i64 @srshr_scalar(ptr %A) nounwind {
 ; CHECK-LABEL: srshr_scalar:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ldr x8, [x0]
-; CHECK-NEXT:    fmov d0, x8
+; CHECK-NEXT:    ldr d0, [x0]
 ; CHECK-NEXT:    srshr d0, d0, #1
 ; CHECK-NEXT:    fmov x0, d0
 ; CHECK-NEXT:    ret
@@ -1241,8 +1237,7 @@ define <1 x i64> @sqshlu1d_constant(ptr %A) nounwind {
 define i64 @sqshlu_i64_constant(ptr %A) nounwind {
 ; CHECK-LABEL: sqshlu_i64_constant:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ldr x8, [x0]
-; CHECK-NEXT:    fmov d0, x8
+; CHECK-NEXT:    ldr d0, [x0]
 ; CHECK-NEXT:    sqshlu d0, d0, #1
 ; CHECK-NEXT:    fmov x0, d0
 ; CHECK-NEXT:    ret
@@ -2737,12 +2732,11 @@ define <1 x i64> @ursra1d(ptr %A, ptr %B) nounwind {
 define i64 @ursra_scalar(ptr %A, ptr %B) nounwind {
 ; CHECK-LABEL: ursra_scalar:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ldr x8, [x0]
+; CHECK-NEXT:    ldr d0, [x0]
 ; CHECK-NEXT:    ldr x9, [x1]
-; CHECK-NEXT:    fmov d1, x8
-; CHECK-NEXT:    fmov d0, x9
-; CHECK-NEXT:    ursra d0, d1, #1
-; CHECK-NEXT:    fmov x0, d0
+; CHECK-NEXT:    urshr d0, d0, #1
+; CHECK-NEXT:    fmov x8, d0
+; CHECK-NEXT:    add x0, x8, x9
 ; CHECK-NEXT:    ret
   %tmp1 = load i64, ptr %A
   %tmp3 = call i64 @llvm.aarch64.neon.urshl.i64(i64 %tmp1, i64 -1)
@@ -2866,12 +2860,11 @@ define <1 x i64> @srsra1d(ptr %A, ptr %B) nounwind {
 define i64 @srsra_scalar(ptr %A, ptr %B) nounwind {
 ; CHECK-LABEL: srsra_scalar:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ldr x8, [x0]
+; CHECK-NEXT:    ldr d0, [x0]
 ; CHECK-NEXT:    ldr x9, [x1]
-; CHECK-NEXT:    fmov d1, x8
-; CHECK-NEXT:    fmov d0, x9
-; CHECK-NEXT:    srsra d0, d1, #1
-; CHECK-NEXT:    fmov x0, d0
+; CHECK-NEXT:    srshr d0, d0, #1
+; CHECK-NEXT:    fmov x8, d0
+; CHECK-NEXT:    add x0, x8, x9
 ; CHECK-NEXT:    ret
   %tmp1 = load i64, ptr %A
   %tmp3 = call i64 @llvm.aarch64.neon.srshl.i64(i64 %tmp1, i64 -1)