Merge branch 'branch-25.06' into cagra-fix-nnd-init

enp1s0 · web-flow · commit 89e10d2eb3b8 · 2025-04-24T02:05:39.000+09:00
diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh
@@ -35,6 +35,14 @@ if [[ "${package_dir}" != "python/libcuvs" ]]; then
     )
 fi
 
+RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}"
+if [[ ${RAPIDS_CUDA_MAJOR} != "11" ]]; then
+    EXCLUDE_ARGS+=(
+      --exclude "libnccl.so.*"
+    )
+    export SKBUILD_CMAKE_ARGS="-DUSE_NCCL_RUNTIME_WHEEL=ON"
+fi
+
 rapids-logger "Building '${package_name}' wheel"
 
 sccache --zero-stats
diff --git a/ci/test_wheel_cuvs.sh b/ci/test_wheel_cuvs.sh
@@ -3,6 +3,12 @@
 
 set -euo pipefail
 
+# Delete system libnccl.so to ensure the wheel is used
+RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}"
+if [[ ${RAPIDS_CUDA_MAJOR} != "11" ]]; then
+  rm -rf /usr/lib64/libnccl*
+fi
+
 mkdir -p ./dist
 RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")"
 RAPIDS_PY_WHEEL_NAME="libcuvs_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 cpp ./local-libcuvs-dep
diff --git a/cpp/include/cuvs/core/detail/interop.hpp b/cpp/include/cuvs/core/detail/interop.hpp
@@ -141,4 +141,30 @@ inline bool is_c_contiguous(DLManagedTensor* managed_tensor)
   return true;
 }
 
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wunused-function"
+static void free_dlmanaged_tensor_shape(DLManagedTensor* tensor)
+{
+  delete[] tensor->dl_tensor.shape;
+}
+#pragma GCC diagnostic pop
+
+template <typename MdspanType, typename = raft::is_mdspan_t<MdspanType>>
+static void to_dlpack(MdspanType src, DLManagedTensor* dst)
+{
+  auto tensor = &dst->dl_tensor;
+
+  tensor->dtype  = data_type_to_DLDataType<typename MdspanType::value_type>();
+  tensor->device = accessor_type_to_DLDevice<typename MdspanType::accessor_type>();
+  tensor->ndim   = MdspanType::extents_type::rank();
+  tensor->data   = src.data_handle();
+
+  tensor->shape = new int64_t[tensor->ndim];
+  dst->deleter  = free_dlmanaged_tensor_shape;
+
+  for (int64_t i = 0; i < tensor->ndim; ++i) {
+    tensor->shape[i] = src.extent(i);
+  }
+}
+
 }  // namespace cuvs::core::detail
diff --git a/cpp/include/cuvs/core/interop.hpp b/cpp/include/cuvs/core/interop.hpp
@@ -68,7 +68,7 @@ inline bool is_c_contiguous(DLManagedTensor* tensor) { return detail::is_c_conti
 inline bool is_f_contiguous(DLManagedTensor* tensor) { return detail::is_f_contiguous(tensor); }
 
 /**
- * @brief Convert a DLManagedTensor to an mdspan
+ * @brief Convert a DLManagedTensor to a mdspan
  * NOTE: This function only supports compact row-major and col-major layouts.
  *
  * @code {.cpp}
@@ -93,6 +93,19 @@ inline MdspanType from_dlpack(DLManagedTensor* managed_tensor)
   return detail::from_dlpack<MdspanType>(managed_tensor);
 }
 
+/**
+ * @brief Convert a mdspan to a DLManagedTensor
+ *
+ * Converts a mdspan to a DLManagedTensor object. This lets us pass non-owning
+ * views from C++ to C code without copying.  Note that returned DLManagedTensor
+ * is a non-owning view, and doesn't ensure that the underlying memory stays valid.
+ */
+template <typename MdspanType, typename = raft::is_mdspan_t<MdspanType>>
+void to_dlpack(MdspanType src, DLManagedTensor* dst)
+{
+  return detail::to_dlpack(src, dst);
+}
+
 /**
  * @}
  */
diff --git a/cpp/include/cuvs/neighbors/nn_descent.h b/cpp/include/cuvs/neighbors/nn_descent.h
@@ -171,11 +171,29 @@ cuvsError_t cuvsNNDescentBuild(cuvsResources_t res,
 /**
  * @brief Get the KNN graph from a built NN-Descent index
  *
+ * @param[in] res cuvsResources_t opaque C handle
  * @param[in] index cuvsNNDescentIndex_t Built NN-Descent index
- * @param[inout] graph Optional preallocated graph on host memory to store output
+ * @param[out] graph Preallocated graph on host memory to store output
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsNNDescentIndexGetGraph(cuvsResources_t res,
+                                       cuvsNNDescentIndex_t index,
+                                       DLManagedTensor* graph);
+
+/**
+ * @brief Get the distances from a build NN_Descent index
+ *
+ * This requires that the `return_distances` parameter was set when building the
+ * graph
+ *
+ * @param[in] res cuvsResources_t opaque C handle
+ * @param[in] index cuvsNNDescentIndex_t Built NN-Descent index
+ * @param[out] distances Preallocated memory to store the output distances tensor
  * @return cuvsError_t
  */
-cuvsError_t cuvsNNDescentIndexGetGraph(cuvsNNDescentIndex_t index, DLManagedTensor* graph);
+cuvsError_t cuvsNNDescentIndexGetDistances(cuvsResources_t res,
+                                           cuvsNNDescentIndex_t index,
+                                           DLManagedTensor* distances);
 #ifdef __cplusplus
 }
 #endif
diff --git a/cpp/src/neighbors/nn_descent_c.cpp b/cpp/src/neighbors/nn_descent_c.cpp
@@ -71,6 +71,57 @@ void* _build(cuvsResources_t res,
     RAFT_FAIL("dataset must be accessible on host or device memory");
   }
 }
