[ADT] Handle uint8_t and uint16_t in countr_zero

[kazutakahirata]

Without this patch, the uint8_t and uint16_t cases are sent to the
fallback route. This patch fixes that by relaxing the "if" condition.

While it's hard to test that the correct control path is taken within
countr_zero, this patch adds a few tests just to verify the
correctness on uint8_t and uint16_t inputs.

[llvmbot]

@llvm/pr-subscribers-llvm-adt

Author: Kazu Hirata (kazutakahirata)

Changes

Without this patch, the uint8_t and uint16_t cases are sent to the
fallback route. This patch fixes that by relaxing the "if" condition.

While it's hard to test that the correct control path is taken within
countr_zero, this patch adds a few tests just to verify the
correctness on uint8_t and uint16_t inputs.

---

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

2 Files Affected:

• (modified) llvm/include/llvm/ADT/bit.h (+1-1)
• (modified) llvm/unittests/ADT/BitTest.cpp (+8)

diff --git a/llvm/include/llvm/ADT/bit.h b/llvm/include/llvm/ADT/bit.h
index d6e33c3e6133a..2ca9b43519740 100644
--- a/llvm/include/llvm/ADT/bit.h
+++ b/llvm/include/llvm/ADT/bit.h
@@ -161,7 +161,7 @@ template <typename T> [[nodiscard]] int countr_zero(T Val) {
     return std::numeric_limits<T>::digits;
 
   // Use the intrinsic if available.
-  if constexpr (sizeof(T) == 4) {
+  if constexpr (sizeof(T) <= 4) {
 #if __has_builtin(__builtin_ctz) || defined(__GNUC__)
     return __builtin_ctz(Val);
 #elif defined(_MSC_VER)
diff --git a/llvm/unittests/ADT/BitTest.cpp b/llvm/unittests/ADT/BitTest.cpp
index 2377ce3b78261..88ae36c44bdb9 100644
--- a/llvm/unittests/ADT/BitTest.cpp
+++ b/llvm/unittests/ADT/BitTest.cpp
@@ -297,6 +297,14 @@ TEST(BitTest, CountrZero) {
   EXPECT_EQ(1, llvm::countr_zero(NZ16));
   EXPECT_EQ(1, llvm::countr_zero(NZ32));
   EXPECT_EQ(1, llvm::countr_zero(NZ64));
+
+  EXPECT_EQ(0, llvm::countr_zero(uint8_t(1)));
+  EXPECT_EQ(3, llvm::countr_zero(uint8_t(8)));
+  EXPECT_EQ(7, llvm::countr_zero(uint8_t(128)));
+
+  EXPECT_EQ(0, llvm::countr_zero(uint16_t(1)));
+  EXPECT_EQ(8, llvm::countr_zero(uint16_t(256)));
+  EXPECT_EQ(15, llvm::countr_zero(uint16_t(32768)));
 }
 
 TEST(BitTest, CountlOne) {