Merge branch 'branch-25.04' into rhdong/logical-merge

ci/build_cpp.sh

@@ -17,7 +17,7 @@ rapids-logger "Begin cpp build"
 
 sccache --zero-stats
 
-RAPIDS_PACKAGE_VERSION=$(rapids-generate-version) rapids-conda-retry mambabuild \
+RAPIDS_PACKAGE_VERSION=$(rapids-generate-version) rapids-conda-retry build \
   conda/recipes/libcuvs
 
 sccache --show-adv-stats

ci/build_python.sh

@@ -25,7 +25,7 @@ sccache --zero-stats
 
 # TODO: Remove `--no-test` flags once importing on a CPU
 # node works correctly
-rapids-conda-retry mambabuild \
+rapids-conda-retry build \
   --no-test \
   --channel "${CPP_CHANNEL}" \
   conda/recipes/cuvs
@@ -34,7 +34,7 @@ sccache --show-adv-stats
 sccache --zero-stats
 
 # Build cuvs-bench for each cuda and python version
-rapids-conda-retry mambabuild \
+rapids-conda-retry build \
   --no-test \
   --channel "${CPP_CHANNEL}" \
   --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
@@ -47,7 +47,7 @@ sccache --zero-stats
 # version
 RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}"
 if [[ ${RAPIDS_CUDA_MAJOR} == "12" ]]; then
-  rapids-conda-retry mambabuild \
+  rapids-conda-retry build \
   --no-test \
   --channel "${CPP_CHANNEL}" \
   --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \

conda/environments/all_cuda-118_arch-aarch64.yaml

@@ -11,7 +11,7 @@ dependencies:
 - c-compiler
 - clang-tools==16.0.6
 - clang==16.0.6
-- cmake>=3.26.4,!=3.30.0
+- cmake>=3.30.4
 - cuda-nvtx=11.8
 - cuda-profiler-api=11.8.86
 - cuda-python>=11.8.5,<12.0a0

conda/environments/all_cuda-118_arch-x86_64.yaml

@@ -11,7 +11,7 @@ dependencies:
 - c-compiler
 - clang-tools==16.0.6
 - clang==16.0.6
-- cmake>=3.26.4,!=3.30.0
+- cmake>=3.30.4
 - cuda-nvtx=11.8
 - cuda-profiler-api=11.8.86
 - cuda-python>=11.8.5,<12.0a0

conda/environments/all_cuda-128_arch-aarch64.yaml

@@ -11,7 +11,7 @@ dependencies:
 - c-compiler
 - clang-tools==16.0.6
 - clang==16.0.6
-- cmake>=3.26.4,!=3.30.0
+- cmake>=3.30.4
 - cuda-cudart-dev
 - cuda-nvcc
 - cuda-nvtx-dev

conda/environments/all_cuda-128_arch-x86_64.yaml

@@ -11,7 +11,7 @@ dependencies:
 - c-compiler
 - clang-tools==16.0.6
 - clang==16.0.6
-- cmake>=3.26.4,!=3.30.0
+- cmake>=3.30.4
 - cuda-cudart-dev
 - cuda-nvcc
 - cuda-nvtx-dev

conda/environments/bench_ann_cuda-118_arch-aarch64.yaml

@@ -12,7 +12,7 @@ dependencies:
 - clang-tools==16.0.6
 - clang==16.0.6
 - click
-- cmake>=3.26.4,!=3.30.0
+- cmake>=3.30.4
 - cuda-nvtx=11.8
 - cuda-profiler-api=11.8.86
 - cuda-python>=11.8.5,<12.0a0

conda/environments/bench_ann_cuda-118_arch-x86_64.yaml

@@ -12,7 +12,7 @@ dependencies:
 - clang-tools==16.0.6
 - clang==16.0.6
 - click
-- cmake>=3.26.4,!=3.30.0
+- cmake>=3.30.4
 - cuda-nvtx=11.8
 - cuda-profiler-api=11.8.86
 - cuda-python>=11.8.5,<12.0a0

conda/environments/bench_ann_cuda-128_arch-aarch64.yaml

@@ -12,7 +12,7 @@ dependencies:
 - clang-tools==16.0.6
 - clang==16.0.6
 - click
-- cmake>=3.26.4,!=3.30.0
+- cmake>=3.30.4
 - cuda-cudart-dev
 - cuda-nvcc
 - cuda-nvtx-dev

conda/environments/bench_ann_cuda-128_arch-x86_64.yaml

