Merge branch 'branch-25.04' into batch-nnd

cpp/src/cluster/detail/kmeans.cuh

@@ -45,6 +45,7 @@
 
 #include <cuda.h>
 #include <thrust/fill.h>
+#include <thrust/iterator/transform_iterator.h>
 #include <thrust/transform.h>
 
 #include <algorithm>
@@ -450,9 +451,8 @@ void kmeans_fit_main(raft::resources const& handle,
     // raft::KeyValuePair and converting them to just return the Key to be used
     // in reduce_rows_by_key prims
     cuvs::cluster::kmeans::detail::KeyValueIndexOp<IndexT, DataT> conversion_op;
-    cub::TransformInputIterator<IndexT,
-                                cuvs::cluster::kmeans::detail::KeyValueIndexOp<IndexT, DataT>,
-                                raft::KeyValuePair<IndexT, DataT>*>
+    thrust::transform_iterator<cuvs::cluster::kmeans::detail::KeyValueIndexOp<IndexT, DataT>,
+                               raft::KeyValuePair<IndexT, DataT>*>
       itr(minClusterAndDistance.data_handle(), conversion_op);
 
     update_centroids(handle,

cpp/src/cluster/detail/kmeans_balanced.cuh

@@ -50,6 +50,7 @@
 #include <rmm/resource_ref.hpp>
 
 #include <thrust/gather.h>
+#include <thrust/iterator/transform_iterator.h>
 #include <thrust/transform.h>
 
 #include <limits>
@@ -341,7 +342,7 @@ void calc_centers_and_sizes(const raft::resources& handle,
       dataset, dim, labels, nullptr, n_rows, dim, n_clusters, centers, stream, reset_counters);
   } else {
     // todo(lsugy): use iterator from KV output of fusedL2NN
-    cub::TransformInputIterator<MathT, MappingOpT, const T*> mapping_itr(dataset, mapping_op);
+    thrust::transform_iterator<MappingOpT, const T*> mapping_itr(dataset, mapping_op);
     raft::linalg::reduce_rows_by_key(
       mapping_itr, dim, labels, nullptr, n_rows, dim, n_clusters, centers, stream, reset_counters);
   }
@@ -966,7 +967,7 @@ auto build_fine_clusters(const raft::resources& handle,
                    "Number of fine clusters must be non-zero for a non-empty mesocluster");
     }
 
