The great Thrust index type fix, part 1: adjacent_difference, reduce.
MODIFY
testing/adjacent_difference.cu
testing/adjacent_difference.cu
#include <thrust/adjacent_difference.h>
#include <thrust/iterator/discard_iterator.h>
#include <thrust/iterator/retag.h>
+#include <thrust/device_malloc.h>
+#include <thrust/device_free.h>
template <class Vector>
void TestAdjacentDifferenceSimple(void)
@@ -176,3 +178,52 @@ void TestAdjacentDifferenceDispatchImplicit()
ASSERT_EQUAL(13, d_input.front());
}
DECLARE_UNITTEST(TestAdjacentDifferenceDispatchImplicit);
+
+struct detect_wrong_difference
+{
+ bool * flag;
+
+ __host__ __device__ detect_wrong_difference operator++() const { return *this; }
+ __host__ __device__ detect_wrong_difference operator*() const { return *this; }
+ template<typename Difference>
+ __host__ __device__ detect_wrong_difference operator+(Difference) const { return *this; }
+ template<typename Index>
+ __host__ __device__ detect_wrong_difference operator[](Index) const { return *this; }
+
+ __device__
+ void operator=(long long difference) const
+ {
+ if (difference != 1)
+ {
+ *flag = false;
+ }
+ }
+};
+
+void TestAdjacentDifferenceWithBigIndexesHelper(int magnitude)
+{
+ thrust::counting_iterator<long long> begin(1);
+ thrust::counting_iterator<long long> end = begin + (1ll << magnitude);
+ ASSERT_EQUAL(thrust::distance(begin, end), 1ll << magnitude);
+
+ thrust::device_ptr<bool> all_differences_correct = thrust::device_malloc<bool>(1);
+ *all_differences_correct = true;
+
+ detect_wrong_difference out = { thrust::raw_pointer_cast(all_differences_correct) };
+
+ thrust::adjacent_difference(thrust::device, begin, end, out);
+
+ bool all_differences_correct_h = *all_differences_correct;
+ thrust::device_free(all_differences_correct);
+
+ ASSERT_EQUAL(all_differences_correct_h, true);
+}
+
+void TestAdjacentDifferenceWithBigIndexes()
+{
+ TestAdjacentDifferenceWithBigIndexesHelper(30);
+ TestAdjacentDifferenceWithBigIndexesHelper(31);
+ TestAdjacentDifferenceWithBigIndexesHelper(32);
+ TestAdjacentDifferenceWithBigIndexesHelper(33);
+}
+DECLARE_UNITTEST(TestAdjacentDifferenceWithBigIndexes);