@@ -12,7 +12,7 @@ dependencies:
 - clang-tools==16.0.6
 - clang==16.0.6
 - click
-- cmake>=3.26.4,!=3.30.0
+- cmake>=3.30.4
 - cuda-cudart-dev
 - cuda-nvcc
 - cuda-nvtx-dev

conda/recipes/cuvs-bench-cpu/conda_build_config.yaml

@@ -13,7 +13,7 @@ c_stdlib_version:
   - "2.28"
 
 cmake_version:
-  - ">=3.26.4,!=3.30.0"
+  - ">=3.30.4"
 
 glog_version:
   - ">=0.6.0"

conda/recipes/cuvs-bench-cpu/meta.yaml

@@ -62,6 +62,7 @@ requirements:
     - pyyaml
     - python
     - requests
+    - sklearn>=1.5
 about:
   home: https://rapids.ai/
   license: Apache-2.0

conda/recipes/cuvs-bench/conda_build_config.yaml

@@ -17,7 +17,7 @@ c_stdlib_version:
   - "2.28"
 
 cmake_version:
-  - ">=3.26.4,!=3.30.0"
+  - ">=3.30.4"
 
 nccl_version:
   - ">=2.19"

conda/recipes/cuvs-bench/meta.yaml

@@ -101,6 +101,7 @@ requirements:
     - python
     - requests
     - rmm ={{ minor_version }}
+    - sklearn>=1.5
 about:
   home: https://rapids.ai/
   license: Apache-2.0

conda/recipes/cuvs/conda_build_config.yaml

@@ -17,4 +17,4 @@ c_stdlib_version:
   - "2.28"
 
 cmake_version:
-  - ">=3.26.4,!=3.30.0"
+  - ">=3.30.4"

conda/recipes/libcuvs/conda_build_config.yaml

@@ -17,7 +17,7 @@ c_stdlib_version:
   - "2.28"
 
 cmake_version:
-  - ">=3.26.4,!=3.30.0"
+  - ">=3.30.4"
 
 h5py_version:
   - ">=3.8.0"

cpp/CMakeLists.txt

@@ -10,7 +10,7 @@
 # is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 # or implied. See the License for the specific language governing permissions and limitations under
 # the License.
-cmake_minimum_required(VERSION 3.26.4 FATAL_ERROR)
+cmake_minimum_required(VERSION 3.30.4 FATAL_ERROR)
 include(../rapids_config.cmake)
 include(rapids-cmake)
 include(rapids-cpm)
@@ -636,6 +636,7 @@ SECTIONS
   if(CUVS_NVTX)
     # This enables NVTX within the project with no option to disable it downstream.
     target_link_libraries(cuvs PUBLIC CUDA::nvtx3)
+    target_compile_definitions(cuvs_objs PUBLIC NVTX_ENABLED)
     target_compile_definitions(cuvs PUBLIC NVTX_ENABLED)
 
     target_link_libraries(cuvs-cagra-search PUBLIC CUDA::nvtx3)

cpp/include/cuvs/neighbors/cagra.hpp

