Fix is_(cuda_)std_pair_like

NVIDIA · Nov 20, 2023 · 614f7bf · 614f7bf
1 parent 31302b4
commit 614f7bf
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 6 deletions.
diff --git a/include/cuco/detail/traits.hpp b/include/cuco/detail/traits.hpp
@@ -20,6 +20,7 @@
 
 #include <cuda/std/tuple>
 #include <cuda/std/type_traits>
+#include <tuple>
 
 namespace cuco::detail {
 
@@ -29,15 +30,24 @@ struct is_std_pair_like : cuda::std::false_type {
 
 template <typename T>
 struct is_std_pair_like<T,
-                        cuda::std::void_t<decltype(cuda::std::get<0>(cuda::std::declval<T>())),
-                                          decltype(cuda::std::get<1>(cuda::std::declval<T>()))>>
-  : cuda::std::conditional_t<cuda::std::tuple_size<T>::value == 2,
-                             cuda::std::true_type,
-                             cuda::std::false_type> {
+                        cuda::std::void_t<decltype(std::get<0>(cuda::std::declval<T>())),
+                                          decltype(std::get<1>(cuda::std::declval<T>()))>>
+  : cuda::std::
+      conditional_t<std::tuple_size<T>::value == 2, cuda::std::true_type, cuda::std::false_type> {
 };
 
 template <typename T, typename = void>
-struct is_thrust_pair_like_impl : cuda::std::false_type {
+struct is_cuda_std_pair_like : cuda::std::false_type {
+};
+
+template <typename T>
+struct is_cuda_std_pair_like<
+  T,
+  cuda::std::void_t<decltype(cuda::std::get<0>(cuda::std::declval<T>())),
+                    decltype(cuda::std::get<1>(cuda::std::declval<T>()))>>
+  : cuda::std::conditional_t<cuda::std::tuple_size<T>::value == 2,
+                             cuda::std::true_type,
+                             cuda::std::false_type> {
 };
 
 template <typename T>

diff --git a/include/cuco/pair.cuh b/include/cuco/pair.cuh
@@ -23,6 +23,7 @@
 #include <thrust/tuple.h>
 
 #include <cuda/std/tuple>
+#include <tuple>
 #include <type_traits>
 
 namespace cuco {
@@ -86,6 +87,19 @@ struct alignas(detail::pair_alignment<First, Second>()) pair {
    * @param p The input pair to copy from
    */
   template <typename T, std::enable_if_t<detail::is_std_pair_like<T>::value>* = nullptr>
+  __host__ __device__ constexpr pair(T const& p)
+    : pair{std::get<0>(thrust::raw_reference_cast(p)), std::get<1>(thrust::raw_reference_cast(p))}
+  {
+  }
+
+  /**
+   * @brief Constructs a pair from the given cuda::std::pair-like `p`.
+   *
+   * @tparam T Type of the pair to copy from
+   *
+   * @param p The input pair to copy from
+   */
+  template <typename T, std::enable_if_t<detail::is_cuda_std_pair_like<T>::value>* = nullptr>
   __host__ __device__ constexpr pair(T const& p)
     : pair{cuda::std::get<0>(thrust::raw_reference_cast(p)),
            cuda::std::get<1>(thrust::raw_reference_cast(p))}