The great Thrust index type fix, part 8: set operations.
MODIFY
testing/set_difference.cu
testing/set_difference.cu
@@ -211,3 +211,28 @@ void TestSetDifferenceMultiset(const size_t n)
}
DECLARE_VARIABLE_UNITTEST(TestSetDifferenceMultiset);
+void TestSetDifferenceWithBigIndexesHelper(int magnitude)
+{
+ thrust::counting_iterator<long long> begin(0);
+ thrust::counting_iterator<long long> end = begin + (1ll << magnitude);
+ thrust::counting_iterator<long long> end_longer = end + 1;
+ ASSERT_EQUAL(thrust::distance(begin, end), 1ll << magnitude);
+
+ thrust::device_vector<long long> result;
+ result.resize(1);
+ thrust::set_difference(thrust::device, begin, end_longer, begin, end, result.begin());
+
+ thrust::host_vector<long long> expected;
+ expected.push_back(*end);
+
+ ASSERT_EQUAL(result, expected);
+}
+
+void TestSetDifferenceWithBigIndexes()
+{
+ TestSetDifferenceWithBigIndexesHelper(30);
+ TestSetDifferenceWithBigIndexesHelper(31);
+ TestSetDifferenceWithBigIndexesHelper(32);
+ TestSetDifferenceWithBigIndexesHelper(33);
+}
+DECLARE_UNITTEST(TestSetDifferenceWithBigIndexes);

MODIFY
testing/set_intersection.cu
testing/set_intersection.cu
@@ -251,3 +251,29 @@ void TestSetIntersectionMultiset(const size_t n)
}
DECLARE_VARIABLE_UNITTEST(TestSetIntersectionMultiset);
+void TestSetDifferenceWithBigIndexesHelper(int magnitude)
+{
+ thrust::counting_iterator<long long> begin1(0);
+ thrust::counting_iterator<long long> begin2 = begin1 + (1ll << magnitude);
+ thrust::counting_iterator<long long> end1 = begin2 + 1;
+ thrust::counting_iterator<long long> end2 = begin2 + (1ll << magnitude);
+ ASSERT_EQUAL(thrust::distance(begin2, end1), 1);
+
+ thrust::device_vector<long long> result;
+ result.resize(1);
+ thrust::set_intersection(thrust::device, begin1, end1, begin2, end2, result.begin());
+
+ thrust::host_vector<long long> expected;
+ expected.push_back(*begin2);
+
+ ASSERT_EQUAL(result, expected);
+}
+
+void TestSetDifferenceWithBigIndexes()
+{
+ TestSetDifferenceWithBigIndexesHelper(30);
+ TestSetDifferenceWithBigIndexesHelper(31);
+ TestSetDifferenceWithBigIndexesHelper(32);
+ TestSetDifferenceWithBigIndexesHelper(33);
+}
+DECLARE_UNITTEST(TestSetDifferenceWithBigIndexes);

