[ADT] Add fshl/fshr operations to APInt

[ckoparkar]

These operations are required for #153151.

[ckoparkar]

/cc @kuhar @RKSimon

[llvmbot]

@llvm/pr-subscribers-llvm-adt

@llvm/pr-subscribers-llvm-support

Author: Chaitanya Koparkar (ckoparkar)

Changes

These operations are required for #153151.

---

Full diff: https://github.com/llvm/llvm-project/pull/153790.diff

3 Files Affected:

• (modified) llvm/include/llvm/ADT/APSInt.h (+20)
• (modified) llvm/lib/Support/APSInt.cpp (+22)
• (modified) llvm/unittests/ADT/APSIntTest.cpp (+85)

diff --git a/llvm/include/llvm/ADT/APSInt.h b/llvm/include/llvm/ADT/APSInt.h
index 88a7a6e71c817..74d4daadcc8f1 100644
--- a/llvm/include/llvm/ADT/APSInt.h
+++ b/llvm/include/llvm/ADT/APSInt.h
@@ -152,6 +152,9 @@ class [[nodiscard]] APSInt : public APInt {
   APSInt operator>>(unsigned Amt) const {
     return IsUnsigned ? APSInt(lshr(Amt), true) : APSInt(ashr(Amt), false);
   }
+  APSInt operator>>(const APSInt &Amt) const {
+    return IsUnsigned ? APSInt(lshr(Amt), true) : APSInt(ashr(Amt), false);
+  }
   APSInt &operator>>=(unsigned Amt) {
     if (IsUnsigned)
       lshrInPlace(Amt);
@@ -211,6 +214,9 @@ class [[nodiscard]] APSInt : public APInt {
   APSInt operator<<(unsigned Bits) const {
     return APSInt(static_cast<const APInt &>(*this) << Bits, IsUnsigned);
   }
+  APSInt operator<<(const APSInt &Amt) const {
+    return APSInt(static_cast<const APInt &>(*this) << Amt, IsUnsigned);
+  }
   APSInt &operator<<=(unsigned Amt) {
     static_cast<APInt &>(*this) <<= Amt;
     return *this;
@@ -387,6 +393,20 @@ template <> struct DenseMapInfo<APSInt, void> {
   }
 };
 
+namespace APSIntOps {
+
+/// Perform a funnel shift left.
+///
+/// fshl(X,Y,Z): (X << (Z % BW)) | (Y >> (BW - (Z % BW)))
+LLVM_ABI APSInt fshl(const APSInt &Hi, const APSInt &Lo, const APSInt &Shift);
+
+/// Perform a funnel shift right.
+///
+/// fshr(X,Y,Z): (X << (BW - (Z % BW))) | (Y >> (Z % BW))
+LLVM_ABI APSInt fshr(const APSInt &Hi, const APSInt &Lo, const APSInt &Shift);
+
+} // namespace APSIntOps
+
 } // end namespace llvm
 
 #endif
diff --git a/llvm/lib/Support/APSInt.cpp b/llvm/lib/Support/APSInt.cpp
index 5a9f44f304a27..8e2dc55bf8bca 100644
--- a/llvm/lib/Support/APSInt.cpp
+++ b/llvm/lib/Support/APSInt.cpp
@@ -41,3 +41,25 @@ void APSInt::Profile(FoldingSetNodeID& ID) const {
   ID.AddInteger((unsigned) (IsUnsigned ? 1 : 0));
   APInt::Profile(ID);
 }
+
+APSInt llvm::APSIntOps::fshl(const APSInt &Hi, const APSInt &Lo,
+                             const APSInt &Shift) {
+  bool IsUnsigned = Hi.isUnsigned();
+  APSInt BitWidth(APInt(Hi.getBitWidth(),
+                        static_cast<uint64_t>(Hi.getBitWidth())),
+                  IsUnsigned);
+  return APSInt((Hi << (Shift % BitWidth)) |
+                (Lo >> (BitWidth - (Shift % BitWidth))),
+                IsUnsigned);
+}
+
+APSInt llvm::APSIntOps::fshr(const APSInt &Hi, const APSInt &Lo,
+                             const APSInt &Shift) {
+  bool IsUnsigned = Hi.isUnsigned();
+  APSInt BitWidth(APInt(Hi.getBitWidth(),
+                        static_cast<uint64_t>(Hi.getBitWidth())),
+                  IsUnsigned);
+  return APSInt((Hi << (BitWidth - (Shift % BitWidth))) |
+                (Lo >> (Shift % BitWidth)),
+                IsUnsigned);
+}
diff --git a/llvm/unittests/ADT/APSIntTest.cpp b/llvm/unittests/ADT/APSIntTest.cpp
index 2d2a64433da94..be303a99b2419 100644
--- a/llvm/unittests/ADT/APSIntTest.cpp
+++ b/llvm/unittests/ADT/APSIntTest.cpp
@@ -285,4 +285,89 @@ TEST(APSIntTest, UnsignedHighBit) {
   EXPECT_TRUE(CharMax.isStrictlyPositive());
 }
 
+// Right shift of unsigned numbers translates to a logical shift.
+TEST(APSIntTest, RightShiftUnsigned) {
+  // Right shift of a negative number.
+  const APInt neg_one(128, static_cast<uint64_t>(-1), /*isSigned*/true);
+  const APSInt neg_one_unsigned(APSInt(neg_one, /*isUnsigned*/true));
+  EXPECT_EQ(0, neg_one_unsigned >> 128);
+
+  APSInt i256(APSInt(APInt::getHighBitsSet(256, 2)));
+  i256 >>= 1;
+  EXPECT_EQ(1U, i256.countl_zero());
+  EXPECT_EQ(253U, i256.countr_zero());
+  EXPECT_EQ(2U, i256.popcount());
+
+  i256 >>= 62;
+  EXPECT_EQ(63U, i256.countl_zero());
+  EXPECT_EQ(191U, i256.countr_zero());
+  EXPECT_EQ(2U, i256.popcount());
+}
+
+// Right shift of signed numbers translates to a arithmetic shift.
+TEST(APSIntTest, RightShiftSigned) {
+  // Right shift of a negative number.
+  const APInt neg_one = APInt(64, static_cast<uint64_t>(-1), /*isSigned*/true);
+  const APSInt neg_one_signed(APSInt(neg_one, /*isUnsigned*/false));
+  EXPECT_EQ(neg_one, neg_one_signed >> 7);
+
+  APSInt i72(APSInt(APInt::getHighBitsSet(72, 1), /*isUnsigned*/false));
+  i72 >>= 46;
+  EXPECT_EQ(47U, i72.countl_one());
+  EXPECT_EQ(25U, i72.countr_zero());
+  EXPECT_EQ(47U, i72.popcount());
+
+  // Ensure we handle large shifts of multi-word.
+  const APSInt signmin128(APSInt(APInt::getSignedMinValue(128), /*isUnsigned*/false));
+  EXPECT_TRUE((signmin128 >> 128).isAllOnes());
+
+  // Ensure we handle large shifts of multi-word.
+  const APSInt umax128(APSInt(APInt::getSignedMaxValue(128), /*isUnsigned*/false));
+  EXPECT_EQ(0, umax128 >> 128);
+}
+
+TEST(APSIntTest, Fshl) {
+  // Unsigned
+  EXPECT_EQ(APSIntOps::fshl(APSInt(APInt(8, 0)), APSInt(APInt(8, 255)),
+                            APSInt(APInt(8, 8))).getExtValue(), 0);
+  EXPECT_EQ(APSIntOps::fshl(APSInt(APInt(8, 255)), APSInt(APInt(8, 0)),
+                            APSInt(APInt(8, 8))).getExtValue(), 255);
+  EXPECT_EQ(APSIntOps::fshl(APSInt(APInt(8, 255)), APSInt(APInt(8, 0)),
+                            APSInt(APInt(8, 15))).getExtValue(), 128);
+  EXPECT_EQ(APSIntOps::fshl(APSInt(APInt(8, 15)), APSInt(APInt(8, 15)),
+                            APSInt(APInt(8, 11))).getExtValue(), 120);
+  EXPECT_EQ(APSIntOps::fshl(APSInt(APInt(8, 2)), APSInt(APInt(8, 1)),
+                            APSInt(APInt(8, 3))).getExtValue(), 16);
+  // Signed
+  EXPECT_EQ(APSIntOps::fshl(APSInt(APInt(32, 0), false),
+                            APSInt(APInt(32, 2147483647), false),
+                            APSInt(APInt(32, 32), false)).getExtValue(), 0);
+  EXPECT_EQ(APSIntOps::fshl(APSInt(APInt(64, 1), false),
+                            APSInt(APInt(64, 2), false),
+                            APSInt(APInt(64, 3), false)).getExtValue(), 8);
+}
+
+TEST(APSIntTest, Fshr) {
+  // Unsigned
+  EXPECT_EQ(APSIntOps::fshr(APSInt(APInt(8, 0)), APSInt(APInt(8, 255)),
+                            APSInt(APInt(8, 8))).getExtValue(), 255);
+  EXPECT_EQ(APSIntOps::fshr(APSInt(APInt(8, 255)), APSInt(APInt(8, 0)),
+                            APSInt(APInt(8, 8))).getExtValue(), 0);
+  EXPECT_EQ(APSIntOps::fshr(APSInt(APInt(8, 255)), APSInt(APInt(8, 0)),
+                            APSInt(APInt(8, 15))).getExtValue(), 254);
+  EXPECT_EQ(APSIntOps::fshr(APSInt(APInt(8, 15)), APSInt(APInt(8, 15)),
+                            APSInt(APInt(8, 11))).getExtValue(), 225);
+  EXPECT_EQ(APSIntOps::fshr(APSInt(APInt(8, 1)), APSInt(APInt(8, 2)),
+                            APSInt(APInt(8, 3))).getExtValue(), 32);
+  // Signed
+  EXPECT_EQ(APSIntOps::fshr(APSInt(APInt(64, 0), false),
+                            APSInt(APInt(64, 9223372036854775807), false),
+                            APSInt(APInt(64, 64), false)).getExtValue(),
+            9223372036854775807);
+  EXPECT_EQ(APSIntOps::fshr(APSInt(APInt(64, 1), false),
+                            APSInt(APInt(64, 2), false),
+                            APSInt(APInt(64, 3), false)).getExtValue(),
+            2305843009213693952);
+}
+
 } // end anonymous namespace

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[XChy]

ping.

[RKSimon]

LGTM - cheers

[kuhar]

LGTM % request for further clarify semantics

[ckoparkar]

All the comments have been addressed. I think this is safe to be merged now. Thanks everyone :-)

[topperc]

LGTM other than the question about the constructor call

[RKSimon]

LGTM