Fix the CUDA backend for_each so that it can handle big indexes.
Bug
Github
MODIFY
testing/for_each.cu
testing/for_each.cu
#include <thrust/device_ptr.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/retag.h>
+#include <thrust/device_malloc.h>
+#include <thrust/device_free.h>
#include <algorithm>
THRUST_DISABLE_MSVC_POSSIBLE_LOSS_OF_DATA_WARNING_BEGIN
@@ -369,3 +371,46 @@ void TestForEachNWithLargeTypes(void)
DECLARE_UNITTEST(TestForEachNWithLargeTypes);
THRUST_DISABLE_MSVC_POSSIBLE_LOSS_OF_DATA_WARNING_END
+
+struct OnlySetWhenExpected
+{
+ unsigned long long expected;
+ bool * flag;
+
+ __device__
+ void operator()(unsigned long long x)
+ {
+ if (x == expected)
+ {
+ *flag = true;
+ }
+ }
+};
+
+void TestForEachWithBigIndexesHelper(int magnitude)
+{
+ thrust::counting_iterator<unsigned long long> begin(0);
+ thrust::counting_iterator<unsigned long long> end = begin + (1ull << magnitude);
+ ASSERT_EQUAL(thrust::distance(begin, end), 1ll << magnitude);
+
+ thrust::device_ptr<bool> has_executed = thrust::device_malloc<bool>(1);
+ *has_executed = false;
+
+ OnlySetWhenExpected fn = { (1ull << magnitude) - 1, thrust::raw_pointer_cast(has_executed) };
+
+ thrust::for_each(thrust::device, begin, end, fn);
+
+ bool has_executed_h = *has_executed;
+ thrust::device_free(has_executed);
+
+ ASSERT_EQUAL(has_executed_h, true);
+}
+
+void TestForEachWithBigIndexes()
+{
+ TestForEachWithBigIndexesHelper(30);
+ TestForEachWithBigIndexesHelper(31);
+ TestForEachWithBigIndexesHelper(32);
+ TestForEachWithBigIndexesHelper(33);
+}
+DECLARE_UNITTEST(TestForEachWithBigIndexes);

MODIFY
thrust/system/cuda/detail/core/agent_launcher.h
thrust/system/cuda/detail/core/agent_launcher.h
@@ -408,7 +408,7 @@ namespace core {
stream(stream_),
name(name_),
debug_sync(debug_sync_),
- grid(static_cast<unsigned int>(count + plan.items_per_tile - 1) / plan.items_per_tile),
+ grid(static_cast<unsigned int>((count + plan.items_per_tile - 1) / plan.items_per_tile)),
vshmem(NULL),
has_shmem((size_t)core::get_max_shared_memory_per_block() >= (size_t)plan.shared_memory_size),
shmem_size(has_shmem ? plan.shared_memory_size : 0)
@@ -429,7 +429,7 @@ namespace core {
stream(stream_),
name(name_),
debug_sync(debug_sync_),
- grid(static_cast<unsigned int>(count + plan.items_per_tile - 1) / plan.items_per_tile),
+ grid(static_cast<unsigned int>((count + plan.items_per_tile - 1) / plan.items_per_tile)),
vshmem(vshmem),
has_shmem((size_t)core::get_max_shared_memory_per_block() >= (size_t)plan.shared_memory_size),
shmem_size(has_shmem ? plan.shared_memory_size : 0)

MODIFY
thrust/system/cuda/detail/parallel_for.h
thrust/system/cuda/detail/parallel_for.h
@@ -93,7 +93,7 @@ namespace __parallel_for {
#pragma unroll
for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM)
{
- int idx = BLOCK_THREADS * ITEM + threadIdx.x;
+ Size idx = BLOCK_THREADS * ITEM + threadIdx.x;
if (IS_FULL_TILE || idx < items_in_tile)
f(tile_base + idx);
}
@@ -103,9 +103,9 @@ namespace __parallel_for {
Size num_items,
char * /*shmem*/ )
{
- Size tile_base = blockIdx.x * ITEMS_PER_TILE;
+ Size tile_base = static_cast<Size>(blockIdx.x) * ITEMS_PER_TILE;
Size num_remaining = num_items - tile_base;
- int items_in_tile = static_cast<int>(
+ Size items_in_tile = static_cast<Size>(
num_remaining < ITEMS_PER_TILE ? num_remaining : ITEMS_PER_TILE);
if (items_in_tile == ITEMS_PER_TILE)

Fix the CUDA backend for_each so that it can handle big indexes.

Bug 2448170
Github #967