MODIFY
thrust/system/cuda/detail/dispatch.h
thrust/system/cuda/detail/dispatch.h
status = call arguments; \
}
+/**
+ * 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.
+ *
+ * This version of the macro supports providing two count variables, which is
+ * necessary for set algorithms.
+ */
+#define THRUST_DOUBLE_INDEX_TYPE_DISPATCH(status, call, count1, count2, arguments) \
+ if (count1 + count2 <= std::numeric_limits<thrust::detail::int32_t>::max()) { \
+ thrust::detail::int32_t THRUST_PP_CAT2(count1, _fixed) = count1; \
+ thrust::detail::int32_t THRUST_PP_CAT2(count2, _fixed) = count2; \
+ status = call arguments; \
+ } \
+ else { \
+ thrust::detail::int64_t THRUST_PP_CAT2(count1, _fixed) = count1; \
+ thrust::detail::int64_t THRUST_PP_CAT2(count2, _fixed) = count2; \
+ 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

MODIFY
thrust/system/cuda/detail/set_operations.h
thrust/system/cuda/detail/set_operations.h
@@ -50,35 +50,36 @@ namespace __set_operations {
template <bool UpperBound,
class IntT,
+ class Size,
class It,
class T,
class Comp>
THRUST_DEVICE_FUNCTION void
binary_search_iteration(It data,
- int &begin,
- int &end,
+ Size &begin,
+ Size &end,
T key,
int shift,
Comp comp)
{
IntT scale = (1 << shift) - 1;
- int mid = (int)((begin + scale * end) >> shift);
+ Size mid = (begin + scale * end) >> shift;
T key2 = data[mid];
bool pred = UpperBound ? !comp(key, key2) : comp(key2, key);
if (pred)
- begin = (int)mid + 1;
+ begin = mid + 1;
else
end = mid;
}
- template <bool UpperBound, class T, class It, class Comp>
- THRUST_DEVICE_FUNCTION int
- binary_search(It data, int count, T key, Comp comp)
+ template <bool UpperBound, class Size, class T, class It, class Comp>
+ THRUST_DEVICE_FUNCTION Size
+ binary_search(It data, Size count, T key, Comp comp)
{
- int begin = 0;
- int end = count;
+ Size begin = 0;
+ Size end = count;
while (begin < end)
binary_search_iteration<UpperBound, int>(data,
begin,
@@ -89,12 +90,12 @@ namespace __set_operations {
return begin;
}
- template <bool UpperBound, class IntT, class T, class It, class Comp>
- THRUST_DEVICE_FUNCTION int
- biased_binary_search(It data, int count, T key, IntT levels, Comp comp)
+ template <bool UpperBound, class IntT, class Size, class T, class It, class Comp>
+ THRUST_DEVICE_FUNCTION Size
+ biased_binary_search(It data, Size count, T key, IntT levels, Comp comp)
{
- int begin = 0;
- int end = count;
+ Size begin = 0;
+ Size end = count;
if (levels >= 4 && begin < end)
binary_search_iteration<UpperBound, IntT>(data, begin, end, key, 9, comp);
@@ -110,18 +111,18 @@ namespace __set_operations {
return begin;
}
- template <bool UpperBound, class It1, class It2, class Comp>
- THRUST_DEVICE_FUNCTION int
- merge_path(It1 a, int aCount, It2 b, int bCount, int diag, Comp comp)
+ template <bool UpperBound, class Size, class It1, class It2, class Comp>
+ THRUST_DEVICE_FUNCTION Size
+ merge_path(It1 a, Size aCount, It2 b, Size bCount, Size diag, Comp comp)
{
typedef typename thrust::iterator_traits<It1>::value_type T;
- int begin = thrust::max(0, diag - bCount);
- int end = thrust::min(diag, aCount);
+ Size begin = thrust::max<Size>(0, diag - bCount);
+ Size end = thrust::min<Size>(diag, aCount);
while (begin < end)
{
- int mid = (begin + end) >> 1;
+ Size mid = (begin + end) >> 1;
T aKey = a[mid];
T bKey = b[diag - 1 - mid];
bool pred = UpperBound ? comp(aKey, bKey) : !comp(bKey, aKey);
@@ -134,7 +135,7 @@ namespace __set_operations {
}
template <class It1, class It2, class Size, class Size2, class CompareOp>
- pair<Size, Size> THRUST_DEVICE_FUNCTION
+ THRUST_DEVICE_FUNCTION pair<Size, Size>
balanced_path(It1 keys1,
It2 keys2,
Size num_keys1,
@@ -434,7 +435,7 @@ namespace __set_operations {
CompareOp compare_op;
SetOp set_op;
pair<Size, Size> *partitions;
- Size *output_count;
+ std::size_t *output_count;
//---------------------------------------------------------------------
// Utility functions
@@ -756,7 +757,7 @@ namespace __set_operations {
CompareOp compare_op_,
SetOp set_op_,
pair<Size, Size> *partitions_,
- Size *output_count_)
+ std::size_t * output_count_)
: storage(storage_),
tile_state(tile_state_),
keys1_in(core::make_load_iterator(ptx_plan(), keys1_)),
@@ -801,7 +802,7 @@ namespace __set_operations {
CompareOp compare_op,
SetOp set_op,
pair<Size, Size> *partitions,
- Size * output_count,
+ std::size_t * output_count,
ScanTileState tile_state,
char * shmem)
{
@@ -1124,7 +1125,7 @@ namespace __set_operations {
Size num_keys2,
KeysOutputIt keys_output,
ValuesOutputIt values_output,
- Size * output_count,
+ std::size_t * output_count,
CompareOp compare_op,
SetOp set_op,
cudaStream_t stream,
@@ -1167,7 +1168,7 @@ namespace __set_operations {
Size num_tiles = (keys_total + tile_size - 1) / tile_size;
size_t tile_agent_storage;
- status = ScanTileState::AllocationSize(static_cast<int>(num_tiles), tile_agent_storage);
+ status = ScanTileState::AllocationSize(num_tiles, tile_agent_storage);
CUDA_CUB_RET_IF_FAIL(status);
size_t vshmem_storage = core::vshmem_size(set_op_plan.shared_memory_size,
@@ -1191,7 +1192,7 @@ namespace __set_operations {
}
ScanTileState tile_state;
- status = tile_state.Init(static_cast<int>(num_tiles), allocations[0], allocation_sizes[0]);
+ status = tile_state.Init(num_tiles, allocations[0], allocation_sizes[0]);
CUDA_CUB_RET_IF_FAIL(status);
pair<Size, Size> *partitions = (pair<Size, Size> *)allocations[1];
@@ -1268,24 +1269,25 @@ namespace __set_operations {
bool debug_sync = THRUST_DEBUG_SYNC_FLAG;
cudaError_t status;
- status = doit_step<HAS_VALUES>(NULL,
+ THRUST_DOUBLE_INDEX_TYPE_DISPATCH(status, doit_step<HAS_VALUES>,
+ num_keys1, num_keys2, (NULL,
temp_storage_bytes,
keys1_first,
keys2_first,
values1_first,
values2_first,
- num_keys1,
- num_keys2,
+ num_keys1_fixed,
+ num_keys2_fixed,
keys_output,
values_output,
- reinterpret_cast<size_type*>(NULL),
+ reinterpret_cast<std::size_t*>(NULL),
compare_op,
set_op,
stream,
- debug_sync);
+ debug_sync));
cuda_cub::throw_on_error(status, "set_operations failed on 1st step");
- size_t allocation_sizes[2] = {sizeof(size_type), temp_storage_bytes};
+ size_t allocation_sizes[2] = {sizeof(std::size_t), temp_storage_bytes};
void * allocations[2] = {NULL, NULL};
size_t storage_size = 0;
@@ -1307,30 +1309,31 @@ namespace __set_operations {
allocation_sizes);
cuda_cub::throw_on_error(status, "set_operations failed on 2nd alias_storage");
- size_type* d_output_count
- = thrust::detail::aligned_reinterpret_cast<size_type*>(allocations[0]);
+ std::size_t* d_output_count
+ = thrust::detail::aligned_reinterpret_cast<std::size_t*>(allocations[0]);
- status = doit_step<HAS_VALUES>(allocations[1],
+ THRUST_DOUBLE_INDEX_TYPE_DISPATCH(status, doit_step<HAS_VALUES>,
+ num_keys1, num_keys2, (allocations[1],
temp_storage_bytes,
keys1_first,
keys2_first,
values1_first,
values2_first,
- num_keys1,
- num_keys2,
+ num_keys1_fixed,
+ num_keys2_fixed,
keys_output,
values_output,
d_output_count,
compare_op,
set_op,
stream,
- debug_sync);
+ debug_sync));
cuda_cub::throw_on_error(status, "set_operations failed on 2nd step");
status = cuda_cub::synchronize(policy);
cuda_cub::throw_on_error(status, "set_operations failed to synchronize");
- size_type output_count = cuda_cub::get_value(policy, d_output_count);
+ std::size_t output_count = cuda_cub::get_value(policy, d_output_count);
return thrust::make_pair(keys_output + output_count, values_output + output_count);
}

The great Thrust index type fix, part 8: set operations.