[AArch64] Add ComputeNumSignBits for VASHR. (llvm#113957)

As with a normal ISD::SRA node, they take the number of sign bits of the
incoming value and increase it by the shifted amount.

Author: davemgreen

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -2536,6 +2536,11 @@ unsigned AArch64TargetLowering::ComputeNumSignBitsForTargetNode(
     case AArch64ISD::FCMLTz:
       // Compares return either 0 or all-ones
       return VTBits;
+    case AArch64ISD::VASHR: {
+      unsigned Tmp =
+          DAG.ComputeNumSignBits(Op.getOperand(0), DemandedElts, Depth + 1);
+      return std::min<uint64_t>(Tmp + Op.getConstantOperandVal(1), VTBits);
+    }
   }
 
   return 1;

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

@@ -3560,4 +3560,16 @@ entry:
   ret <4 x i16> %vrshrn_n1
 }
 
+define <8 x i16> @signbits_vashr(<8 x i16> %a)  {
+; CHECK-LABEL: signbits_vashr:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    sshr.8h v0, v0, #8
+; CHECK-NEXT:    sshr.8h v0, v0, #9
+; CHECK-NEXT:    ret
+  %b = call <8 x i16> @llvm.aarch64.neon.sshl.v8i16(<8 x i16> %a, <8 x i16> <i16 -8, i16 -8, i16 -8, i16 -8, i16 -8, i16 -8, i16 -8, i16 -8>)
+  %c = call <8 x i16> @llvm.aarch64.neon.sshl.v8i16(<8 x i16> %b, <8 x i16> <i16 -9, i16 -9, i16 -9, i16 -9, i16 -9, i16 -9, i16 -9, i16 -9>)
+  %d = ashr <8 x i16> %c, <i16 7, i16 7, i16 7, i16 7, i16 7, i16 7, i16 7, i16 7>
+  ret <8 x i16> %d
+}
+
 declare <2 x i64> @llvm.aarch64.neon.addp.v2i64(<2 x i64>, <2 x i64>)

llvm/unittests/CodeGen/AArch64SelectionDAGTest.cpp

@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "../lib/Target/AArch64/AArch64ISelLowering.h"
 #include "llvm/Analysis/MemoryLocation.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/AsmParser/Parser.h"
@@ -167,6 +168,18 @@ TEST_F(AArch64SelectionDAGTest, ComputeNumSignBits_EXTRACT_SUBVECTOR) {
   EXPECT_EQ(DAG->ComputeNumSignBits(Op, DemandedElts), 7u);
 }
 
+TEST_F(AArch64SelectionDAGTest, ComputeNumSignBits_VASHR) {
+  SDLoc Loc;
+  auto VecVT = MVT::v8i8;
+  auto Shift = DAG->getConstant(4, Loc, MVT::i32);
+  auto Vec0 = DAG->getConstant(1, Loc, VecVT);
+  auto Op1 = DAG->getNode(AArch64ISD::VASHR, Loc, VecVT, Vec0, Shift);
+  EXPECT_EQ(DAG->ComputeNumSignBits(Op1), 8u);
+  auto VecA = DAG->getConstant(0xaa, Loc, VecVT);
+  auto Op2 = DAG->getNode(AArch64ISD::VASHR, Loc, VecVT, VecA, Shift);
+  EXPECT_EQ(DAG->ComputeNumSignBits(Op2), 5u);
+}
+
 TEST_F(AArch64SelectionDAGTest, SimplifyDemandedVectorElts_EXTRACT_SUBVECTOR) {
   TargetLowering TL(*TM);