Use cuda/iterator in cub/test

This PR replaces thrust::constant_iterator with cuda::constant_iterator.
Same for counting_iterator and transform_iterator.

Replacing thrust::zip_iterator with cuda::zip_iterator was not always
successful, see #6400

Replacing thrust::make_permutation_iterator with cuda::make_permutation_iterator
also ran into compilation errors.

Author: oleksandr-pavlyk

cub/test/catch2_large_array_sort_helper.cuh

@@ -7,13 +7,12 @@
 
 #include <thrust/fill.h>
 #include <thrust/functional.h>
-#include <thrust/iterator/counting_iterator.h>
-#include <thrust/iterator/transform_iterator.h>
 #include <thrust/memory.h>
 #include <thrust/random.h>
 #include <thrust/shuffle.h>
 #include <thrust/tabulate.h>
 
+#include <cuda/iterator>
 #include <cuda/std/iterator>
 #include <cuda/std/limits>
 #include <cuda/std/type_traits>
@@ -184,9 +183,8 @@ struct large_array_sort_helper
     TIME(timer.print_elapsed_seconds_and_reset("Device Alloc"));
 
     { // Place the sorted keys into keys_out
-      auto key_iter = thrust::make_transform_iterator(
-        thrust::make_counting_iterator(std::size_t{0}),
-        detail::key_sort_ref_key_transform<KeyType>(num_items, is_descending));
+      auto key_iter = cuda::make_transform_iterator(
+        cuda::counting_iterator(std::size_t{0}), detail::key_sort_ref_key_transform<KeyType>(num_items, is_descending));
       thrust::copy(c2h::device_policy, key_iter, key_iter + num_items, keys_out.begin());
     }
 
@@ -212,9 +210,8 @@ struct large_array_sort_helper
   void verify_unstable_key_sort(std::size_t num_items, bool is_descending, const c2h::device_vector<KeyType>& keys)
   {
     TIME(c2h::cpu_timer timer);
-    auto key_iter = thrust::make_transform_iterator(
-      thrust::make_counting_iterator(std::size_t{0}),
-      detail::key_sort_ref_key_transform<KeyType>{num_items, is_descending});
+    auto key_iter = cuda::make_transform_iterator(
+      cuda::counting_iterator(std::size_t{0}), detail::key_sort_ref_key_transform<KeyType>{num_items, is_descending});
     REQUIRE(thrust::equal(c2h::device_policy, keys.cbegin(), keys.cend(), key_iter));
     TIME(timer.print_elapsed_seconds_and_reset("Validate keys"));
   }
@@ -352,8 +349,8 @@ struct large_array_sort_helper
 
     TIME(c2h::cpu_timer timer);
 
