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

Author: griwes

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);

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);

thrust/system/cuda/detail/dispatch.h

@@ -34,6 +34,26 @@
         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

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);
   }