[REVIEW] Move from thrust::make_counting_iterator to cuda::make_counting_iterator

cpp/src/cluster/detail/agglomerative.cuh

@@ -13,14 +13,14 @@
 
 #include <rmm/device_uvector.hpp>
 
+#include <cuda/iterator>
 #include <cuda/std/tuple>
 #include <thrust/device_ptr.h>
 #include <thrust/execution_policy.h>
 #include <thrust/extrema.h>
 #include <thrust/fill.h>
 #include <thrust/for_each.h>
 #include <thrust/functional.h>
-#include <thrust/iterator/counting_iterator.h>
 #include <thrust/iterator/zip_iterator.h>
 #include <thrust/sort.h>
 
@@ -290,7 +290,7 @@ void extract_flattened_clusters(raft::resources const& handle,
     thrust::fill(thrust_policy, tmp_labels.data(), tmp_labels.data() + n_vertices, -1);
 
     // Write labels for cluster roots to "labels"
-    thrust::counting_iterator<uint> first(0);
+    auto first = cuda::make_counting_iterator<uint>(0);
 
     auto z_iter = thrust::make_zip_iterator(
       cuda::std::make_tuple(first, label_roots.data() + (label_roots.size() - n_clusters)));

cpp/src/distance/detail/sparse/coo_spmv_strategies/hash_strategy.cuh

@@ -12,8 +12,8 @@
 #include <raft/util/cuda_dev_essentials.cuh>
 
 #include <cuco/static_map.cuh>
+#include <cuda/iterator>
 #include <thrust/copy.h>
-#include <thrust/iterator/counting_iterator.h>
 
 #include <cooperative_groups.h>
 #include <rmm/device_uvector.hpp>
@@ -68,16 +68,16 @@ class hash_strategy : public coo_spmv_strategy<value_idx, value_t, tpb> {
     auto policy = raft::resource::get_thrust_policy(this->config.handle);
 
     auto less                   = thrust::copy_if(policy,
-                                thrust::make_counting_iterator(value_idx(0)),
-                                thrust::make_counting_iterator(n_rows),
+                                cuda::make_counting_iterator(value_idx(0)),
+                                cuda::make_counting_iterator(n_rows),
                                 mask_indptr.data(),
                                 fits_in_hash_table(indptr, 0, capacity_threshold * map_size));
     std::get<0>(n_rows_divided) = less - mask_indptr.data();
 
     auto more = thrust::copy_if(
       policy,
-      thrust::make_counting_iterator(value_idx(0)),
-      thrust::make_counting_iterator(n_rows),
+      cuda::make_counting_iterator(value_idx(0)),
+      cuda::make_counting_iterator(n_rows),
       less,
       fits_in_hash_table(
         indptr, capacity_threshold * map_size, std::numeric_limits<value_idx>::max()));

cpp/src/distance/detail/sparse/l2_distance.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -19,8 +19,8 @@
 
 #include <rmm/device_uvector.hpp>
 
+#include <cuda/iterator>
 #include <thrust/for_each.h>
-#include <thrust/iterator/counting_iterator.h>
 
 #include <algorithm>
 #include <nvfunctional>
@@ -470,7 +470,7 @@ class russelrao_expanded_distances_t : public distances_t<value_t> {
       raft::resource::get_cuda_stream(config_->handle));
 
     auto exec_policy  = rmm::exec_policy(raft::resource::get_cuda_stream(config_->handle));
-    auto diags        = thrust::counting_iterator<value_idx>(0);
+    auto diags        = cuda::make_counting_iterator<value_idx>(0);
     value_idx b_nrows = config_->b_nrows;
     thrust::for_each(exec_policy, diags, diags + config_->a_nrows, [=] __device__(value_idx input) {
       out_dists[input * b_nrows + input] = 0.0;

cpp/src/neighbors/detail/knn_brute_force.cuh

@@ -37,11 +37,11 @@
 #include <raft/util/cudart_utils.hpp>
 #include <raft/util/popc.cuh>
 
+#include <cuda/iterator>
 #include <cuda_fp16.h>
 #include <rmm/cuda_device.hpp>
 #include <rmm/device_uvector.hpp>
 #include <thrust/for_each.h>
-#include <thrust/iterator/counting_iterator.h>
 
 #include <cstdint>
 #include <iostream>
@@ -227,7 +227,7 @@ void tiled_brute_force_knn(const raft::resources& handle,
       }
 
       auto distances_ptr        = temp_distances.data();
-      auto count                = thrust::make_counting_iterator<IndexType>(0);
+      auto count                = cuda::make_counting_iterator<IndexType>(0);
       DistanceT masked_distance = select_min ? std::numeric_limits<DistanceT>::infinity()
                                              : std::numeric_limits<DistanceT>::lowest();
 
@@ -277,7 +277,7 @@ void tiled_brute_force_knn(const raft::resources& handle,
         DistanceT* out_distances      = temp_out_distances.data();
         IndexType* out_indices        = temp_out_indices.data();
 
-        auto count = thrust::make_counting_iterator<IndexType>(0);
+        auto count = cuda::make_counting_iterator<IndexType>(0);
         thrust::for_each(raft::resource::get_thrust_policy(handle),
                          count,
                          count + current_query_size * current_k,

cpp/src/sparse/neighbors/detail/cross_component_nn.cuh

@@ -26,11 +26,11 @@
 
 #include <rmm/device_uvector.hpp>
 
+#include <cuda/iterator>
 #include <cuda/std/tuple>
 #include <thrust/copy.h>
 #include <thrust/device_ptr.h>
 #include <thrust/gather.h>
-#include <thrust/iterator/counting_iterator.h>
 #include <thrust/iterator/zip_iterator.h>
 #include <thrust/scan.h>
 #include <thrust/scatter.h>
@@ -432,8 +432,8 @@ void sort_by_color(raft::resources const& handle,
                    value_idx* src_indices,
                    size_t n_rows)
 {
-  auto exec_policy = raft::resource::get_thrust_policy(handle);
-  thrust::counting_iterator<value_idx> arg_sort_iter(0);
+  auto exec_policy   = raft::resource::get_thrust_policy(handle);
+  auto arg_sort_iter = cuda::make_counting_iterator<value_idx>(0);
   thrust::copy(exec_policy, arg_sort_iter, arg_sort_iter + n_rows, src_indices);
 
   auto keys = thrust::make_zip_iterator(

cpp/tests/test_utils.cuh

@@ -7,6 +7,7 @@
 
 #include "test_utils.h"
 
+#include <cuda/iterator>
 #include <gtest/gtest.h>
 #include <iostream>
 #include <memory>
@@ -15,7 +16,6 @@
 #include <raft/util/cudart_utils.hpp>
 #include <rmm/exec_policy.hpp>
 #include <thrust/for_each.h>
-#include <thrust/iterator/counting_iterator.h>
 
 #include <fstream>
 #include <sstream>
@@ -260,7 +260,7 @@ void gen_uniform(const raft::resources& handle,
 
   const T1* d_keys   = keys.data();
   const T2* d_values = values.data();
-  auto counting      = thrust::make_counting_iterator<IdxT>(0);
+  auto counting      = cuda::make_counting_iterator<IdxT>(0);
   thrust::for_each(rmm::exec_policy(stream),
                    counting,
                    counting + len,

examples/cpp/src/common.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -17,7 +17,6 @@
 
 #include <thrust/copy.h>
 #include <thrust/device_ptr.h>
-#include <thrust/iterator/counting_iterator.h>
 
 #include <fstream>