@@ -343,6 +343,8 @@ struct index : cuvs::neighbors::index {
   using search_params_type = cagra::search_params;
   using index_type         = IdxT;
   using value_type         = T;
+  using dataset_index_type = int64_t;
+
   static_assert(!raft::is_narrowing_v<uint32_t, IdxT>,
                 "IdxT must be able to represent all values of uint32_t");
 
@@ -516,14 +518,14 @@ struct index : cuvs::neighbors::index {
    */
   template <typename DatasetT>
   auto update_dataset(raft::resources const& res, DatasetT&& dataset)
-    -> std::enable_if_t<std::is_base_of_v<cuvs::neighbors::dataset<int64_t>, DatasetT>>
+    -> std::enable_if_t<std::is_base_of_v<cuvs::neighbors::dataset<dataset_index_type>, DatasetT>>
   {
     dataset_ = std::make_unique<DatasetT>(std::move(dataset));
   }
 
   template <typename DatasetT>
   auto update_dataset(raft::resources const& res, std::unique_ptr<DatasetT>&& dataset)
-    -> std::enable_if_t<std::is_base_of_v<neighbors::dataset<int64_t>, DatasetT>>
+    -> std::enable_if_t<std::is_base_of_v<neighbors::dataset<dataset_index_type>, DatasetT>>
   {
     dataset_ = std::move(dataset);
   }
@@ -567,7 +569,7 @@ struct index : cuvs::neighbors::index {
   cuvs::distance::DistanceType metric_;
   raft::device_matrix<IdxT, int64_t, raft::row_major> graph_;
   raft::device_matrix_view<const IdxT, int64_t, raft::row_major> graph_view_;
-  std::unique_ptr<neighbors::dataset<int64_t>> dataset_;
+  std::unique_ptr<neighbors::dataset<dataset_index_type>> dataset_;
 };
 /**
  * @}

cpp/src/neighbors/detail/cagra/cagra_merge.cuh

@@ -43,14 +43,16 @@ index<T, IdxT> merge(raft::resources const& handle,
                      const cagra::merge_params& params,
                      std::vector<cuvs::neighbors::cagra::index<T, IdxT>*>& indices)
 {
+  using cagra_index_t = cuvs::neighbors::cagra::index<T, IdxT>;
+  using ds_idx_type   = typename cagra_index_t::dataset_index_type;
+
   std::size_t dim              = 0;
   std::size_t new_dataset_size = 0;
   int64_t stride               = -1;
 
-  for (auto index : indices) {
+  for (cagra_index_t* index : indices) {
     RAFT_EXPECTS(index != nullptr,
                  "Null pointer detected in 'indices'. Ensure all elements are valid before usage.");
-    using ds_idx_type = decltype(index->data().n_rows());
     if (auto* strided_dset = dynamic_cast<const strided_dataset<T, ds_idx_type>*>(&index->data());
         strided_dset != nullptr) {
       if (dim == 0) {
@@ -74,8 +76,7 @@ index<T, IdxT> merge(raft::resources const& handle,
   IdxT offset = 0;
 
   auto merge_dataset = [&](T* dst) {
-    for (auto index : indices) {
-      using ds_idx_type  = decltype(index->data().n_rows());
+    for (cagra_index_t* index : indices) {
       auto* strided_dset = dynamic_cast<const strided_dataset<T, ds_idx_type>*>(&index->data());
 
       RAFT_CUDA_TRY(cudaMemcpy2DAsync(dst + offset * dim,

cpp/src/neighbors/detail/cagra/graph_core.cuh

@@ -24,6 +24,7 @@
 #include <raft/core/resources.hpp>
 
 // TODO: This shouldn't be invoking anything from spatial/knn
+#include "../../../core/nvtx.hpp"
 #include "../ann_utils.cuh"
 
 #include <raft/util/bitonic_sort.cuh>
@@ -1086,6 +1087,8 @@ void optimize(
   const uint64_t graph_size          = new_graph.extent(0);
   auto input_graph_ptr               = knn_graph.data_handle();
   auto output_graph_ptr              = new_graph.data_handle();
+  raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> fun_scope(
+    "cagra::graph::optimize(%zu, %zu, %u)", graph_size, input_graph_degree, output_graph_degree);
 
   // MST optimization
   auto mst_graph               = raft::make_host_matrix<IdxT, int64_t, raft::row_major>(0, 0);
@@ -1096,6 +1099,8 @@ void optimize(
     mst_graph_num_edges_ptr[i] = 0;
   }
   if (guarantee_connectivity) {
+    raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> block_scope(
+      "cagra::graph::optimize/check_connectivity");
     mst_graph =
       raft::make_host_matrix<IdxT, int64_t, raft::row_major>(graph_size, output_graph_degree);
     RAFT_LOG_INFO("MST optimization is used to guarantee graph connectivity.");
@@ -1110,6 +1115,8 @@ void optimize(
   }
 
   {
+    raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> block_scope(
+      "cagra::graph::optimize/prune");
     //
     // Prune kNN graph
     //
@@ -1270,6 +1277,8 @@ void optimize(
   auto rev_graph_count = raft::make_host_vector<uint32_t, int64_t>(graph_size);
 
   {
+    raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> block_scope(
+      "cagra::graph::optimize/reverse");
     //
     // Make reverse graph
     //
@@ -1335,6 +1344,8 @@ void optimize(
   }
 
   {
+    raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> block_scope(
+      "cagra::graph::optimize/combine");
     //
     // Create search graphs from MST and pruned and reverse graphs
     //
@@ -1435,6 +1446,8 @@ void optimize(
 
   // Check number of incoming edges
   {
+    raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> block_scope(
+      "cagra::graph::optimize/check_edges");
     auto in_edge_count     = raft::make_host_vector<uint32_t, int64_t>(graph_size);
     auto in_edge_count_ptr = in_edge_count.data_handle();
 #pragma omp parallel for
@@ -1477,6 +1490,8 @@ void optimize(
 
   // Check duplication and out-of-range indices
   {
+    raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> block_scope(
+      "cagra::graph::optimize/check_duplicates");
     uint64_t num_dup = 0;
     uint64_t num_oor = 0;
 #pragma omp parallel for reduction(+ : num_dup) reduction(+ : num_oor)

cpp/src/neighbors/detail/nn_descent.cuh

@@ -1047,24 +1047,32 @@ void GnndGraph<Index_t>::init_random_graph()
   for (size_t seg_idx = 0; seg_idx < static_cast<size_t>(num_segments); seg_idx++) {
     // random sequence (range: 0~nrow)
     // segment_x stores neighbors which id % num_segments == x
-    std::vector<Index_t> rand_seq(nrow / num_segments);
+    std::vector<Index_t> rand_seq((nrow + num_segments - 1) / num_segments);
     std::iota(rand_seq.begin(), rand_seq.end(), 0);
     auto gen = std::default_random_engine{seg_idx};
     std::shuffle(rand_seq.begin(), rand_seq.end(), gen);
 
 #pragma omp parallel for
     for (size_t i = 0; i < nrow; i++) {
-      size_t base_idx      = i * node_degree + seg_idx * segment_size;
-      auto h_neighbor_list = h_graph + base_idx;
-      auto h_dist_list     = h_dists.data_handle() + base_idx;
+      size_t base_idx         = i * node_degree + seg_idx * segment_size;
+      auto h_neighbor_list    = h_graph + base_idx;
+      auto h_dist_list        = h_dists.data_handle() + base_idx;
+      size_t idx              = base_idx;
+      size_t self_in_this_seg = 0;
       for (size_t j = 0; j < static_cast<size_t>(segment_size); j++) {
-        size_t idx = base_idx + j;
         Index_t id = rand_seq[idx % rand_seq.size()] * num_segments + seg_idx;
         if ((size_t)id == i) {
-          id = rand_seq[(idx + segment_size) % rand_seq.size()] * num_segments + seg_idx;
+          idx++;
+          id               = rand_seq[idx % rand_seq.size()] * num_segments + seg_idx;
+          self_in_this_seg = 1;
         }
-        h_neighbor_list[j].id_with_flag() = id;
-        h_dist_list[j]                    = std::numeric_limits<DistData_t>::max();
+
+        h_neighbor_list[j].id_with_flag() =
+          j < (rand_seq.size() - self_in_this_seg) && size_t(id) < nrow
+            ? id
+            : std::numeric_limits<Index_t>::max();
+        h_dist_list[j] = std::numeric_limits<DistData_t>::max();
+        idx++;
       }
     }
   }

cpp/src/neighbors/hnsw.cpp

@@ -15,20 +15,22 @@
  */
 
 #include "detail/hnsw.hpp"
+#include "../core/nvtx.hpp"
 #include <cstdint>
 #include <cuvs/neighbors/hnsw.hpp>
 #include <sys/types.h>
 
 namespace cuvs::neighbors::hnsw {
 
-#define CUVS_INST_HNSW_FROM_CAGRA(T)                                                  \
-  std::unique_ptr<index<T>> from_cagra(                                               \
-    raft::resources const& res,                                                       \
-    const index_params& params,                                                       \
-    const cuvs::neighbors::cagra::index<T, uint32_t>& cagra_index,                    \
-    std::optional<raft::host_matrix_view<const T, int64_t, raft::row_major>> dataset) \
-  {                                                                                   \
-    return detail::from_cagra<T>(res, params, cagra_index, dataset);                  \
+#define CUVS_INST_HNSW_FROM_CAGRA(T)                                                           \
+  std::unique_ptr<index<T>> from_cagra(                                                        \
+    raft::resources const& res,                                                                \
+    const index_params& params,                                                                \
+    const cuvs::neighbors::cagra::index<T, uint32_t>& cagra_index,                             \
+    std::optional<raft::host_matrix_view<const T, int64_t, raft::row_major>> dataset)          \
+  {                                                                                            \
+    raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> fun_scope("hnsw::from_cagra"); \
+    return detail::from_cagra<T>(res, params, cagra_index, dataset);                           \
   }
 
 CUVS_INST_HNSW_FROM_CAGRA(float);