-    cub::TransformInputIterator<MathT, MappingOpT, const T*> mapping_itr(dataset_mptr, mapping_op);
+    thrust::transform_iterator<MappingOpT, const T*> mapping_itr(dataset_mptr, mapping_op);
     raft::matrix::gather(mapping_itr, dim, n_rows, mc_trainset_ids, k, mc_trainset, stream);
     if (params.metric == cuvs::distance::DistanceType::L2Expanded ||
         params.metric == cuvs::distance::DistanceType::L2SqrtExpanded ||

cpp/src/cluster/detail/kmeans_common.cuh

@@ -45,6 +45,7 @@
 #include <cuda.h>
 #include <thrust/fill.h>
 #include <thrust/for_each.h>
+#include <thrust/iterator/transform_iterator.h>
 
 #include <algorithm>
 #include <cmath>
@@ -199,8 +200,7 @@ void computeClusterCost(raft::resources const& handle,
 {
   cudaStream_t stream = raft::resource::get_cuda_stream(handle);
 
-  cub::TransformInputIterator<OutputT, MainOpT, InputT*> itr(minClusterDistance.data_handle(),
-                                                             main_op);
+  thrust::transform_iterator<MainOpT, InputT*> itr(minClusterDistance.data_handle(), main_op);
 
   size_t temp_storage_bytes = 0;
   RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(nullptr,
@@ -659,9 +659,8 @@ void countSamplesInCluster(raft::resources const& handle,
   // and converting them to just return the Key to be used in reduce_rows_by_key
   // prims
   cuvs::cluster::kmeans::detail::KeyValueIndexOp<IndexT, DataT> conversion_op;
-  cub::TransformInputIterator<IndexT,
-                              cuvs::cluster::kmeans::detail::KeyValueIndexOp<IndexT, DataT>,
-                              raft::KeyValuePair<IndexT, DataT>*>
+  thrust::transform_iterator<cuvs::cluster::kmeans::detail::KeyValueIndexOp<IndexT, DataT>,
+                             raft::KeyValuePair<IndexT, DataT>*>
     itr(minClusterAndDistance.data_handle(), conversion_op);
 
   // count # of samples in each cluster

cpp/src/cluster/detail/kmeans_mg.cuh

@@ -41,6 +41,7 @@
 #include <cuda/functional>
 #include <thrust/execution_policy.h>
 #include <thrust/fill.h>
+#include <thrust/iterator/transform_iterator.h>
 #include <thrust/reduce.h>
 #include <thrust/scan.h>
 #include <thrust/transform.h>
@@ -615,9 +616,7 @@ void fit(const raft::resources& handle,
     // cub::KeyValuePair and converting them to just return the Key to be used
     // in reduce_rows_by_key prims
     KeyValueIndexOp<IndexT, DataT> conversion_op;
-    cub::TransformInputIterator<IndexT,
-                                KeyValueIndexOp<IndexT, DataT>,
-                                raft::KeyValuePair<IndexT, DataT>*>
+    thrust::transform_iterator<KeyValueIndexOp<IndexT, DataT>, raft::KeyValuePair<IndexT, DataT>*>
       itr(minClusterAndDistance.data_handle(), conversion_op);
 
     workspace.resize(n_samples, stream);

cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh

@@ -62,6 +62,7 @@
 
 #include <cuda_fp16.h>
 #include <thrust/extrema.h>
+#include <thrust/iterator/transform_iterator.h>
 #include <thrust/scan.h>
 
 #include <memory>
@@ -96,8 +97,8 @@ void select_residuals(raft::resources const& handle,
   rmm::device_uvector<float> tmp(size_t(n_rows) * size_t(dim), stream, device_memory);
   // Note: the number of rows of the input dataset isn't actually n_rows, but matrix::gather doesn't
   // need to know it, any strictly positive number would work.
-  cub::TransformInputIterator<float, utils::mapping<float>, const T*> mapping_itr(
-    dataset, utils::mapping<float>{});
+  thrust::transform_iterator<utils::mapping<float>, const T*, thrust::use_default, float>
+    mapping_itr(dataset, utils::mapping<float>{});
   raft::matrix::gather(mapping_itr, (IdxT)dim, n_rows, row_ids, n_rows, tmp.data(), stream);
 
   raft::matrix::linewise_op(handle,

cpp/tests/cluster/kmeans.cu

@@ -28,6 +28,7 @@
 #include <rmm/device_uvector.hpp>
 
 #include <thrust/fill.h>
+#include <thrust/iterator/transform_iterator.h>
 
 #include <gtest/gtest.h>
 
@@ -201,9 +202,8 @@ class KmeansTest : public ::testing::TestWithParam<KmeansInputs<T>> {
   //      params.batch_centroids,
   //      workspace);
   //    raft::cluster::kmeans::KeyValueIndexOp<int, T> conversion_op;
-  //    cub::TransformInputIterator<int,
-  //                                raft::cluster::kmeans::KeyValueIndexOp<int, T>,
-  //                                raft::KeyValuePair<int, T>*>
+  //    thrust::transform_iterator<raft::cluster::kmeans::KeyValueIndexOp<int, T>,
+  //                               raft::KeyValuePair<int, T>*>
   //      itr(minClusterAndDistance.data_handle(), conversion_op);
   //
   //    auto sampleCountInCluster = raft::make_device_vector<T, int>(handle, params.n_clusters);