Proclaim pair and tuple trivially relocatable

NVIDIA · Jul 23, 2024 · 45e168f · 45e168f
1 parent 53fe08f
commit 45e168f
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 0 deletions.
diff --git a/thrust/testing/type_traits.cu b/thrust/testing/type_traits.cu
@@ -5,8 +5,14 @@
 #include <thrust/iterator/iterator_traits.h>
 #include <thrust/iterator/transform_iterator.h>
 #include <thrust/iterator/zip_iterator.h>
+#include <thrust/pair.h>
+#include <thrust/tuple.h>
 #include <thrust/type_traits/is_contiguous_iterator.h>
 
+#include <cuda/std/complex>
+#include <cuda/std/tuple>
+#include <cuda/std/utility>
+
 #include <unittest/unittest.h>
 
 void TestIsContiguousIterator()
@@ -146,3 +152,32 @@ void TestIsCommutative()
   }
 }
 DECLARE_UNITTEST(TestIsCommutative);
+
+void TestTriviallyRelocatable()
+{
+  static_assert(thrust::is_trivially_relocatable<int>::value, "");
+#if THRUST_DEVICE_SYSTEM == THRUST_DEVICE_SYSTEM_CUDA
+  static_assert(thrust::is_trivially_relocatable<__half>::value, "");
+  static_assert(thrust::is_trivially_relocatable<int1>::value, "");
+  static_assert(thrust::is_trivially_relocatable<int2>::value, "");
+  static_assert(thrust::is_trivially_relocatable<int3>::value, "");
+  static_assert(thrust::is_trivially_relocatable<int4>::value, "");
+  static_assert(thrust::is_trivially_relocatable<__int128>::value, "");
+#endif // THRUST_DEVICE_SYSTEM == THRUST_DEVICE_SYSTEM_CUDA
+  static_assert(thrust::is_trivially_relocatable<thrust::complex<float>>::value, "");
+  static_assert(thrust::is_trivially_relocatable<::cuda::std::complex<float>>::value, "");
+  static_assert(thrust::is_trivially_relocatable<thrust::pair<int, int>>::value, "");
+  static_assert(thrust::is_trivially_relocatable<::cuda::std::pair<int, int>>::value, "");
+  static_assert(thrust::is_trivially_relocatable<thrust::tuple<int, float, char>>::value, "");
+  static_assert(thrust::is_trivially_relocatable<::cuda::std::tuple<int, float, char>>::value, "");
+  static_assert(thrust::is_trivially_relocatable<
+                  ::cuda::std::tuple<thrust::pair<int, thrust::tuple<int, ::cuda::std::tuple<>>>,
+                                     thrust::tuple<::cuda::std::pair<int, thrust::tuple<>>, int>>>::value,
+                "");
+
+  static_assert(!thrust::is_trivially_relocatable<thrust::pair<int, std::string>>::value, "");
+  static_assert(!thrust::is_trivially_relocatable<::cuda::std::pair<int, std::string>>::value, "");
+  static_assert(!thrust::is_trivially_relocatable<thrust::tuple<int, float, std::string>>::value, "");
+  static_assert(!thrust::is_trivially_relocatable<::cuda::std::tuple<int, float, std::string>>::value, "");
+};
+DECLARE_UNITTEST(TestTriviallyRelocatable);
diff --git a/thrust/thrust/pair.h b/thrust/thrust/pair.h
@@ -30,6 +30,9 @@
 #  pragma system_header
 #endif // no system header
 
+#include <thrust/type_traits/is_trivially_relocatable.h>
+
+#include <cuda/std/__type_traits/conjunction.h>
 #include <cuda/std/utility>
 
 THRUST_NAMESPACE_BEGIN
@@ -117,6 +120,11 @@ make_pair(T1&& t1, T2&& t2)
 
 using _CUDA_VSTD::get;
 
+template <typename T, typename U>
+struct proclaim_trivially_relocatable<pair<T, U>>
+    : ::cuda::std::conjunction<is_trivially_relocatable<T>, is_trivially_relocatable<U>>
+{};
+
 /*! \endcond
  */
 

diff --git a/thrust/thrust/tuple.h b/thrust/thrust/tuple.h
@@ -39,6 +39,8 @@
 #  pragma system_header
 #endif // no system header
 
+#include <thrust/type_traits/is_trivially_relocatable.h>
+
 #include <cuda/std/tuple>
 #include <cuda/std/type_traits>
 #include <cuda/std/utility>
@@ -232,6 +234,10 @@ inline _CCCL_HOST_DEVICE tuple<Ts&...> tie(Ts&... ts) noexcept
 
 using _CUDA_VSTD::get;
 
+template <typename... Ts>
+struct proclaim_trivially_relocatable<tuple<Ts...>> : ::cuda::std::conjunction<is_trivially_relocatable<Ts>...>
+{};
+
 /*! \endcond
  */
 

diff --git a/thrust/thrust/type_traits/is_trivially_relocatable.h b/thrust/thrust/type_traits/is_trivially_relocatable.h
@@ -36,6 +36,10 @@
 #include <thrust/detail/type_traits.h>
 #include <thrust/type_traits/is_contiguous_iterator.h>
 
+#include <cuda/std/__fwd/pair.h>
+#include <cuda/std/__fwd/tuple.h>
+#include <cuda/std/__type_traits/conjunction.h>
+
 #include <type_traits>
 
 THRUST_NAMESPACE_BEGIN
@@ -285,6 +289,18 @@ THRUST_PROCLAIM_TRIVIALLY_RELOCATABLE(double3)
 THRUST_PROCLAIM_TRIVIALLY_RELOCATABLE(double4)
 #endif
 
+THRUST_NAMESPACE_BEGIN
+template <typename T, typename U>
+struct proclaim_trivially_relocatable<::cuda::std::pair<T, U>>
+    : ::cuda::std::conjunction<is_trivially_relocatable<T>, is_trivially_relocatable<U>>
+{};
+
+template <typename... Ts>
+struct proclaim_trivially_relocatable<::cuda::std::tuple<Ts...>>
+    : ::cuda::std::conjunction<is_trivially_relocatable<Ts>...>
+{};
+THRUST_NAMESPACE_END
+
 /*! \endcond
  */