The great Thrust index type fix, part 9: exclusive_scan, inclusive_scan.
MODIFY
testing/scan.cu
testing/scan.cu
#include <thrust/functional.h>
#include <thrust/iterator/discard_iterator.h>
#include <thrust/iterator/retag.h>
+#include <thrust/device_malloc.h>
+#include <thrust/device_free.h>
template<typename T>
@@ -572,3 +574,93 @@ void TestInclusiveScanWithIndirection(void)
}
DECLARE_INTEGRAL_VECTOR_UNITTEST(TestInclusiveScanWithIndirection);
+struct only_set_when_expected_it
+{
+ long long expected;
+ bool * flag;
+
+ __host__ __device__ only_set_when_expected_it operator++() const { return *this; }
+ __host__ __device__ only_set_when_expected_it operator*() const { return *this; }
+ template<typename Difference>
+ __host__ __device__ only_set_when_expected_it operator+(Difference) const { return *this; }
+ template<typename Index>
+ __host__ __device__ only_set_when_expected_it operator[](Index) const { return *this; }
+
+ __device__
+ void operator=(long long value) const
+ {
+ if (value == expected)
+ {
+ *flag = true;
+ }
+ }
+};
+
+namespace thrust
+{
+template<>
+struct iterator_traits<only_set_when_expected_it>
+{
+ typedef long long value_type;
+ typedef only_set_when_expected_it reference;
+};
+}
+
+void TestInclusiveScanWithBigIndexesHelper(int magnitude)
+{
+ thrust::constant_iterator<long long> begin(1);
+ thrust::constant_iterator<long long> end = begin + (1ll << magnitude);
+ ASSERT_EQUAL(thrust::distance(begin, end), 1ll << magnitude);
+
+ thrust::device_ptr<bool> has_executed = thrust::device_malloc<bool>(1);
+ *has_executed = false;
+
+ only_set_when_expected_it out = { (1ll << magnitude), thrust::raw_pointer_cast(has_executed) };
+
+ thrust::inclusive_scan(thrust::device, begin, end, out);
+
+ bool has_executed_h = *has_executed;
+ thrust::device_free(has_executed);
+
+ ASSERT_EQUAL(has_executed_h, true);
+}
+
+void TestInclusiveScanWithBigIndexes()
+{
+ TestInclusiveScanWithBigIndexesHelper(30);
+ TestInclusiveScanWithBigIndexesHelper(31);
+ TestInclusiveScanWithBigIndexesHelper(32);
+ TestInclusiveScanWithBigIndexesHelper(33);
+}
+
+DECLARE_UNITTEST(TestInclusiveScanWithBigIndexes);
+
+void TestExclusiveScanWithBigIndexesHelper(int magnitude)
+{
+ thrust::constant_iterator<long long> begin(1);
+ thrust::constant_iterator<long long> end = begin + (1ll << magnitude);
+ ASSERT_EQUAL(thrust::distance(begin, end), 1ll << magnitude);
+
+ thrust::device_ptr<bool> has_executed = thrust::device_malloc<bool>(1);
+ *has_executed = false;
+
+ only_set_when_expected_it out = { (1ll << magnitude) - 1, thrust::raw_pointer_cast(has_executed) };
+
+ thrust::exclusive_scan(thrust::device, begin, end, out,0ll);
+
+ bool has_executed_h = *has_executed;
+ thrust::device_free(has_executed);
+
+ ASSERT_EQUAL(has_executed_h, true);
+}
+
+void TestExclusiveScanWithBigIndexes()
+{
+ TestExclusiveScanWithBigIndexesHelper(30);
+ TestExclusiveScanWithBigIndexesHelper(31);
+ TestExclusiveScanWithBigIndexesHelper(32);
+ TestExclusiveScanWithBigIndexesHelper(33);
+}
+
+DECLARE_UNITTEST(TestExclusiveScanWithBigIndexes);
+

MODIFY
thrust/system/cuda/detail/scan.h
thrust/system/cuda/detail/scan.h
#include <cub/device/device_scan.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/detail/mpl/math.h>
#include <thrust/detail/minmax.h>
#include <thrust/distance.h>
+#include <thrust/iterator/iterator_traits.h>
THRUST_BEGIN_NS
template <typename DerivedPolicy,
@@ -710,15 +712,18 @@ namespace __scan {
bool debug_sync = THRUST_DEBUG_SYNC_FLAG;
cudaError_t status;
- status = doit_step<Inclusive>(NULL,
+ THRUST_INDEX_TYPE_DISPATCH(status,
+ doit_step<Inclusive>,
+ num_items,
+ (NULL,
storage_size,
input_it,
- num_items,
+ num_items_fixed,
add_init_to_exclusive_scan,
output_it,
scan_op,
stream,
- debug_sync);
+ debug_sync));
cuda_cub::throw_on_error(status, "scan failed on 1st step");
// Allocate temporary storage.
@@ -726,15 +731,18 @@ namespace __scan {
tmp(policy, storage_size);
void *ptr = static_cast<void*>(tmp.data().get());
- status = doit_step<Inclusive>(ptr,
+ THRUST_INDEX_TYPE_DISPATCH(status,
+ doit_step<Inclusive>,
+ num_items,
+ (ptr,
storage_size,
input_it,
- num_items,
+ num_items_fixed,
add_init_to_exclusive_scan,
output_it,
scan_op,
stream,
- debug_sync);
+ debug_sync));
cuda_cub::throw_on_error(status, "scan failed on 2nd step");
status = cuda_cub::synchronize(policy);
@@ -798,7 +806,8 @@ inclusive_scan(execution_policy<Derived> &policy,
OutputIt result,
ScanOp scan_op)
{
- int num_items = static_cast<int>(thrust::distance(first, last));
+ typedef typename thrust::iterator_traits<InputIt>::difference_type diff_t;
+ diff_t num_items = thrust::distance(first, last);
return cuda_cub::inclusive_scan_n(policy, first, num_items, result, scan_op);
}
@@ -873,7 +882,8 @@ exclusive_scan(execution_policy<Derived> &policy,
T init,
ScanOp scan_op)
{
- int num_items = static_cast<int>(thrust::distance(first, last));
+ typedef typename thrust::iterator_traits<InputIt>::difference_type diff_t;
+ diff_t num_items = thrust::distance(first, last);
return cuda_cub::exclusive_scan_n(policy, first, num_items, result, init, scan_op);
}

The great Thrust index type fix, part 9: exclusive_scan, inclusive_scan.