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

Bug 2448170
Github NVIDIA#967

Author: griwes

testing/for_each.cu

@@ -3,6 +3,8 @@
 #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
@@ -352,3 +354,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);

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)

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)