API for CPU side output indices and distances of all_neighbors

c/include/cuvs/neighbors/all_neighbors.h

@@ -26,7 +26,8 @@ extern "C" {
  * provide the dataset on host.
  *
  * Notes:
- * - Outputs (indices, distances, core_distances) are expected to be on device memory.
+ * - Outputs (indices, distances) can be on host memory (numpy arrays)
+ *   or device memory (CUDA arrays). core_distances can only be on device memory.
  * - Host variant accepts host-resident dataset; device variant accepts device-resident dataset.
  * - For batching, `overlap_factor < n_clusters` must hold.
  * - When `core_distances` is provided, mutual-reachability distances are produced (see alpha).
@@ -92,16 +93,16 @@ cuvsError_t cuvsAllNeighborsIndexParamsDestroy(cuvsAllNeighborsIndexParams_t ind
  * resources
  * @param[in] params          Build parameters (see cuvsAllNeighborsIndexParams)
  * @param[in] dataset         2D tensor [num_rows x dim] on host or device (auto-detected)
- * @param[out] indices        2D tensor [num_rows x k] on device (int64)
- * @param[out] distances      Optional 2D tensor [num_rows x k] on device (float32); can be NULL
+ * @param[out] indices        2D tensor [num_rows x k] on host or device (int64)
+ * @param[out] distances      Optional 2D tensor [num_rows x k] on host or device (float32); can be NULL
  * @param[out] core_distances Optional 1D tensor [num_rows] on device (float32); can be NULL
  * @param[in] alpha           Mutual-reachability scaling; used only when core_distances is provided
  *
  * The function automatically detects whether the dataset is host-resident or device-resident
  * and calls the appropriate implementation. For host datasets, it partitions data into
  * `n_clusters` clusters and assigns each row to `overlap_factor` nearest clusters. For device
  * datasets, `n_clusters` must be 1 (no batching); `overlap_factor` is ignored.
- * Outputs always reside in device memory.
+ * Outputs can be on host memory (numpy arrays) or device memory (CUDA arrays).
  */
 cuvsError_t cuvsAllNeighborsBuild(cuvsResources_t res,
                                   cuvsAllNeighborsIndexParams_t params,

c/src/neighbors/all_neighbors.cpp

@@ -80,33 +80,36 @@ static cuvs::neighbors::all_neighbors::all_neighbors_params convert_params(
   return out;
 }
 
-static void ensure_indices_dtype_and_device_compatibility(DLManagedTensor* indices)
+static void ensure_indices_dtype_compatibility(DLManagedTensor* indices)
 {
   auto dtype = indices->dl_tensor.dtype;
   RAFT_EXPECTS(dtype.code == kDLInt && dtype.bits == 64, "indices must be int64 output tensor");
-  RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(indices->dl_tensor),
-               "indices tensor must be device-compatible");
+  RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(indices->dl_tensor) ||
+               cuvs::core::is_dlpack_host_compatible(indices->dl_tensor),
+               "indices tensor must be either device-compatible or host-compatible");
 }
 
-static void ensure_optional_distance_dtype_and_device_compatibility(DLManagedTensor* distances)
+static void ensure_optional_distance_dtype_compatibility(DLManagedTensor* distances)
 {
   if (distances == nullptr) { return; }
   auto dtype = distances->dl_tensor.dtype;
   RAFT_EXPECTS(dtype.code == kDLFloat && dtype.bits == 32,
                "distances must be float32 output tensor");
-  RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(distances->dl_tensor),
-               "distances tensor must be device-compatible");
+  RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(distances->dl_tensor) ||
+               cuvs::core::is_dlpack_host_compatible(distances->dl_tensor),
+               "distances tensor must be either device-compatible or host-compatible");
 }
 
-static void ensure_optional_core_distance_dtype_and_device_compatibility(
-  DLManagedTensor* core_distances)
+static void ensure_optional_core_distance_dtype_compatibility(DLManagedTensor* core_distances)
 {
   if (core_distances == nullptr) { return; }
   auto dtype = core_distances->dl_tensor.dtype;
   RAFT_EXPECTS(dtype.code == kDLFloat && dtype.bits == 32,
                "core_distances must be float32 output tensor");
-  RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(core_distances->dl_tensor),
-               "core_distances tensor must be device-compatible");
+  RAFT_EXPECTS(
+    cuvs::core::is_dlpack_device_compatible(core_distances->dl_tensor) ||
+      cuvs::core::is_dlpack_host_compatible(core_distances->dl_tensor),
+    "core_distances tensor must be either device-compatible or host-compatible");
 }
 
 template <typename T>