MODIFY
thrust/system/cuda/detail/adjacent_difference.h
thrust/system/cuda/detail/adjacent_difference.h
#include <cub/block/block_adjacent_difference.cuh>
#include <thrust/system/cuda/detail/core/agent_launcher.h>
#include <thrust/system/cuda/detail/par_to_seq.h>
+#include <thrust/system/cuda/detail/dispatch.h>
#include <thrust/functional.h>
#include <thrust/distance.h>
#include <thrust/detail/mpl/math.h>
@@ -257,8 +258,8 @@ namespace __adjacent_difference {
template <bool IS_LAST_TILE>
void THRUST_DEVICE_FUNCTION
- consume_tile(Size num_remaining,
- Size tile_idx,
+ consume_tile(int num_remaining,
+ int tile_idx,
Size tile_base)
{
if (tile_idx == 0)
@@ -279,7 +280,7 @@ namespace __adjacent_difference {
consume_range(Size num_items)
{
int tile_idx = blockIdx.x;
- Size tile_base = tile_idx * ITEMS_PER_TILE;
+ Size tile_base = static_cast<Size>(tile_idx) * ITEMS_PER_TILE;
Size num_remaining = num_items - tile_base;
if (num_remaining > ITEMS_PER_TILE) // not a last tile
@@ -349,7 +350,7 @@ namespace __adjacent_difference {
char * /*shmem*/)
{
int tile_idx = blockIdx.x * blockDim.x + threadIdx.x;
- int tile_base = tile_idx * items_per_tile;
+ Size tile_base = static_cast<Size>(tile_idx) * items_per_tile;
if (tile_base > 0 && tile_idx < num_tiles)
result[tile_idx] = first[tile_base - 1];
}
@@ -391,8 +392,8 @@ namespace __adjacent_difference {
AgentPlan init_plan = init_agent::get_plan();
- size_t tile_size = difference_plan.items_per_tile;
- size_t num_tiles = (num_items + tile_size - 1) / tile_size;
+ Size tile_size = difference_plan.items_per_tile;
+ Size num_tiles = (num_items + tile_size - 1) / tile_size;
size_t tmp1 = num_tiles * sizeof(input_type);
size_t vshmem_size = core::vshmem_size(difference_plan.shared_memory_size,
@@ -448,14 +449,9 @@ namespace __adjacent_difference {
bool debug_sync = THRUST_DEBUG_SYNC_FLAG;
cudaError_t status;
- status = doit_step(NULL,
- storage_size,
- first,
- result,
- binary_op,
- num_items,
- stream,
- debug_sync);
+ THRUST_INDEX_TYPE_DISPATCH(status, doit_step, num_items,
+ (NULL, storage_size, first, result, binary_op,
+ num_items_fixed, stream, debug_sync));
cuda_cub::throw_on_error(status, "adjacent_difference failed on 1st step");
// Allocate temporary storage.
@@ -463,14 +459,9 @@ namespace __adjacent_difference {
tmp(policy, storage_size);
void *ptr = static_cast<void*>(tmp.data().get());
- status = doit_step(ptr,
- storage_size,
- first,
- result,
- binary_op,
- num_items,
- stream,
- debug_sync);
+ THRUST_INDEX_TYPE_DISPATCH(status, doit_step, num_items,
+ (ptr, storage_size, first, result, binary_op,
+ num_items_fixed, stream, debug_sync));
cuda_cub::throw_on_error(status, "adjacent_difference failed on 2nd step");
status = cuda_cub::synchronize(policy);

ADD
thrust/system/cuda/detail/dispatch.h
+/*
+ * Copyright 2018 NVIDIA Corporation
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <thrust/detail/preprocessor.h>
+
+/**
+ * Dispatch between 32-bit and 64-bit index based versions of the same algorithm
+ * implementation. This version assumes that callables for both branches consist
+ * of the same tokens, and is intended to be used with Thrust-style dispatch
+ * interfaces, that always deduce the size type from the arguments.
+ */
+#define THRUST_INDEX_TYPE_DISPATCH(status, call, count, arguments) \
+ if (count <= std::numeric_limits<thrust::detail::int32_t>::max()) { \
+ thrust::detail::int32_t THRUST_PP_CAT2(count, _fixed) = count; \
+ status = call arguments; \
+ } \
+ else { \
+ thrust::detail::int64_t THRUST_PP_CAT2(count, _fixed) = count; \
+ status = call arguments; \
+ }
+
+/**
+ * Dispatch between 32-bit and 64-bit index based versions of the same algorithm
+ * implementation. This version allows using different token sequences for callables
+ * in both branches, and is intended to be used with CUB-style dispatch interfaces,
+ * where the "simple" interface always forces the size to be `int` (making it harder
+ * for us to use), but the complex interface that we end up using doesn't actually
+ * provide a way to fully deduce the type from just the call, making the size type
+ * appear in the token sequence of the callable.
+ *
+ * See reduce_n_impl to see an example of how this is meant to be used.
+ */
+#define THRUST_INDEX_TYPE_DISPATCH2(status, call_32, call_64, count, arguments) \
+ if (count <= std::numeric_limits<thrust::detail::int32_t>::max()) { \
+ thrust::detail::int32_t THRUST_PP_CAT2(count, _fixed) = count; \
+ status = call_32 arguments; \
+ } \
+ else { \
+ thrust::detail::int64_t THRUST_PP_CAT2(count, _fixed) = count; \
+ status = call_64 arguments; \
+ }
+

MODIFY
thrust/system/cuda/detail/reduce.h
thrust/system/cuda/detail/reduce.h
#include <cub/device/device_reduce.cuh>
#include <thrust/system/cuda/detail/par_to_seq.h>
#include <thrust/system/cuda/detail/get_value.h>
+#include <thrust/system/cuda/detail/dispatch.h>
#include <thrust/functional.h>
#include <thrust/system/cuda/detail/core/agent_launcher.h>
#include <thrust/detail/minmax.h>
@@ -930,21 +931,22 @@ T reduce_n_impl(execution_policy<Derived>& policy,
BinaryOp binary_op)
{
cudaStream_t stream = cuda_cub::stream(policy);
+ cudaError_t status;
// Determine temporary device storage requirements.
size_t tmp_size = 0;
- cuda_cub::throw_on_error(
- cub::DeviceReduce::Reduce(NULL,
- tmp_size,
- first,
- reinterpret_cast<T*>(NULL),
+
+ THRUST_INDEX_TYPE_DISPATCH2(status,
+ cub::DeviceReduce::Reduce,
+ (cub::DispatchReduce<
+ InputIt, T*, Size, BinaryOp
+ >::Dispatch),
num_items,
- binary_op,
- init,
- stream,
- THRUST_DEBUG_SYNC_FLAG),
- "after reduction step 1");
+ (NULL, tmp_size, first, reinterpret_cast<T*>(NULL),
+ num_items_fixed, binary_op, init, stream,
+ THRUST_DEBUG_SYNC_FLAG));
+ cuda_cub::throw_on_error(status, "after reduction step 1");
// Allocate temporary storage.
@@ -963,17 +965,16 @@ T reduce_n_impl(execution_policy<Derived>& policy,
// make this guarantee.
T* ret_ptr = thrust::detail::aligned_reinterpret_cast<T*>(tmp.data().get());
void* tmp_ptr = static_cast<void*>((tmp.data() + sizeof(T)).get());
- cuda_cub::throw_on_error(
- cub::DeviceReduce::Reduce(tmp_ptr,
- tmp_size,
- first,
- ret_ptr,
+ THRUST_INDEX_TYPE_DISPATCH2(status,
+ cub::DeviceReduce::Reduce,
+ (cub::DispatchReduce<
+ InputIt, T*, Size, BinaryOp
+ >::Dispatch),
num_items,
- binary_op,
- init,
- stream,
- THRUST_DEBUG_SYNC_FLAG),
- "after reduction step 2");
+ (tmp_ptr, tmp_size, first, ret_ptr,
+ num_items_fixed, binary_op, init, stream,
+ THRUST_DEBUG_SYNC_FLAG));
+ cuda_cub::throw_on_error(status, "after reduction step 2");
// Synchronize the stream and get the value.

The great Thrust index type fix, part 1: adjacent_difference, reduce.