All-neighbors graph on host memory (#2313)

Closes #1903

This PR adds support for CPU kNN graph for batched all-neighbors. Passing host matrices for the indices and distances avoids having to fit the entire graph on GPU. Small perf gap and expected peak GPU memory reductions below:

<img width="2384" height="1487" alt="results" src="https://github.com/user-attachments/assets/884f9716-18f1-4a1f-9993-4d4604761c6a" />

Authors:
  - Jinsol Park (https://github.com/jinsolp)

Approvers:
  - Dante Gama Dessavre (https://github.com/dantegd)

URL: #2313

Author: jinsolp

c/include/cuvs/neighbors/all_neighbors.h

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -28,8 +28,11 @@ extern "C" {
  * provide the dataset on host.
  *
  * Notes:
- * - Outputs (indices, distances, core_distances) are expected to be on device memory.
- * - Host variant accepts host-resident dataset; device variant accepts device-resident dataset.
+ * - A host-resident dataset accepts either host- or device-resident outputs (indices, distances,
+ *   core_distances); a device-resident dataset requires device-resident outputs. All provided
+ *   outputs must share the same memory space.
+ * - With batched builds (`n_clusters > 1`), host-resident outputs avoid materializing the full
+ *   [num_rows x k] graph on the GPU at once.
  * - For batching, `overlap_factor < n_clusters` must hold.
  * - When `core_distances` is provided, mutual-reachability distances are produced (see alpha).
  */
@@ -94,16 +97,19 @@ CUVS_EXPORT cuvsError_t cuvsAllNeighborsIndexParamsDestroy(cuvsAllNeighborsIndex
  * 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] core_distances Optional 1D tensor [num_rows] on device (float32); can be NULL
+ * @param[out] indices        2D tensor [num_rows x k] (int64), host or device
+ * @param[out] distances      Optional 2D tensor [num_rows x k] (float32), host or device; can be
+ *                            NULL
+ * @param[out] core_distances Optional 1D tensor [num_rows] (float32), host or device; 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.
+ *
+ * Output memory space: a host dataset supports host- or device-resident outputs; a device dataset
+ * requires device-resident outputs. All provided outputs must share the same memory space.
  */
 CUVS_EXPORT cuvsError_t cuvsAllNeighborsBuild(cuvsResources_t res,
                                   cuvsAllNeighborsIndexParams_t params,

c/src/neighbors/all_neighbors.cpp

@@ -1,10 +1,11 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #include <cstdint>
 #include <dlpack/dlpack.h>
+#include <type_traits>
 
 #include <raft/core/error.hpp>
 #include <raft/core/mdspan_types.hpp>
@@ -80,36 +81,50 @@ 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(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");
 }
 
-static void ensure_optional_distance_dtype_and_device_compatibility(DLManagedTensor* distances)
+static void ensure_optional_distance_dtype(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");
 }
 
-static void ensure_optional_core_distance_dtype_and_device_compatibility(
-  DLManagedTensor* core_distances)
+static void ensure_optional_core_distance_dtype(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");
 }
 
-template <typename T>
+// Validate that the outputs (indices/distances/core_distances) all live in the same memory space
+static bool validate_output_memory_space(DLManagedTensor* indices,
+                                         DLManagedTensor* distances,
+                                         DLManagedTensor* core_distances)
+{
+  const bool host = cuvs::core::is_dlpack_host_compatible(indices->dl_tensor);
+  RAFT_EXPECTS(host || cuvs::core::is_dlpack_device_compatible(indices->dl_tensor),
+               "indices tensor must be host- or device-compatible");
+  auto same_space = [&](DLManagedTensor* t, const char* name) {
+    if (t == nullptr) { return; }
+    RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(t->dl_tensor) == host,
+                 "%s tensor must be in the same memory space (host or device) as indices",
+                 name);
+  };
+  same_space(distances, "distances");
+  same_space(core_distances, "core_distances");
+  return host;
+}
+
+// Build with a host-resident dataset. HostOutput selects whether the outputs live on host or
+// device.
+template <typename T, bool HostOutput>
 void _build_host(cuvsResources_t res,
                  cuvsAllNeighborsIndexParams_t params,
                  DLManagedTensor* dataset_tensor,
@@ -124,24 +139,23 @@ 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);
-
-  // Check dependencies between parameters
-  if (core_distances_tensor != nullptr && distances_tensor == nullptr) {
-    RAFT_FAIL("distances tensor must be provided when core_distances tensor is provided");
-  }
-
   int64_t n_rows = dlt.shape[0];
   int64_t n_cols = dlt.shape[1];
 
   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>;
+  using indices_mdspan_t =
+    std::conditional_t<HostOutput,
+                       raft::host_matrix_view<int64_t, int64_t, raft::row_major>,
+                       raft::device_matrix_view<int64_t, int64_t, raft::row_major>>;
+  using distances_mdspan_t =
+    std::conditional_t<HostOutput,
+                       raft::host_matrix_view<float, int64_t, raft::row_major>,
+                       raft::device_matrix_view<float, int64_t, raft::row_major>>;
+  using core_mdspan_t = std::conditional_t<HostOutput,
+                                           raft::host_vector_view<float, int64_t>,
+                                           raft::device_vector_view<float, int64_t>>;
 
   auto dataset = cuvs::core::from_dlpack<dataset_mdspan_t>(dataset_tensor);
   auto indices = cuvs::core::from_dlpack<indices_mdspan_t>(indices_tensor);
@@ -160,6 +174,7 @@ void _build_host(cuvsResources_t res,
     cpp_res, cpp_params, dataset, indices, distances, core_distances, alpha);
 }
 
+// Build with a device-resident dataset. Outputs are always device-resident.
 template <typename T>
 void _build_device(cuvsResources_t device_res,
                    cuvsAllNeighborsIndexParams_t params,
@@ -175,15 +190,6 @@ 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);
-
-  // Check dependencies between parameters
-  if (core_distances_tensor != nullptr && distances_tensor == nullptr) {
-    RAFT_FAIL("distances tensor must be provided when core_distances tensor is provided");
-  }
-
   int64_t n_rows = dlt.shape[0];
   int64_t n_cols = dlt.shape[1];
 
@@ -250,17 +256,33 @@ extern "C" cuvsError_t cuvsAllNeighborsBuild(cuvsResources_t res,
   return cuvs::core::translate_exceptions([=] {
     auto dataset = dataset_tensor->dl_tensor;
 
+    ensure_indices_dtype(indices_tensor);
+    ensure_optional_distance_dtype(distances_tensor);
+    ensure_optional_core_distance_dtype(core_distances_tensor);
+    if (core_distances_tensor != nullptr && distances_tensor == nullptr) {
+      RAFT_FAIL("distances tensor must be provided when core_distances tensor is provided");
+    }
+
+     const bool host_output =
+      validate_output_memory_space(indices_tensor, distances_tensor, core_distances_tensor);
+
     if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) {
       // Check if dataset is host-compatible or device-compatible
       if (cuvs::core::is_dlpack_host_compatible(dataset)) {
-        _build_host<float>(res,
-                           params,
-                           dataset_tensor,
-                           indices_tensor,
-                           distances_tensor,
-                           core_distances_tensor,
-                           alpha);
+        // Host dataset supports both host- and device-resident outputs.
+        if (host_output) {
+          _build_host<float, true>(
+            res, params, dataset_tensor, indices_tensor, distances_tensor, core_distances_tensor,
+            alpha);
+        } else {
+          _build_host<float, false>(
+            res, params, dataset_tensor, indices_tensor, distances_tensor, core_distances_tensor,
+            alpha);
+        }
       } else if (cuvs::core::is_dlpack_device_compatible(dataset)) {
+        RAFT_EXPECTS(!host_output,
+                    "A device-resident dataset requires device-resident outputs; put the dataset "
+                    "on host to produce host-resident outputs.");
         _build_device<float>(res,
                              params,
                              dataset_tensor,

cpp/include/cuvs/neighbors/all_neighbors.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -119,8 +119,8 @@ 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] indices nearest neighbor indices of shape [n_row x k] on device memory
+ * @param[out] distances nearest neighbor distances [n_row x k] on device memory
  * @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.
@@ -135,6 +135,40 @@ 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)
+ *
+ * Usage example:
+ * @code{.cpp}
+ *  using namespace cuvs::neighbors;
+ *  all_neighbors::all_neighbors_params params;
+ *  params.n_clusters = 4;
+ *  auto indices   = raft::make_host_matrix<int64_t, int64_t>(n_row, k);
+ *  auto distances = raft::make_host_matrix<float, int64_t>(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] on host memory
+ * @param[out] distances nearest neighbor distances [n_row x k] on host memory
+ * @param[out] core_distances array for core distances of size [n_row] on host memory. Requires
+ * distances matrix to compute core_distances. If core_distances is given, the resulting indices and
+ * distances will be 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) params.n_clusters should be 1 for data on device. To use a larger params.n_clusters for
@@ -155,8 +189,8 @@ 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] indices nearest neighbor indices of shape [n_row x k] on device memory
+ * @param[out] distances nearest neighbor distances [n_row x k] on device memory
  * @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.

cpp/src/neighbors/all_neighbors/all_neighbors.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -30,6 +30,18 @@ namespace cuvs::neighbors::all_neighbors {
   {                                                                                      \
     return all_neighbors::detail::build<T, IdxT>(                                        \
       handle, params, dataset, indices, distances, core_distances, alpha);               \
+  }                                                                                      \
+                                                                                         \
+  void build(const raft::resources& handle,                                              \
+             const all_neighbors_params& params,                                         \
+             raft::host_matrix_view<const T, IdxT, row_major> dataset,                   \
+             raft::host_matrix_view<IdxT, IdxT, row_major> indices,                      \
+             std::optional<raft::host_matrix_view<T, IdxT, row_major>> distances,        \
+             std::optional<raft::host_vector_view<T, IdxT, row_major>> core_distances,   \
+             T alpha)                                                                    \
+  {                                                                                      \
+    return all_neighbors::detail::build<T, IdxT>(                                        \
+      handle, params, dataset, indices, distances, core_distances, alpha);               \
   }
 
 CUVS_INST_ALL_NEIGHBORS(float, int64_t);