+
+template <typename output_mdspan_type>
+void _get_graph(cuvsResources_t res, cuvsNNDescentIndex_t index, DLManagedTensor* graph)
+{
+  auto dtype = index->dtype;
+  if ((dtype.code == kDLUInt) && (dtype.bits == 32)) {
+    auto index_ptr = reinterpret_cast<cuvs::neighbors::nn_descent::index<uint32_t>*>(index->addr);
+    auto dst       = cuvs::core::from_dlpack<output_mdspan_type>(graph);
+    auto src       = index_ptr->graph();
+    auto res_ptr   = reinterpret_cast<raft::resources*>(res);
+
+    RAFT_EXPECTS(src.extent(0) == dst.extent(0), "Output graph has incorrect number of rows");
+    RAFT_EXPECTS(src.extent(1) == dst.extent(1), "Output graph has incorrect number of cols");
+
+    cudaMemcpyAsync(dst.data_handle(),
+                    src.data_handle(),
+                    dst.extent(0) * dst.extent(1) * sizeof(uint32_t),
+                    cudaMemcpyDefault,
+                    raft::resource::get_cuda_stream(*res_ptr));
+  } else {
+    RAFT_FAIL("Unsupported nn-descent index dtype: %d and bits: %d", dtype.code, dtype.bits);
+  }
+}
+
+template <typename output_mdspan_type>
+void _get_distances(cuvsResources_t res, cuvsNNDescentIndex_t index, DLManagedTensor* distances)
+{
+  auto dtype = index->dtype;
+  if ((dtype.code == kDLUInt) && (dtype.bits == 32)) {
+    auto index_ptr = reinterpret_cast<cuvs::neighbors::nn_descent::index<uint32_t>*>(index->addr);
+    auto src       = index_ptr->distances();
+    if (!src.has_value()) {
+      RAFT_FAIL("nn-descent index doesn't contain distances - set return_distances when building");
+    }
+
+    auto res_ptr = reinterpret_cast<raft::resources*>(res);
+    auto dst     = cuvs::core::from_dlpack<output_mdspan_type>(distances);
+
+    RAFT_EXPECTS(src->extent(0) == dst.extent(0), "Output distances has incorrect number of rows");
+    RAFT_EXPECTS(src->extent(1) == dst.extent(1), "Output distances has incorrect number of cols");
+
+    cudaMemcpyAsync(dst.data_handle(),
+                    src->data_handle(),
+                    dst.extent(0) * dst.extent(1) * sizeof(float),
+                    cudaMemcpyDefault,
+                    raft::resource::get_cuda_stream(*res_ptr));
+
+  } else {
+    RAFT_FAIL("Unsupported nn-descent index dtype: %d and bits: %d", dtype.code, dtype.bits);
+  }
+}
 }  // namespace
 
 extern "C" cuvsError_t cuvsNNDescentIndexCreate(cuvsNNDescentIndex_t* index)
@@ -146,22 +197,32 @@ extern "C" cuvsError_t cuvsNNDescentIndexParamsDestroy(cuvsNNDescentIndexParams_
   return cuvs::core::translate_exceptions([=] { delete params; });
 }
 
