[llvm] [AArch64] Add computeKnownBits unit test coverage for AArch64ISD::VASHR/VLSHR/VSHL (PR #156631)

[Simon Pilgrim via llvm-commits]

https://github.com/RKSimon updated https://github.com/llvm/llvm-project/pull/156631

>From 4b3a77f785a6ee79e40baf2ad580d379ed341c25 Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev at redking.me.uk>
Date: Wed, 3 Sep 2025 10:36:12 +0100
Subject: [PATCH] [AArch64] Add computeKnownBits unit test coverage for
 AArch64ISD::VASHR/VLSHR/VSHL

---
 .../AArch64/AArch64SelectionDAGTest.cpp       | 60 +++++++++++++++++++
 1 file changed, 60 insertions(+)

diff --git a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
index 675fdc7b0db3b..18c8d4a69e7a8 100644
--- a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
+++ b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
@@ -505,6 +505,66 @@ TEST_F(AArch64SelectionDAGTest, ComputeKnownBits_USUBO_CARRY) {
   EXPECT_EQ(Known.One, APInt(8, 0x31));
 }
 
+// Piggy-backing on the AArch64 tests to verify SelectionDAG::computeKnownBits.
+TEST_F(AArch64SelectionDAGTest, ComputeKnownBits_VASHR) {
+  SDLoc Loc;
+  KnownBits Known;
+  auto VecVT = MVT::v8i8;
+  auto Shift0 = DAG->getConstant(4, Loc, MVT::i32);
+  auto Vec0 = DAG->getConstant(0x80, Loc, VecVT);
+  auto Op0 = DAG->getNode(AArch64ISD::VASHR, Loc, VecVT, Vec0, Shift0);
+  Known = DAG->computeKnownBits(Op0);
+  EXPECT_EQ(Known.Zero, APInt(8, 0x07));
+  EXPECT_EQ(Known.One, APInt(8, 0xF8));
+
+  auto Shift1 = DAG->getConstant(7, Loc, MVT::i32);
+  auto Vec1 = DAG->getConstant(0xF7, Loc, VecVT);
+  auto Op1 = DAG->getNode(AArch64ISD::VASHR, Loc, VecVT, Vec1, Shift1);
+  Known = DAG->computeKnownBits(Op1);
+  EXPECT_EQ(Known.Zero, APInt(8, 0x00));
+  EXPECT_EQ(Known.One, APInt(8, 0xFF));
+}
+
+// Piggy-backing on the AArch64 tests to verify SelectionDAG::computeKnownBits.
+TEST_F(AArch64SelectionDAGTest, ComputeKnownBits_VLSHR) {
+  SDLoc Loc;
+  KnownBits Known;
+  auto VecVT = MVT::v8i8;
+  auto Shift0 = DAG->getConstant(4, Loc, MVT::i32);
+  auto Vec0 = DAG->getConstant(0x80, Loc, VecVT);
+  auto Op0 = DAG->getNode(AArch64ISD::VLSHR, Loc, VecVT, Vec0, Shift0);
+  Known = DAG->computeKnownBits(Op0);
+  EXPECT_EQ(Known.Zero, APInt(8, 0xF7));
+  EXPECT_EQ(Known.One, APInt(8, 0x08));
+
+  auto Shift1 = DAG->getConstant(7, Loc, MVT::i32);
+  auto Vec1 = DAG->getConstant(0xF7, Loc, VecVT);
+  auto Op1 = DAG->getNode(AArch64ISD::VLSHR, Loc, VecVT, Vec1, Shift1);
+  Known = DAG->computeKnownBits(Op1);
+  EXPECT_EQ(Known.Zero, APInt(8, 0xFE));
+  EXPECT_EQ(Known.One, APInt(8, 0x1));
+}
+
+// Piggy-backing on the AArch64 tests to verify SelectionDAG::computeKnownBits.
+TEST_F(AArch64SelectionDAGTest, ComputeKnownBits_VSHL) {
+  SDLoc Loc;
+  KnownBits Known;
+  auto VecVT = MVT::v8i8;
+  auto Shift0 = DAG->getConstant(4, Loc, MVT::i32);
+  auto Vec0 = DAG->getConstant(0x02, Loc, VecVT);
+  auto Op0 = DAG->getNode(AArch64ISD::VSHL, Loc, VecVT, Vec0, Shift0);
+  Known = DAG->computeKnownBits(Op0);
+  EXPECT_EQ(Known.Zero, APInt(8, 0xDF));
+  EXPECT_EQ(Known.One, APInt(8, 0x20));
+
+  auto Shift1 = DAG->getConstant(7, Loc, MVT::i32);
+  auto Vec1 = DAG->getConstant(0xF7, Loc, VecVT);
+  auto Op1 = DAG->getNode(AArch64ISD::VSHL, Loc, VecVT, Vec1, Shift1);
+  Known = DAG->computeKnownBits(Op1);
+  EXPECT_EQ(Known.Zero, APInt(8, 0x7F));
+  EXPECT_EQ(Known.One, APInt(8, 0x80));
+}
+
 TEST_F(AArch64SelectionDAGTest, isSplatValue_Fixed_BUILD_VECTOR) {
   TargetLowering TL(*TM);