[llvm] [KnownBits] Fix add() SelfAdd assertion for bitwidths >= 512 (PR #202769)

[Paweł Bylica via llvm-commits]

https://github.com/chfast updated https://github.com/llvm/llvm-project/pull/202769

>From 49053d41159c19e59dba558a15622ba67e73368b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82=20Bylica?= <pawel at hepcolgum.band>
Date: Wed, 10 Jun 2026 17:22:00 +0200
Subject: [PATCH] [KnownBits] Fix getMaxShiftAmount assertion for narrow shift
 amounts

KnownBits::{shl,lshr,ashr} call getMaxShiftAmount(RHS.getMaxValue(),
BitWidth), which for power-of-two BitWidth does
MaxValue.extractBitsAsZExtValue(Log2_32(BitWidth), 0). Nothing requires
the shift amount (RHS) to be as wide as the shifted value, so when its
bitwidth is smaller than Log2_32(BitWidth) the extraction trips the
"Illegal bit extraction" assertion in APInt::extractBitsAsZExtValue.

This is reachable from KnownBits::add() with SelfAdd=true, which lowers
X+X to shl(X, 1) using a fixed 8-bit shift amount: for source widths
>= 512, Log2_32(BitWidth) >= 9 > 8 and the assertion fires.

Fix it centrally in getMaxShiftAmount by clamping the extraction width
to the shift amount's bitwidth. A shift amount that narrow is already
< BitWidth, so the clamped value remains a valid upper bound. This
protects all three shift callers rather than just the add() path.

Found via fuzzing. Adds a SelfAddWide regression test covering widths
256/512/1024 with both minimal and maximal constants.
---
 llvm/lib/Support/KnownBits.cpp           |  5 ++++-
 llvm/unittests/Support/KnownBitsTest.cpp | 19 +++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Support/KnownBits.cpp b/llvm/lib/Support/KnownBits.cpp
index b1b985b80cffd..5becdac89fcaa 100644
--- a/llvm/lib/Support/KnownBits.cpp
+++ b/llvm/lib/Support/KnownBits.cpp
@@ -399,7 +399,10 @@ KnownBits KnownBits::abds(KnownBits LHS, KnownBits RHS) {
 
 static unsigned getMaxShiftAmount(const APInt &MaxValue, unsigned BitWidth) {
   if (isPowerOf2_32(BitWidth))
-    return MaxValue.extractBitsAsZExtValue(Log2_32(BitWidth), 0);
+    // Clamp to the shift amount's width: a narrower amount is already
+    // < BitWidth, so this stays a valid upper bound.
+    return MaxValue.extractBitsAsZExtValue(
+        std::min(Log2_32(BitWidth), MaxValue.getBitWidth()), 0);
   // This is only an approximate upper bound.
   return MaxValue.getLimitedValue(BitWidth - 1);
 }
diff --git a/llvm/unittests/Support/KnownBitsTest.cpp b/llvm/unittests/Support/KnownBitsTest.cpp
index 0303b2f121c3b..581efad2aee8b 100644
--- a/llvm/unittests/Support/KnownBitsTest.cpp
+++ b/llvm/unittests/Support/KnownBitsTest.cpp
@@ -179,6 +179,25 @@ TEST(KnownBitsTest, SelfAddExhaustive) {
   });
 }
 
+TEST(KnownBitsTest, SelfAddWide) {
+  // SelfAdd lowers to shl(X, 1), whose getMaxShiftAmount extracts
+  // Log2_32(BitWidth) bits from the shift amount. For source widths >= 512 that
+  // is more than the 8-bit shift amount used by the lowering, which used to
+  // assert in APInt::extractBitsAsZExtValue.
+  // 1 << 23 is IntegerType::MAX_INT_BITS, the widest integer IR can produce.
+  for (unsigned Bits : {256u, 512u, 1024u, 1u << 23}) {
+    for (const APInt &X : {APInt(Bits, 1), APInt::getMaxValue(Bits)}) {
+      KnownBits Known = KnownBits::makeConstant(X);
+      KnownBits Computed = KnownBits::add(Known, Known, /*NSW=*/false,
+                                          /*NUW=*/false, /*SelfAdd=*/true);
+      // X + X == X << 1 (truncated to Bits); the result is fully known.
+      APInt Doubled = X << 1;
+      EXPECT_EQ(Computed.One, Doubled);
+      EXPECT_EQ(Computed.Zero, ~Doubled);
+    }
+  }
+}
+
 TEST(KnownBitsTest, AddCarryExhaustive) {
   unsigned Bits = 4;
   ForeachKnownBits(Bits, [&](const KnownBits &Known1) {