-extern "C" cuvsError_t cuvsNNDescentIndexGetGraph(cuvsNNDescentIndex_t index,
+extern "C" cuvsError_t cuvsNNDescentIndexGetGraph(cuvsResources_t res,
+                                                  cuvsNNDescentIndex_t index,
                                                   DLManagedTensor* graph)
 {
   return cuvs::core::translate_exceptions([=] {
-    auto dtype = index->dtype;
-    if ((dtype.code == kDLUInt) && (dtype.bits == 32)) {
-      auto index_ptr = reinterpret_cast<cuvs::neighbors::nn_descent::index<uint32_t>*>(index->addr);
+    if (cuvs::core::is_dlpack_device_compatible(graph->dl_tensor)) {
+      using output_mdspan_type = raft::device_matrix_view<uint32_t, int64_t, raft::row_major>;
+      _get_graph<output_mdspan_type>(res, index, graph);
+    } else {
       using output_mdspan_type = raft::host_matrix_view<uint32_t, int64_t, raft::row_major>;
-      auto dst                 = cuvs::core::from_dlpack<output_mdspan_type>(graph);
-      auto src                 = index_ptr->graph();
+      _get_graph<output_mdspan_type>(res, index, graph);
+    }
+  });
+}
 
-      RAFT_EXPECTS(src.extent(0) == dst.extent(0), "Output graph has incorrect number of rows");
-      RAFT_EXPECTS(src.extent(1) == dst.extent(1), "Output graph has incorrect number of cols");
-      std::copy(src.data_handle(), src.data_handle() + dst.size(), dst.data_handle());
+extern "C" cuvsError_t cuvsNNDescentIndexGetDistances(cuvsResources_t res,
+                                                      cuvsNNDescentIndex_t index,
+                                                      DLManagedTensor* distances)
+{
+  return cuvs::core::translate_exceptions([=] {
+    if (cuvs::core::is_dlpack_device_compatible(distances->dl_tensor)) {
+      using output_mdspan_type = raft::device_matrix_view<float, int64_t, raft::row_major>;
+      _get_distances<output_mdspan_type>(res, index, distances);
     } else {
-      RAFT_FAIL("Unsupported nn-descent index dtype: %d and bits: %d", dtype.code, dtype.bits);
+      using output_mdspan_type = raft::host_matrix_view<float, int64_t, raft::row_major>;
+      _get_distances<output_mdspan_type>(res, index, distances);
     }
   });
 }
diff --git a/dependencies.yaml b/dependencies.yaml
@@ -17,6 +17,7 @@ files:
       - depends_on_cupy
       - depends_on_librmm
       - depends_on_pylibraft
+      - depends_on_nccl
       - docs
       - rapids_build
       - run_py_cuvs
@@ -42,6 +43,7 @@ files:
       - depends_on_pylibraft
       - depends_on_libcuvs
       - depends_on_librmm
+      - depends_on_nccl
       - rapids_build
       - rapids_build_setuptools
   test_cpp:
@@ -94,6 +96,7 @@ files:
       - rust
       - depends_on_libcuvs
       - depends_on_libraft
+      - depends_on_nccl
   go:
     output: conda
     matrix:
@@ -107,6 +110,7 @@ files:
       - go
       - depends_on_libcuvs
       - depends_on_libraft
+      - depends_on_nccl
   py_build_libcuvs:
     output: pyproject
     pyproject_dir: python/libcuvs
@@ -123,6 +127,7 @@ files:
     includes:
       - depends_on_libraft
       - depends_on_librmm
+      - depends_on_nccl
       - rapids_build
   py_run_libcuvs:
     output: pyproject
@@ -133,6 +138,7 @@ files:
       - cuda_wheels
       - depends_on_libraft
       - depends_on_librmm
+      - depends_on_nccl
   py_build_cuvs:
     output: pyproject
     pyproject_dir: python/cuvs
@@ -226,7 +232,6 @@ dependencies:
         packages:
           - c-compiler
           - cxx-compiler
-          - nccl>=2.19
     specific:
       - output_types: conda
         matrices:
@@ -706,3 +711,18 @@ dependencies:
             packages:
               - pylibraft-cu11==25.6.*,>=0.0.0a0
           - {matrix: null, packages: [*pylibraft_unsuffixed]}
+  depends_on_nccl:
+    common:
+      - output_types: conda
+        packages:
+          - &nccl_unsuffixed nccl>=2.19
+    specific:
+      - output_types: [pyproject, requirements]
+        matrices:
+          - matrix:
+              cuda: "12.*"
+              cuda_suffixed: "true"
+            packages:
+              - nvidia-nccl-cu12>=2.19
+          - matrix:
+            packages:
diff --git a/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pxd b/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pxd
@@ -53,9 +53,14 @@ cdef extern from "cuvs/neighbors/nn_descent.h" nogil:
 
     cuvsError_t cuvsNNDescentIndexDestroy(cuvsNNDescentIndex_t index)
 
-    cuvsError_t cuvsNNDescentIndexGetGraph(cuvsNNDescentIndex_t index,
+    cuvsError_t cuvsNNDescentIndexGetGraph(cuvsResources_t res,
+                                           cuvsNNDescentIndex_t index,
                                            DLManagedTensor * output)
 
+    cuvsError_t cuvsNNDescentIndexGetDistances(cuvsResources_t res,
+                                               cuvsNNDescentIndex_t index,
+                                               DLManagedTensor * output)
+
     cuvsError_t cuvsNNDescentBuild(cuvsResources_t res,
                                    cuvsNNDescentIndexParams* params,
                                    DLManagedTensor* dataset,
diff --git a/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pyx b/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pyx
@@ -88,7 +88,8 @@ cdef class IndexParams:
                  intermediate_graph_degree=None,
                  max_iterations=None,
                  termination_threshold=None,
-                 n_clusters=None
+                 n_clusters=None,
+                 return_distances=None
                  ):
         if metric is not None:
             self.params.metric = <cuvsDistanceType>DISTANCE_TYPES[metric]
@@ -102,11 +103,8 @@ cdef class IndexParams:
             self.params.termination_threshold = termination_threshold
         if n_clusters is not None:
             self.params.n_clusters = n_clusters
-
-        # setting this parameter to true will cause an exception in the c++
-        # api (`Using return_distances set to true requires distance view to
-        # be allocated.`) - so instead force to be false here
-        self.params.return_distances = False
+        if return_distances is not None:
+            self.params.return_distances = return_distances
 
     @property
     def metric(self):
@@ -163,13 +161,39 @@ cdef class Index:
 
     @property
     def graph(self):
+        return self._get_graph()
+
+    @property
+    def distances(self):
+        return self._get_distances()
+
+    @auto_sync_resources
+    def _get_graph(self, resources=None):
         if not self.trained:
             raise ValueError("Index needs to be built before getting graph")
 
+        cdef cuvsResources_t res = <cuvsResources_t>resources.get_c_obj()
+
         output = np.empty((self.num_rows, self.graph_degree), dtype='uint32')
         ai = wrap_array(output)
         cdef cydlpack.DLManagedTensor* output_dlpack = cydlpack.dlpack_c(ai)
-        check_cuvs(cuvsNNDescentIndexGetGraph(self.index, output_dlpack))
+        check_cuvs(cuvsNNDescentIndexGetGraph(res, self.index, output_dlpack))
+        return output
+
+    @auto_sync_resources
+    def _get_distances(self, resources=None):
+        if not self.trained:
+            raise ValueError("Index needs to be built before getting"
+                             " distances")
+
+        cdef cuvsResources_t res = <cuvsResources_t>resources.get_c_obj()
+
+        output = np.empty((self.num_rows, self.graph_degree), dtype='float32')
+        ai = wrap_array(output)
+        cdef cydlpack.DLManagedTensor* output_dlpack = cydlpack.dlpack_c(ai)
+        check_cuvs(cuvsNNDescentIndexGetDistances(res,
+                                                  self.index,
+                                                  output_dlpack))
         return output
 
     def __repr__(self):
@@ -221,6 +245,16 @@ def build(IndexParams index_params, dataset, graph=None, resources=None):
 
     cdef cydlpack.DLManagedTensor* graph_dlpack = NULL
     if graph is not None:
+        if params.return_distances:
+            # When using a pre-existing graph - having return_distances set to
+            # true will cause an exception in the C++ api
+            # (`Using return_distances set to true requires distance view to
+            # be allocated.`). Raise a more informative error here instead of
+            # the C++ exception
+            raise ValueError("Can't use return_distances with an existing"
+                             " graph. Either set params.return_distances to"
+                             " False, or set graph to None")
+
         graph_ai = wrap_array(graph)
         graph_dlpack = cydlpack.dlpack_c(graph_ai)
 
diff --git a/python/cuvs/cuvs/tests/test_nn_descent.py b/python/cuvs/cuvs/tests/test_nn_descent.py
diff --git a/python/libcuvs/CMakeLists.txt b/python/libcuvs/CMakeLists.txt
