[llvm] [ADT] Handle uint8_t and uint16_t in countr_zero (PR #158518)

[Kazu Hirata via llvm-commits]

https://github.com/kazutakahirata created https://github.com/llvm/llvm-project/pull/158518

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.


>From 095b813885f4b61833c0e0cfbbf22cf51b5bd5cd Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu at google.com>
Date: Sun, 14 Sep 2025 14:59:08 -0700
Subject: [PATCH] [ADT] Handle uint8_t and uint16_t in countr_zero

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.
---
 llvm/include/llvm/ADT/bit.h    | 2 +-
 llvm/unittests/ADT/BitTest.cpp | 8 ++++++++
 2 files changed, 9 insertions(+), 1 deletion(-)

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) {