uses dynamic dispatch to unsigned type

NVIDIA · Aug 19, 2024 · ebea6bd · ebea6bd
1 parent 4aaac08
commit ebea6bd
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 4 deletions.
diff --git a/thrust/thrust/system/cuda/detail/dispatch.h b/thrust/thrust/system/cuda/detail/dispatch.h
@@ -90,6 +90,20 @@
     status                             = call_64 arguments;                     \
   }
 
+/// Like \ref THRUST_INDEX_TYPE_DISPATCH2 but dispatching to uint32_t and uint64_t, respectively, depending on the
+/// `count` argument.
+#define THRUST_UNSIGNED_INDEX_TYPE_DISPATCH2(status, call_32, call_64, count, arguments) \
+  if (count <= thrust::detail::integer_traits<std::uint32_t>::const_max)                 \
+  {                                                                                      \
+    auto THRUST_PP_CAT2(count, _fixed) = static_cast<std::uint32_t>(count);              \
+    status                             = call_32 arguments;                              \
+  }                                                                                      \
+  else                                                                                   \
+  {                                                                                      \
+    auto THRUST_PP_CAT2(count, _fixed) = static_cast<std::uint64_t>(count);              \
+    status                             = call_64 arguments;                              \
+  }
+
 /// Like \ref THRUST_INDEX_TYPE_DISPATCH2 but uses two counts.
 #define THRUST_DOUBLE_INDEX_TYPE_DISPATCH2(status, call_32, call_64, count1, count2, arguments) \
   if (count1 + count2 <= thrust::detail::integer_traits<std::int32_t>::const_max)               \

diff --git a/thrust/thrust/system/cuda/detail/scan.h b/thrust/thrust/system/cuda/detail/scan.h
@@ -72,7 +72,7 @@ _CCCL_HOST_DEVICE OutputIt inclusive_scan_n_impl(
   // Determine temporary storage requirements:
   size_t tmp_size = 0;
   {
-    THRUST_INDEX_TYPE_DISPATCH2(
+    THRUST_UNSIGNED_INDEX_TYPE_DISPATCH2(
       status,
       Dispatch32::Dispatch,
       Dispatch64::Dispatch,
@@ -88,7 +88,7 @@ _CCCL_HOST_DEVICE OutputIt inclusive_scan_n_impl(
   {
     // Allocate temporary storage:
     thrust::detail::temporary_array<std::uint8_t, Derived> tmp{policy, tmp_size};
-    THRUST_INDEX_TYPE_DISPATCH2(
+    THRUST_UNSIGNED_INDEX_TYPE_DISPATCH2(
       status,
       Dispatch32::Dispatch,
       Dispatch64::Dispatch,
@@ -122,7 +122,7 @@ _CCCL_HOST_DEVICE OutputIt exclusive_scan_n_impl(
   // Determine temporary storage requirements:
   size_t tmp_size = 0;
   {
-    THRUST_INDEX_TYPE_DISPATCH2(
+    THRUST_UNSIGNED_INDEX_TYPE_DISPATCH2(
       status,
       Dispatch32::Dispatch,
       Dispatch64::Dispatch,
@@ -138,7 +138,7 @@ _CCCL_HOST_DEVICE OutputIt exclusive_scan_n_impl(
   {
     // Allocate temporary storage:
     thrust::detail::temporary_array<std::uint8_t, Derived> tmp{policy, tmp_size};
-    THRUST_INDEX_TYPE_DISPATCH2(
+    THRUST_UNSIGNED_INDEX_TYPE_DISPATCH2(
       status,
       Dispatch32::Dispatch,
       Dispatch64::Dispatch,