Fixes compiler error for extended fp type data gen (#666)

* fixes ctor compiler error for extended fp type data gen * fixes msvc ctor invalid redeclaration
NVIDIA · Nov 7, 2023 · 591dc78 · 591dc78
1 parent a9d7d8e
commit 591dc78
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 6 deletions.
diff --git a/cub/test/bfloat16.h b/cub/test/bfloat16.h
@@ -32,11 +32,14 @@
  * Utilities for interacting with the opaque CUDA __nv_bfloat16 type
  */
 
-#include <stdint.h>
+#include <cub/util_type.cuh>
+
 #include <cuda_bf16.h>
-#include <iosfwd>
 
-#include <cub/util_type.cuh>
+#include <cuda/std/type_traits>
+
+#include <cstdint>
+#include <iosfwd>
 
 #ifdef __GNUC__
 // There's a ton of type-punning going on in this file.
@@ -77,6 +80,16 @@ struct bfloat16_t
         *this = bfloat16_t(float(a));
     }
 
+    /// Constructor from unsigned long long int
+    template < typename T,
+               typename = typename ::cuda::std::enable_if<
+                 ::cuda::std::is_same<T, unsigned long long int>::value
+                 && (!::cuda::std::is_same<std::size_t, unsigned long long int>::value)>::type>
+    __host__ __device__ __forceinline__ bfloat16_t(T a)
+    {
+      *this = bfloat16_t(float(a));
+    }
+
     /// Default constructor
     bfloat16_t() = default;
 

diff --git a/cub/test/c2h/generators.cu b/cub/test/c2h/generators.cu
@@ -259,16 +259,17 @@ void generator_t::operator()(seed_t seed, thrust::device_vector<T> &data, T min,
 template <typename T>
 struct count_to_item_t
 {
-  std::size_t n;
+  unsigned long long int n;
 
-  count_to_item_t(std::size_t n)
+  count_to_item_t(unsigned long long int n)
       : n(n)
   {}
 
   template <typename CounterT>
   __device__ T operator()(CounterT id)
   {
-    return static_cast<T>(static_cast<std::size_t>(id) % n);
+    // This has to be a type for which extended floating point types like __nv_fp8_e5m2 provide an overload
+    return static_cast<T>(static_cast<unsigned long long int>(id) % n);
   }
 };
 

diff --git a/cub/test/half.h b/cub/test/half.h
@@ -37,6 +37,8 @@
 
 #include <cuda_fp16.h>
 
+#include <cuda/std/type_traits>
+
 #include <cstdint>
 #include <cstring>
 #include <iosfwd>
@@ -80,6 +82,16 @@ struct half_t
         *this = half_t(float(a));
     }
 
+    /// Constructor from unsigned long long int
+    template < typename T,
+               typename = typename ::cuda::std::enable_if<
+                 ::cuda::std::is_same<T, unsigned long long int>::value
+                 && (!::cuda::std::is_same<std::size_t, unsigned long long int>::value)>::type>
+    __host__ __device__ __forceinline__ half_t(T a)
+    {
+      *this = half_t(float(a));
+    }
+
     /// Default constructor
     half_t() = default;