-    auto ref_key_begin = thrust::make_transform_iterator(
-      thrust::make_counting_iterator(std::size_t{0}),
+    auto ref_key_begin = cuda::make_transform_iterator(
+      cuda::counting_iterator(std::size_t{0}),
       detail::pair_sort_ref_key_transform<KeyType>(num_items, num_summaries, is_descending));
 
     REQUIRE(thrust::equal(c2h::device_policy, keys.cbegin(), keys.cend(), ref_key_begin));

cub/test/catch2_large_problem_helper.cuh

@@ -6,9 +6,6 @@
 #include <cub/util_type.cuh>
 
 #include <thrust/equal.h>
-#include <thrust/iterator/constant_iterator.h>
-#include <thrust/iterator/counting_iterator.h>
-#include <thrust/iterator/transform_iterator.h>
 
 #include <cuda/iterator>
 #include <cuda/std/__algorithm/clamp.h>
@@ -43,8 +40,8 @@ struct concat_iterators_op
 template <typename FirstSegmentItT, typename SecondSegmentItT>
 auto make_concat_iterators_op(FirstSegmentItT first_it, SecondSegmentItT second_it, ::cuda::std::int64_t num_first_items)
 {
-  return thrust::make_transform_iterator(
-    thrust::make_counting_iterator(::cuda::std::int64_t{0}),
+  return cuda::make_transform_iterator(
+    cuda::counting_iterator(::cuda::std::int64_t{0}),
     concat_iterators_op<FirstSegmentItT, SecondSegmentItT>{first_it, second_it, num_first_items});
 }
 
@@ -108,7 +105,7 @@ struct large_problem_test_helper
   {
     auto correctness_flags_end = correctness_flags.cbegin() + (num_elements / bits_per_element);
     const bool all_correct =
-      thrust::equal(correctness_flags.cbegin(), correctness_flags_end, thrust::make_constant_iterator(0xFFFFFFFFU));
+      thrust::equal(correctness_flags.cbegin(), correctness_flags_end, cuda::constant_iterator(0xFFFFFFFFU));
 
     if (!all_correct)
     {

cub/test/catch2_segmented_sort_helper.cuh

@@ -6,15 +6,14 @@
 
 #include <thrust/device_ptr.h>
 #include <thrust/for_each.h>
-#include <thrust/iterator/constant_iterator.h>
-#include <thrust/iterator/counting_iterator.h>
 #include <thrust/logical.h>
 #include <thrust/random.h>
 #include <thrust/scan.h>
 #include <thrust/sequence.h>
 #include <thrust/sort.h>
 #include <thrust/unique.h>
 
+#include <cuda/iterator>
 #include <cuda/std/limits>
 #include <cuda/std/tuple>
 #include <cuda/std/type_traits>
@@ -115,7 +114,7 @@ public:
     REQUIRE(count <= max_histo_size);
 
     // Verify keys are sorted using prior histogram computation
-    auto index_it = thrust::make_counting_iterator(std::size_t{0});
+    auto index_it = cuda::counting_iterator(std::size_t{0});
     c2h::device_vector<key_t> unique_keys_out(count);
     c2h::device_vector<std::size_t> unique_indexes_out(count);
     thrust::unique_by_key_copy(
@@ -180,8 +179,7 @@ public:
 
   void prepare_input_data(c2h::device_vector<key_t>& in_keys) const
   {
-    auto data_gen_it =
-      thrust::make_transform_iterator(thrust::make_counting_iterator(std::size_t{0}), mod_n<key_t>{sequence_length});
+    auto data_gen_it = cuda::transform_iterator(cuda::counting_iterator(std::size_t{0}), mod_n<key_t>{sequence_length});
     thrust::copy_n(data_gen_it, in_keys.size(), in_keys.begin());
   }
 
@@ -197,7 +195,7 @@ public:
     REQUIRE(count <= sequence_length * num_segments);
 
     // // Verify keys are sorted using prior histogram computation
-    auto index_it = thrust::make_counting_iterator(std::size_t{0});
+    auto index_it = cuda::counting_iterator(std::size_t{0});
     c2h::device_vector<key_t> unique_keys_out(count);
     c2h::device_vector<std::size_t> unique_indexes_out(count);
     thrust::unique_by_key_copy(
@@ -510,15 +508,15 @@ void generate_unsorted_derived_inputs(
 
   // Build keys in reversed order from how they'll eventually be sorted:
   thrust::for_each(c2h::nosync_device_policy,
-                   thrust::make_counting_iterator(0),
-                   thrust::make_counting_iterator(num_segments),
+                   cuda::counting_iterator(0),
+                   cuda::counting_iterator(num_segments),
                    segment_filler<KeyT>{keys, offsets, !descending_sort});
   if constexpr (sort_pairs)
   {
     // Values are generated in reversed order from keys:
     thrust::for_each(c2h::nosync_device_policy,
-                     thrust::make_counting_iterator(0),
-                     thrust::make_counting_iterator(num_segments),
+                     cuda::counting_iterator(0),
+                     cuda::counting_iterator(num_segments),
                      segment_filler<ValueT>{values, offsets, descending_sort});
   }
 
@@ -544,8 +542,8 @@ void validate_sorted_derived_outputs(
   const int* offsets     = thrust::raw_pointer_cast(d_offsets.data());
 
   REQUIRE(thrust::all_of(c2h::device_policy,
-                         thrust::make_counting_iterator(0),
-                         thrust::make_counting_iterator(num_segments),
+                         cuda::counting_iterator(0),
+                         cuda::counting_iterator(num_segments),
                          segment_checker<KeyT, ValueT, STABLE>{keys, values, offsets, descending_sort}));
 }
 
@@ -734,8 +732,8 @@ void validate_sorted_random_outputs(
 
       REQUIRE(thrust::all_of(
         c2h::device_policy,
-        thrust::make_counting_iterator(0),
-        thrust::make_counting_iterator(num_segments),
+        cuda::counting_iterator(0),
+        cuda::counting_iterator(num_segments),
         unstable_segmented_value_checker<KeyT, ValueT>{
           ref_keys, ref_values, test_values, d_segment_begin, d_segment_end}));
     }
@@ -1642,7 +1640,7 @@ inline int generate_unspecified_segments_offsets(
   // calculation below.
   c2h::gen(make_offset_eraser_seed(seed), erase_indices, 1, num_segments - 2);
 
-  auto const_zero_begin = thrust::make_constant_iterator<int>(0);
+  auto const_zero_begin = cuda::constant_iterator<int>(0);
   auto const_zero_end   = const_zero_begin + erase_indices.size();
 
   thrust::scatter(

cub/test/catch2_test_block_run_length_decode.cu

@@ -7,9 +7,7 @@
 #include <cub/device/device_scan.cuh>
 #include <cub/util_allocator.cuh>
 
-#include <thrust/iterator/counting_iterator.h>
-#include <thrust/iterator/transform_iterator.h>
-
+#include <cuda/iterator>
 #include <cuda/std/type_traits>
 
 #include <c2h/catch2_test_helper.h>
@@ -294,6 +292,16 @@ struct ModOp
   }
 };
 
+template <typename TargetT>
+struct CastOp
+{
+  template <typename U>
+  __host__ __device__ __forceinline__ TargetT operator()(const U& x) const
+  {
+    return static_cast<TargetT>(x);
+  }
+};
+
 template <uint32_t RUNS_PER_THREAD,
           uint32_t DECODED_ITEMS_PER_THREAD,
           uint32_t BLOCK_DIM_X,
@@ -308,11 +316,12 @@ void TestAlgorithmSpecialisation()
 
   using RunItemT      = float;
   using RunLengthT    = uint32_t;
-  using ItemItT       = thrust::counting_iterator<RunItemT>;
-  using RunLengthsItT = thrust::transform_iterator<ModOp, thrust::counting_iterator<RunLengthT>>;
+  using CasterOp      = CastOp<RunItemT>;
+  using ItemItT       = cuda::transform_iterator<CasterOp, cuda::counting_iterator<uint64_t>>;
+  using RunLengthsItT = cuda::transform_iterator<ModOp, cuda::counting_iterator<RunLengthT>>;
 
-  ItemItT d_unique_items(1000U);
-  RunLengthsItT d_run_lengths(thrust::counting_iterator<RunLengthT>(0), ModOp{});
+  ItemItT d_unique_items(cuda::counting_iterator<uint64_t>(1000U), CasterOp{});
+  RunLengthsItT d_run_lengths(cuda::counting_iterator<RunLengthT>(0), ModOp{});
 
   constexpr uint32_t num_runs   = 10000;
   constexpr uint32_t num_blocks = (num_runs + (RUNS_PER_BLOCK - 1U)) / RUNS_PER_BLOCK;

cub/test/catch2_test_device_adjacent_difference_substract_left.cu

@@ -5,7 +5,7 @@
 
 #include <cub/device/device_adjacent_difference.cuh>
 
-#include <thrust/iterator/discard_iterator.h>
+#include <cuda/iterator>
 
 #include <algorithm>
 #include <numeric>
@@ -59,7 +59,7 @@ C2H_TEST("DeviceAdjacentDifference::SubtractLeftCopy does not change the input",
   c2h::gen(C2H_SEED(2), in);
 
   c2h::device_vector<type> reference = in;
-  adjacent_difference_subtract_left_copy(in.begin(), thrust::discard_iterator<>(), num_items, cuda::std::minus<>{});
+  adjacent_difference_subtract_left_copy(in.begin(), cuda::discard_iterator(), num_items, cuda::std::minus<>{});
 
   REQUIRE(reference == in);
 }
@@ -252,7 +252,7 @@ C2H_TEST("DeviceAdjacentDifference::SubtractLeftCopy works with large indexes",
   c2h::device_vector<int> error(1);
   int* d_error = thrust::raw_pointer_cast(error.data());
   adjacent_difference_subtract_left_copy(
-    thrust::counting_iterator<cuda::std::size_t>{0}, thrust::discard_iterator<>{}, num_items, check_difference{d_error});
+    cuda::counting_iterator<cuda::std::size_t>{0}, cuda::discard_iterator{}, num_items, check_difference{d_error});
   const int h_error = error[0];
   REQUIRE(h_error == 0);
 }
@@ -280,8 +280,8 @@ C2H_TEST("DeviceAdjacentDifference::SubtractLeftCopy uses right number of invoca
   const int num_items = GENERATE_COPY(take(2, random(1, 1000000)));
   c2h::device_vector<unsigned long long> counts(1, 0);
   adjacent_difference_subtract_left_copy(
-    thrust::counting_iterator<cuda::std::size_t>{0},
-    thrust::discard_iterator<>(),
+    cuda::counting_iterator<cuda::std::size_t>{0},
+    cuda::discard_iterator(),
     num_items,
     invocation_counter{thrust::raw_pointer_cast(counts.data())});
 

cub/test/catch2_test_device_adjacent_difference_substract_right.cu

@@ -5,7 +5,7 @@
 
 #include <cub/device/device_adjacent_difference.cuh>
 
-#include <thrust/iterator/discard_iterator.h>
+#include <cuda/iterator>
 
 #include <algorithm>
 #include <numeric>
@@ -59,7 +59,7 @@ C2H_TEST("DeviceAdjacentDifference::SubtractRightCopy does not change the input"
   c2h::gen(C2H_SEED(2), in);
 
   c2h::device_vector<type> reference = in;
-  adjacent_difference_subtract_right_copy(in.begin(), thrust::discard_iterator<>(), num_items, cuda::std::minus<>{});
+  adjacent_difference_subtract_right_copy(in.begin(), cuda::discard_iterator(), num_items, cuda::std::minus<>{});
 
   REQUIRE(reference == in);
 }
@@ -304,7 +304,7 @@ C2H_TEST("DeviceAdjacentDifference::SubtractRightCopy works with large indexes",
   c2h::device_vector<int> error(1);
   int* d_error = thrust::raw_pointer_cast(error.data());
   adjacent_difference_subtract_right_copy(
-    thrust::counting_iterator<cuda::std::size_t>{0}, thrust::discard_iterator<>{}, num_items, check_difference{d_error});
+    cuda::counting_iterator<cuda::std::size_t>{0}, cuda::discard_iterator{}, num_items, check_difference{d_error});
   const int h_error = error[0];
   REQUIRE(h_error == 0);
 }
@@ -333,8 +333,8 @@ C2H_TEST("DeviceAdjacentDifference::SubtractRightCopy uses right number of invoc
   const int num_items = GENERATE_COPY(take(2, random(1, 1000000)));
   c2h::device_vector<unsigned long long> counts(1, 0);
   adjacent_difference_subtract_right_copy(
-    thrust::counting_iterator<cuda::std::size_t>{0},
-    thrust::discard_iterator<>(),
+    cuda::counting_iterator<cuda::std::size_t>{0},
+    cuda::discard_iterator(),
     num_items,
     invocation_counter{thrust::raw_pointer_cast(counts.data())});