Core: Count bits with the CUDA intrinsics in device code

libeigen/eigen!3339

Co-authored-by: Rasmus Munk Larsen <rmlarsen@gmail.com>

Eigen/src/Core/MathFunctions.h

diff --git a/Eigen/src/Core/MathFunctions.h b/Eigen/src/Core/MathFunctions.h
index 7ede5d8..ad3118d 100644
--- a/Eigen/src/Core/MathFunctions.h
+++ b/Eigen/src/Core/MathFunctions.h
@@ -666,7 +666,49 @@
   return count_bits_impl<BitsType>::popcount(bits);
 }
 
-#if EIGEN_COMP_GNUC || EIGEN_COMP_CLANG
+#if defined(EIGEN_CUDA_ARCH)
+
+// nvcc's device code has only __builtin_clz, __builtin_clzll, __builtin_popcount and __builtin_popcountll. The other
+// builtins are host functions there, and with warning 20014 silenced, a call from here compiles and returns garbage:
+// ctz at every width, and clz and popcount where long has 64 bits. __clz and __clzll return the full width for zero.
+// From CUDA 13.2, their parameter and return types are unsigned on aarch64 with GCC 11.4 or later as the host compiler
+// and signed elsewhere and in earlier toolkits; the conversions below keep the bits under either declaration.
+template <typename BitsType>
+struct count_bits_impl<
+    BitsType, std::enable_if_t<std::is_integral<BitsType>::value && sizeof(BitsType) <= sizeof(unsigned int)>> {
+  static constexpr int kNumBits = static_cast<int>(sizeof(BitsType) * CHAR_BIT);
+  static EIGEN_DEVICE_FUNC inline int clz(BitsType bits) {
+    static constexpr int kLeadingBitsOffset = (sizeof(unsigned int) - sizeof(BitsType)) * CHAR_BIT;
+    return static_cast<int>(__clz(static_cast<unsigned int>(bits))) - kLeadingBitsOffset;
+  }
+
+  static EIGEN_DEVICE_FUNC inline int ctz(BitsType bits) {
+    return bits == 0 ? kNumBits : __ffs(static_cast<int>(static_cast<unsigned int>(bits))) - 1;
+  }
+
+  static EIGEN_DEVICE_FUNC inline int popcount(BitsType bits) { return __popc(static_cast<unsigned int>(bits)); }
+};
+
+template <typename BitsType>
+struct count_bits_impl<BitsType,
+                       std::enable_if_t<std::is_integral<BitsType>::value && sizeof(unsigned int) < sizeof(BitsType) &&
+                                        sizeof(BitsType) <= sizeof(unsigned long long)>> {
+  static constexpr int kNumBits = static_cast<int>(sizeof(BitsType) * CHAR_BIT);
+  static EIGEN_DEVICE_FUNC inline int clz(BitsType bits) {
+    static constexpr int kLeadingBitsOffset = (sizeof(unsigned long long) - sizeof(BitsType)) * CHAR_BIT;
+    return static_cast<int>(__clzll(static_cast<unsigned long long>(bits))) - kLeadingBitsOffset;
+  }
+
+  static EIGEN_DEVICE_FUNC inline int ctz(BitsType bits) {
+    return bits == 0 ? kNumBits : __ffsll(static_cast<long long>(static_cast<unsigned long long>(bits))) - 1;
+  }
+
+  static EIGEN_DEVICE_FUNC inline int popcount(BitsType bits) {
+    return __popcll(static_cast<unsigned long long>(bits));
+  }
+};
+
+#elif EIGEN_COMP_GNUC || EIGEN_COMP_CLANG
 
 template <typename BitsType>
 struct count_bits_impl<
@@ -795,7 +837,7 @@
 
 #endif  // _WIN64
 
-#endif  // EIGEN_COMP_GNUC || EIGEN_COMP_CLANG
+#endif  // EIGEN_CUDA_ARCH
 
 template <typename BitsType>
 struct log_2_impl {

test/gpu_basic.cu

diff --git a/test/gpu_basic.cu b/test/gpu_basic.cu
index 319d45d..83a1694 100644
--- a/test/gpu_basic.cu
+++ b/test/gpu_basic.cu
@@ -685,6 +685,46 @@
 #endif
 }
 
+// clz, ctz and popcount of the low 8, 16, 32 and 64 bits of each input. The clz test checks the host counts.
+struct bit_count_test {
+  template <typename T>
+  EIGEN_DEVICE_FUNC static void count(T bits, int* out) {
+    out[0] = internal::clz(bits);
+    out[1] = internal::ctz(bits);
+    out[2] = internal::popcount(bits);
+  }
+
+  EIGEN_DEVICE_FUNC void operator()(int i, const uint64_t* in, int* out) const {
+    const uint64_t bits = in[i];
+    count(static_cast<uint8_t>(bits), out + 12 * i);
+    count(static_cast<uint16_t>(bits), out + 12 * i + 3);
+    count(static_cast<uint32_t>(bits), out + 12 * i + 6);
+    count(bits, out + 12 * i + 9);
+  }
+};
+
+void test_bit_count() {
+  const int n = 3 * 64 + 2;
+  Eigen::Array<uint64_t, Eigen::Dynamic, 1> in(n);
+  Eigen::ArrayXi out_ref(12 * n), out_gpu(12 * n);
+  for (int k = 0; k < 64; ++k) {
+    in(3 * k) = uint64_t(1) << k;
+    in(3 * k + 1) = ~(uint64_t(1) << k);
+    in(3 * k + 2) = internal::random<uint64_t>();
+  }
+  in(n - 2) = 0;
+  in(n - 1) = ~uint64_t(0);
+  out_ref.setConstant(-1);
+  out_gpu.setConstant(-1);
+
+  run_on_cpu(bit_count_test(), n, in, out_ref);
+  run_on_gpu(bit_count_test(), n, in, out_gpu);
+
+#if !defined(EIGEN_GPU_COMPILE_PHASE)
+  VERIFY_IS_CWISE_EQUAL(out_ref, out_gpu);
+#endif
+}
+
 template <typename Type1, typename Type2>
 bool verifyIsApproxWithInfsNans(const Type1& a, const Type2& b,
                                 typename Type1::Scalar* = 0)  // Enabled for Eigen's type only
@@ -892,6 +932,7 @@
 
   CALL_SUBTEST(test_custom_less_scalar_minmax());
   CALL_SUBTEST(test_float_nan_minmax());
+  CALL_SUBTEST(test_bit_count());
 
   // `preverse` on every GPU packet type. 32 elements cover several packets of
   // each, including `Packet4h2`, the widest at 8.