@@ -124,9 +127,9 @@ void _build_host(cuvsResources_t res,
   RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(dlt),
                "Host build expects host-compatible dataset tensor");
 
-  ensure_indices_dtype_and_device_compatibility(indices_tensor);
-  ensure_optional_distance_dtype_and_device_compatibility(distances_tensor);
-  ensure_optional_core_distance_dtype_and_device_compatibility(core_distances_tensor);
+  ensure_indices_dtype_compatibility(indices_tensor);
+  ensure_optional_distance_dtype_compatibility(distances_tensor);
+  ensure_optional_core_distance_dtype_compatibility(core_distances_tensor);
 
   // Check dependencies between parameters
   if (core_distances_tensor != nullptr && distances_tensor == nullptr) {
@@ -138,26 +141,63 @@ void _build_host(cuvsResources_t res,
 
   auto cpp_params = convert_params(params, n_rows, n_cols);
 
-  using dataset_mdspan_t   = raft::host_matrix_view<const T, int64_t, raft::row_major>;
-  using indices_mdspan_t   = raft::device_matrix_view<int64_t, int64_t, raft::row_major>;
-  using distances_mdspan_t = raft::device_matrix_view<float, int64_t, raft::row_major>;
-  using core_mdspan_t      = raft::device_vector_view<float, int64_t>;
+  using dataset_mdspan_t = raft::host_matrix_view<const T, int64_t, raft::row_major>;
 
-  auto dataset = cuvs::core::from_dlpack<dataset_mdspan_t>(dataset_tensor);
-  auto indices = cuvs::core::from_dlpack<indices_mdspan_t>(indices_tensor);
+  bool indices_is_host   = cuvs::core::is_dlpack_host_compatible(indices_tensor->dl_tensor);
+  bool distances_is_host = distances_tensor ? cuvs::core::is_dlpack_host_compatible(distances_tensor->dl_tensor) : indices_is_host;
 
-  std::optional<distances_mdspan_t> distances = std::nullopt;
-  if (distances_tensor) {
-    distances = cuvs::core::from_dlpack<distances_mdspan_t>(distances_tensor);
+  if (distances_tensor && distances_is_host != indices_is_host) {
+    RAFT_FAIL("distances and indices must be on the same memory location (both host or both device)");
   }
 
-  std::optional<core_mdspan_t> core_distances = std::nullopt;
   if (core_distances_tensor) {
-    core_distances = cuvs::core::from_dlpack<core_mdspan_t>(core_distances_tensor);
+    bool core_distances_is_host =
+      cuvs::core::is_dlpack_host_compatible(core_distances_tensor->dl_tensor);
+    RAFT_EXPECTS(core_distances_is_host == indices_is_host,
+                 "core_distances must be on the same memory location as indices and distances");
   }
 
-  cuvs::neighbors::all_neighbors::build(
-    cpp_res, cpp_params, dataset, indices, distances, core_distances, alpha);
+  auto dataset = cuvs::core::from_dlpack<dataset_mdspan_t>(dataset_tensor);
+
+  if (indices_is_host) {
+    using indices_mdspan_t   = raft::host_matrix_view<int64_t, int64_t, raft::row_major>;
+    using distances_mdspan_t = raft::host_matrix_view<float, int64_t, raft::row_major>;
+    using core_mdspan_t      = raft::host_vector_view<float, int64_t>;
+
+    auto indices = cuvs::core::from_dlpack<indices_mdspan_t>(indices_tensor);
+
+    std::optional<distances_mdspan_t> distances = std::nullopt;
+    if (distances_tensor) {
+      distances = cuvs::core::from_dlpack<distances_mdspan_t>(distances_tensor);
+    }
+
+    std::optional<core_mdspan_t> core_distances = std::nullopt;
+    if (core_distances_tensor) {
+      core_distances = cuvs::core::from_dlpack<core_mdspan_t>(core_distances_tensor);
+    }
+
+    cuvs::neighbors::all_neighbors::build(
+      cpp_res, cpp_params, dataset, indices, distances, core_distances, alpha);
+  } else {
+    using indices_mdspan_t   = raft::device_matrix_view<int64_t, int64_t, raft::row_major>;
+    using distances_mdspan_t = raft::device_matrix_view<float, int64_t, raft::row_major>;
+    using core_mdspan_t      = raft::device_vector_view<float, int64_t>;
+
+    auto indices = cuvs::core::from_dlpack<indices_mdspan_t>(indices_tensor);
+
+    std::optional<distances_mdspan_t> distances = std::nullopt;
+    if (distances_tensor) {
+      distances = cuvs::core::from_dlpack<distances_mdspan_t>(distances_tensor);
+    }
+
+    std::optional<core_mdspan_t> core_distances = std::nullopt;
+    if (core_distances_tensor) {
+      core_distances = cuvs::core::from_dlpack<core_mdspan_t>(core_distances_tensor);
+    }
+
+    cuvs::neighbors::all_neighbors::build(
+      cpp_res, cpp_params, dataset, indices, distances, core_distances, alpha);
+  }
 }
 
 template <typename T>
@@ -175,9 +215,9 @@ void _build_device(cuvsResources_t device_res,
   RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(dlt),
                "Device build expects device-compatible dataset tensor");
 
-  ensure_indices_dtype_and_device_compatibility(indices_tensor);
-  ensure_optional_distance_dtype_and_device_compatibility(distances_tensor);
-  ensure_optional_core_distance_dtype_and_device_compatibility(core_distances_tensor);
+  ensure_indices_dtype_compatibility(indices_tensor);
+  ensure_optional_distance_dtype_compatibility(distances_tensor);
+  ensure_optional_core_distance_dtype_compatibility(core_distances_tensor);
 
   // Check dependencies between parameters
   if (core_distances_tensor != nullptr && distances_tensor == nullptr) {
@@ -189,26 +229,62 @@ void _build_device(cuvsResources_t device_res,
 
   auto cpp_params = convert_params(params, n_rows, n_cols);
 
-  using dataset_mdspan_t   = raft::device_matrix_view<const T, int64_t, raft::row_major>;
-  using indices_mdspan_t   = raft::device_matrix_view<int64_t, int64_t, raft::row_major>;
-  using distances_mdspan_t = raft::device_matrix_view<float, int64_t, raft::row_major>;
-  using core_mdspan_t      = raft::device_vector_view<float, int64_t>;
+  using dataset_mdspan_t = raft::device_matrix_view<const T, int64_t, raft::row_major>;
+  auto dataset           = cuvs::core::from_dlpack<dataset_mdspan_t>(dataset_tensor);
 
-  auto dataset = cuvs::core::from_dlpack<dataset_mdspan_t>(dataset_tensor);
-  auto indices = cuvs::core::from_dlpack<indices_mdspan_t>(indices_tensor);
+  bool indices_is_host   = cuvs::core::is_dlpack_host_compatible(indices_tensor->dl_tensor);
+  bool distances_is_host = distances_tensor ? cuvs::core::is_dlpack_host_compatible(distances_tensor->dl_tensor) : indices_is_host;
 
-  std::optional<distances_mdspan_t> distances = std::nullopt;
-  if (distances_tensor) {
-    distances = cuvs::core::from_dlpack<distances_mdspan_t>(distances_tensor);
+  if (distances_tensor && distances_is_host != indices_is_host) {
+    RAFT_FAIL("distances and indices must be on the same memory location (both host or both device)");
   }
 
-  std::optional<core_mdspan_t> core_distances = std::nullopt;
   if (core_distances_tensor) {
-    core_distances = cuvs::core::from_dlpack<core_mdspan_t>(core_distances_tensor);
+    bool core_distances_is_host =
+      cuvs::core::is_dlpack_host_compatible(core_distances_tensor->dl_tensor);
+    RAFT_EXPECTS(core_distances_is_host == indices_is_host,
+                 "core_distances must be on the same memory location as indices and distances");
   }
 
-  cuvs::neighbors::all_neighbors::build(
-    cpp_res, cpp_params, dataset, indices, distances, core_distances, alpha);
+  if (indices_is_host) {
+    using indices_mdspan_t   = raft::host_matrix_view<int64_t, int64_t, raft::row_major>;
+    using distances_mdspan_t = raft::host_matrix_view<float, int64_t, raft::row_major>;
+    using core_mdspan_t      = raft::host_vector_view<float, int64_t>;
+
+    auto indices = cuvs::core::from_dlpack<indices_mdspan_t>(indices_tensor);
+
+    std::optional<distances_mdspan_t> distances = std::nullopt;
+    if (distances_tensor) {
+      distances = cuvs::core::from_dlpack<distances_mdspan_t>(distances_tensor);
+    }
+
+    std::optional<core_mdspan_t> core_distances = std::nullopt;
+    if (core_distances_tensor) {
+      core_distances = cuvs::core::from_dlpack<core_mdspan_t>(core_distances_tensor);
+    }
+
+    cuvs::neighbors::all_neighbors::build(
+      cpp_res, cpp_params, dataset, indices, distances, core_distances, alpha);
+  } else {
+    using indices_mdspan_t   = raft::device_matrix_view<int64_t, int64_t, raft::row_major>;
+    using distances_mdspan_t = raft::device_matrix_view<float, int64_t, raft::row_major>;
+    using core_mdspan_t      = raft::device_vector_view<float, int64_t>;
+
+    auto indices = cuvs::core::from_dlpack<indices_mdspan_t>(indices_tensor);
+
+    std::optional<distances_mdspan_t> distances = std::nullopt;
+    if (distances_tensor) {
+      distances = cuvs::core::from_dlpack<distances_mdspan_t>(distances_tensor);
+    }
+
+    std::optional<core_mdspan_t> core_distances = std::nullopt;
+    if (core_distances_tensor) {
+      core_distances = cuvs::core::from_dlpack<core_mdspan_t>(core_distances_tensor);
+    }
+
+    cuvs::neighbors::all_neighbors::build(
+      cpp_res, cpp_params, dataset, indices, distances, core_distances, alpha);
+  }
 }
 
 }  // namespace

cpp/include/cuvs/neighbors/all_neighbors.hpp

@@ -116,11 +116,11 @@ struct all_neighbors_params {
  *               to build all-neighbors knn graph
  * @param[in] dataset raft::host_matrix_view input dataset expected to be located
  *                in host memory
- * @param[out] indices nearest neighbor indices of shape [n_row x k]
- * @param[out] distances nearest neighbor distances [n_row x k]
- * @param[out] core_distances array for core distances of size [n_row]. Requires distances matrix to
- * compute core_distances. If core_distances is given, the resulting indices and distances will be
- * mutual reachability space.
+ * @param[out] indices nearest neighbor indices of shape [n_row x k] in device memory
+ * @param[out] distances nearest neighbor distances [n_row x k] in device memory
+ * @param[out] core_distances core distances of size [n_row] in device memory. Requires distances
+ * matrix to compute core_distances. If core_distances is given, the resulting indices and distances
+ * will be in mutual reachability space.
  * @param[in] alpha distance scaling parameter as used in robust single linkage.
  */
 void build(
@@ -152,11 +152,11 @@ void build(
  *               to build all-neighbors knn graph
  * @param[in] dataset raft::device_matrix_view input dataset expected to be located
  *                in device memory
- * @param[out] indices nearest neighbor indices of shape [n_row x k]
- * @param[out] distances nearest neighbor distances [n_row x k]
- * @param[out] core_distances array for core distances of size [n_row]. Requires distances matrix to
- * compute core_distances. If core_distances is given, the resulting indices and distances will be
- * mutual reachability space.
+ * @param[out] indices nearest neighbor indices of shape [n_row x k] in device memory
+ * @param[out] distances nearest neighbor distances [n_row x k] in device memory
+ * @param[out] core_distances core distances of size [n_row] in device memory. Requires distances
+ * matrix to compute core_distances. If core_distances is given, the resulting indices and distances
+ * will be in mutual reachability space.
  * @param[in] alpha distance scaling parameter as used in robust single linkage.
  */
 void build(
@@ -168,5 +168,76 @@ void build(
   std::optional<raft::device_vector_view<float, int64_t, row_major>> core_distances = std::nullopt,
   float alpha                                                                       = 1.0);
 
+/**
+ * @brief Builds an approximate all-neighbors knn graph (find nearest neighbors for all the training
+ * vectors) with host memory output buffers.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *  using namespace cuvs::neighbors;
+ *  // use default index parameters
+ *  all_neighbors::all_neighbors_params params;
+ *  auto indices = raft::make_host_matrix<int64_t, int64_t>(handle, n_row, k);
+ *  auto distances = raft::make_host_matrix<float, int64_t>(handle, n_row, k);
+ *  all_neighbors::build(res, params, dataset, indices.view(), distances.view());
+ * @endcode
+ *
+ * @param[in] handle raft::resources is an object managing resources
+ * @param[in] params an instance of all_neighbors::all_neighbors_params that are parameters
+ *               to build all-neighbors knn graph
+ * @param[in] dataset raft::host_matrix_view input dataset expected to be located
+ *                in host memory
+ * @param[out] indices nearest neighbor indices of shape [n_row x k] in host memory
+ * @param[out] distances nearest neighbor distances [n_row x k] in host memory
+ * @param[out] core_distances core distances of size [n_row] in host memory. Requires distances
+ * matrix to compute core_distances. If core_distances is given, the resulting indices and distances
+ * will be in mutual reachability space.
+ * @param[in] alpha distance scaling parameter as used in robust single linkage.
+ */
+void build(
+  const raft::resources& handle,
+  const all_neighbors_params& params,
+  raft::host_matrix_view<const float, int64_t, row_major> dataset,
+  raft::host_matrix_view<int64_t, int64_t, row_major> indices,
+  std::optional<raft::host_matrix_view<float, int64_t, row_major>> distances      = std::nullopt,
+  std::optional<raft::host_vector_view<float, int64_t, row_major>> core_distances = std::nullopt,
+  float alpha                                                                     = 1.0);
+
+/**
+ * @brief Builds an approximate all-neighbors knn graph (find nearest neighbors for all the training
+ * vectors) with host memory output buffers. params.n_clusters should be 1 for data on device. To
+ * use a larger params.n_clusters for efficient device memory usage, put data on host RAM.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *  using namespace cuvs::neighbors;
+ *  // use default index parameters
+ *  all_neighbors::all_neighbors_params params;
+ *  auto indices = raft::make_host_matrix<int64_t, int64_t>(handle, n_row, k);
+ *  auto distances = raft::make_host_matrix<float, int64_t>(handle, n_row, k);
+ *  all_neighbors::build(res, params, dataset, indices.view(), distances.view());
+ * @endcode
+ *
+ * @param[in] handle raft::resources is an object managing resources
+ * @param[in] params an instance of all_neighbors::all_neighbors_params that are parameters
+ *               to build all-neighbors knn graph
+ * @param[in] dataset raft::device_matrix_view input dataset expected to be located
+ *                in device memory
+ * @param[out] indices nearest neighbor indices of shape [n_row x k] in host memory
+ * @param[out] distances nearest neighbor distances [n_row x k] in host memory
+ * @param[out] core_distances core distances of size [n_row] in host memory. Requires distances
+ * matrix to compute core_distances. If core_distances is given, the resulting indices and distances
+ * will be in mutual reachability space.
+ * @param[in] alpha distance scaling parameter as used in robust single linkage.
+ */
+void build(
+  const raft::resources& handle,
+  const all_neighbors_params& params,
+  raft::device_matrix_view<const float, int64_t, row_major> dataset,
+  raft::host_matrix_view<int64_t, int64_t, row_major> indices,
+  std::optional<raft::host_matrix_view<float, int64_t, row_major>> distances      = std::nullopt,
+  std::optional<raft::host_vector_view<float, int64_t, row_major>> core_distances = std::nullopt,
+  float alpha                                                                     = 1.0);
+
 /** @} */
 }  // namespace cuvs::neighbors::all_neighbors