[llvm] [ADT] Add fshl/fshr operations to APSInt (PR #153790)

[Chaitanya Koparkar via llvm-commits]

https://github.com/ckoparkar updated https://github.com/llvm/llvm-project/pull/153790

>From 526ee9dc5e20f0f6d3bd2237fc4857aba5a8dd5e Mon Sep 17 00:00:00 2001
From: Chaitanya Koparkar <ckoparkar at gmail.com>
Date: Fri, 15 Aug 2025 07:24:12 -0400
Subject: [PATCH] [ADT] Add fshl/fshr operations to APInt

---
 llvm/include/llvm/ADT/APInt.h    | 11 ++++++
 llvm/lib/Support/APInt.cpp       | 16 ++++++++
 llvm/unittests/ADT/APIntTest.cpp | 66 ++++++++++++++++++++++++++++++++
 3 files changed, 93 insertions(+)

diff --git a/llvm/include/llvm/ADT/APInt.h b/llvm/include/llvm/ADT/APInt.h
index 4d60ad4459ad6..a5f0fe54dfd5b 100644
--- a/llvm/include/llvm/ADT/APInt.h
+++ b/llvm/include/llvm/ADT/APInt.h
@@ -2410,6 +2410,17 @@ LLVM_ABI std::optional<unsigned> GetMostSignificantDifferentBit(const APInt &A,
 /// A.getBitwidth() or NewBitWidth must be a whole multiples of the other.
 LLVM_ABI APInt ScaleBitMask(const APInt &A, unsigned NewBitWidth,
                             bool MatchAllBits = false);
+
+/// Perform a funnel shift left.
+///
+/// fshl(X,Y,Z): (X << (Z % BW)) | (Y >> (BW - (Z % BW)))
+LLVM_ABI APInt fshl(const APInt &Hi, const APInt &Lo, const APInt &Shift);
+
+/// Perform a funnel shift right.
+///
+/// fshr(X,Y,Z): (X << (BW - (Z % BW))) | (Y >> (Z % BW))
+LLVM_ABI APInt fshr(const APInt &Hi, const APInt &Lo, const APInt &Shift);
+
 } // namespace APIntOps
 
 // See friend declaration above. This additional declaration is required in
diff --git a/llvm/lib/Support/APInt.cpp b/llvm/lib/Support/APInt.cpp
index 1547f48bc7ac0..8b2a383a9cdd9 100644
--- a/llvm/lib/Support/APInt.cpp
+++ b/llvm/lib/Support/APInt.cpp
@@ -3169,3 +3169,19 @@ APInt APIntOps::pow(const APInt &X, int64_t N) {
   }
   return Acc;
 }
+
+APInt llvm::APIntOps::fshl(const APInt &Hi, const APInt &Lo,
+                           const APInt &Shift) {
+  APInt BitWidth(
+      APInt(Hi.getBitWidth(), static_cast<uint64_t>(Hi.getBitWidth())));
+  return APInt((Hi.shl(Shift.urem(BitWidth))) |
+               (Lo.lshr(BitWidth - Shift.urem(BitWidth))));
+}
+
+APInt llvm::APIntOps::fshr(const APInt &Hi, const APInt &Lo,
+                           const APInt &Shift) {
+  APInt BitWidth(
+      APInt(Hi.getBitWidth(), static_cast<uint64_t>(Hi.getBitWidth())));
+  return APInt((Hi.shl(BitWidth - Shift.urem(BitWidth))) |
+               (Lo.lshr(Shift.urem(BitWidth))));
+}
diff --git a/llvm/unittests/ADT/APIntTest.cpp b/llvm/unittests/ADT/APIntTest.cpp
index acc6a09acbf4d..89e2ba4560b28 100644
--- a/llvm/unittests/ADT/APIntTest.cpp
+++ b/llvm/unittests/ADT/APIntTest.cpp
@@ -3744,4 +3744,70 @@ TEST(APIntTest, TryExt) {
   ASSERT_EQ(42, APInt(128, -1).trySExtValue().value_or(42));
 }
 
+TEST(APIntTest, Fshl) {
+  EXPECT_EQ(APIntOps::fshl(APInt(8, 0), APInt(8, 255), APInt(8, 8))
+                .getZExtValue(),
+            0);
+  EXPECT_EQ(APIntOps::fshl(APInt(8, 255), APInt(8, 0), APInt(8, 8))
+                .getZExtValue(),
+            255);
+  EXPECT_EQ(APIntOps::fshl(APInt(8, 255), APInt(8, 0), APInt(8, 15))
+                .getZExtValue(),
+            128);
+  EXPECT_EQ(APIntOps::fshl(APInt(8, 15), APInt(8, 15), APInt(8, 11))
+                .getZExtValue(),
+            120);
+  EXPECT_EQ(APIntOps::fshl(APInt(8, 2), APInt(8, 1), APInt(8, 3))
+                .getZExtValue(),
+            16);
+  EXPECT_EQ(APIntOps::fshl(APInt(32, 0, /*isSigned*/ true),
+                           APInt(32, 2147483647, /*isSigned*/ true),
+                           APInt(32, 32, /*isSigned*/ true))
+                .getSExtValue(),
+            0);
+  EXPECT_EQ(APIntOps::fshl(APInt(64, 1, /*isSigned*/ true),
+                           APInt(64, 2, /*isSigned*/ true),
+                           APInt(64, 3, /*isSigned*/ true))
+                .getSExtValue(),
+            8);
+  EXPECT_EQ(APIntOps::fshl(APInt(16, -2, /*isSigned*/ true),
+                           APInt(16, -1, /*isSigned*/ true),
+                           APInt(16, 3, /*isSigned*/ true))
+                .getSExtValue(),
+            -9);
+}
+
+TEST(APIntTest, Fshr) {
+  EXPECT_EQ(APIntOps::fshr(APInt(8, 0), APInt(8, 255), APInt(8, 8))
+                .getZExtValue(),
+            255);
+  EXPECT_EQ(APIntOps::fshr(APInt(8, 255), APInt(8, 0), APInt(8, 8))
+                .getZExtValue(),
+            0);
+  EXPECT_EQ(APIntOps::fshr(APInt(8, 255), APInt(8, 0), APInt(8, 15))
+                .getZExtValue(),
+            254);
+  EXPECT_EQ(APIntOps::fshr(APInt(8, 15), APInt(8, 15), APInt(8, 11))
+                .getZExtValue(),
+            225);
+  EXPECT_EQ(APIntOps::fshr(APInt(8, 1), APInt(8, 2), APInt(8, 3))
+                .getZExtValue(),
+            32);
+  EXPECT_EQ(APIntOps::fshr(APInt(64, 0, /*isSigned*/ true),
+                           APInt(64, 9223372036854775807, /*isSigned*/ true),
+                           APInt(64, 64, /*isSigned*/ true))
+                .getSExtValue(),
+            9223372036854775807);
+  EXPECT_EQ(APIntOps::fshr(APInt(64, 1, /*isSigned*/ true),
+                           APInt(64, 2, /*isSigned*/ true),
+                           APInt(64, 3, /*isSigned*/ true))
+                .getSExtValue(),
+            2305843009213693952);
+  EXPECT_EQ(APIntOps::fshr(APInt(16, -2, /*isSigned*/ true),
+                           APInt(16, -1, /*isSigned*/ true),
+                           APInt(16, 3, /*isSigned*/ true))
+                .getSExtValue(),
+            -8193);
+}
+
 } // end anonymous namespace