[WIP] CAGRA ACE improvement: Global Rev-edge and exact initial kNN graph

c/include/cuvs/neighbors/cagra.h

@@ -149,8 +149,8 @@ struct cuvsAceParams {
    *
    * Small values might improve recall but potentially degrade performance and
    * increase memory usage. Partitions should not be too small to prevent issues
-   * in KNN graph construction. The partition size is on average 2 * (n_rows /
-   * npartitions) * dim * sizeof(T). 2 is because of the core and augmented
+   * in KNN graph construction. The partition size is on average `(1 + nsubpartitions) *
+   * (n_rows / npartitions) * dim * sizeof(T)`, accounting for core and augmented
    * vectors. Please account for imbalance in the partition sizes (up to 3x in
    * our tests).
    *
@@ -199,6 +199,17 @@ struct cuvsAceParams {
    * Useful for testing or when running alongside other memory-intensive processes.
    */
   double max_gpu_memory_gb;
+  /**
+   * Number of additional nearest partitions assigned to each vector. Default: 1.
+   * Must be positive and smaller than npartitions (unless automatically selected).
+   * Increasing this improves overlap at the cost of build time, memory, and disk space.
+   */
+  size_t nsubpartitions;
+  /**
+   * Build pruned partial graphs, then add reverse edges globally. Default: false.
+   * In disk mode, incoming edges are buffered to temporary files and merged on the CPU.
+   */
+  bool add_global_reverse_edges;
 };
 
 typedef struct cuvsAceParams* cuvsAceParams_t;

c/src/neighbors/cagra.cpp

@@ -817,6 +817,8 @@ static void _set_graph_build_params(
       if (params.graph_build_params) {
         auto ace_params_c             = static_cast<cuvsAceParams*>(params.graph_build_params);
         ace_p.npartitions         = ace_params_c->npartitions;
+        ace_p.nsubpartitions = ace_params_c->nsubpartitions;
+        ace_p.add_global_reverse_edges = ace_params_c->add_global_reverse_edges;
         ace_p.ef_construction     = ace_params_c->ef_construction;
         ace_p.build_dir           = std::string(ace_params_c->build_dir);
         ace_p.use_disk            = ace_params_c->use_disk;
@@ -1455,6 +1457,8 @@ static void _populate_cagra_index_params_from_cpp(cuvsCagraIndexParams_t c_param
         cpp_params.graph_build_params);
     cuvsAceParams* c_ace_params = new cuvsAceParams;
     c_ace_params->npartitions = ace_params.npartitions;
+    c_ace_params->nsubpartitions = ace_params.nsubpartitions;
+    c_ace_params->add_global_reverse_edges = ace_params.add_global_reverse_edges;
     c_ace_params->ef_construction = ace_params.ef_construction;
     c_ace_params->build_dir = ace_params.build_dir.empty() ? nullptr : strdup(ace_params.build_dir.c_str());
     c_ace_params->use_disk = ace_params.use_disk;
@@ -2342,7 +2346,9 @@ extern "C" cuvsError_t cuvsAceParamsCreate(cuvsAceParams_t* params)
                                 .build_dir           = build_dir,
                                 .use_disk            = ps.use_disk,
                                 .max_host_memory_gb  = ps.max_host_memory_gb,
-                                .max_gpu_memory_gb   = ps.max_gpu_memory_gb};
+                                .max_gpu_memory_gb   = ps.max_gpu_memory_gb,
+                                .nsubpartitions = ps.nsubpartitions,
+                                .add_global_reverse_edges = ps.add_global_reverse_edges};
   });
 }
 

c/tests/neighbors/ann_cagra_c.cu

@@ -1186,13 +1186,19 @@ TEST(CagraC, BuildSearchACEMemory)
 
   cuvsAceParams_t ace_params;
   cuvsAceParamsCreate(&ace_params);
+  EXPECT_EQ(ace_params->nsubpartitions, 1);
+  EXPECT_FALSE(ace_params->add_global_reverse_edges);
+  ace_params->nsubpartitions = 0;
+  ace_params->add_global_reverse_edges = true;
   ace_params->npartitions = 2;
   ace_params->ef_construction = 120;
   ace_params->use_disk = false;
 
   build_params->graph_build_params = ace_params;
   cuvsDataset_t dataset_view;
   ASSERT_EQ(cuvsDatasetMakeStandardView(res, &dataset_tensor, &dataset_view), CUVS_SUCCESS);
+  EXPECT_EQ(cuvsCagraBuild(res, build_params, dataset_view, index), CUVS_ERROR);
+  ace_params->nsubpartitions = 1;
   ASSERT_EQ(cuvsCagraBuild(res, build_params, dataset_view, index), CUVS_SUCCESS);
 
   // Host build yields a host index. Attach a caller-provided device padded dataset

cpp/include/cuvs/neighbors/cagra.hpp

@@ -180,20 +180,25 @@ struct ace_params {
    *
    * Small values might improve recall but potentially degrade performance and
    * increase memory usage. Partitions should not be too small to prevent issues
-   * in KNN graph construction. The partition size is on average 2 * (n_rows / npartitions) * dim *
-   * sizeof(T). 2 is because of the core and augmented vectors. Please account for imbalance in the
-   * partition sizes (up to 3x in our tests).
+   * in KNN graph construction. The partition size is on average `(1 + nsubpartitions) * (n_rows /
+   * npartitions) * dim * sizeof(T)`, accounting for the core and augmented vectors. Please account
+   * for imbalance in the partition sizes (up to 3x in our tests).
    *
    * If the specified number of partitions results in partitions that exceed
    * available memory, the value will be automatically increased to fit memory
    * constraints and a warning will be issued.
    */
   size_t npartitions = 0;
+  /** Number of additional nearest partitions assigned to each vector. Must be positive and
+   * smaller than npartitions (unless npartitions is auto-selected). Increasing this improves
+   * overlap at the cost of build time, memory, and temporary disk space. Default: 1.
+   */
+  size_t nsubpartitions = 1;
   /**
    * The index quality for the ACE build.
    *
-   * Bigger values increase the index quality. At some point, increasing this will no longer improve
-   * the quality.
+   * Bigger values increase the index quality. At some point, increasing this will no longer
+   * improve the quality.
    */
   size_t ef_construction = 120;
   /**
@@ -205,6 +210,9 @@ struct ace_params {
    * already exist, but ACE's named artifacts must not already exist. Simultaneous
    * builds must use different directories. On failure, ACE removes only artifacts
    * it created and never deletes unrelated directory contents.
+   * With `add_global_reverse_edges`, temporary incoming-edge records also require up to
+   * `dataset_size * graph_degree * (2 * sizeof(IdxT) + sizeof(uint32_t))` bytes, plus any
+   * record alignment padding. These records are removed after the merge.
    */
   std::string build_dir = "/tmp/ace_build";
   /**
@@ -214,6 +222,20 @@ struct ace_params {
    */
   bool use_disk = false;
 
+  /**
+   * Build pruned core-plus-argument partial graphs, then add reverse edges globally.
+   *
+   * Each partition first builds an intermediate kNN graph and prunes it without a local
+   * reverse-edge merge. Then build the reverse-edge graph globally and merge it into the pruned
+   * graph. In disk mode (explicit or automatic), incoming edges are buffered
+   * to temporary files while building each pruned partition, then merged on the CPU one
+   * destination partition at a time, retaining at most `ceil(graph_degree / 2)` incoming candidates
+   * per node in memory.
+   * The merge preserves edge-rank priority; ties use reordered source ID in disk mode
+   * and are GPU-scheduling-dependent in memory mode. Disabled by default.
+   */
+  bool add_global_reverse_edges = false;
+
   /**
    * Maximum host memory to use for ACE build in GiB.
    *
@@ -330,8 +352,8 @@ struct index_params : cuvs::neighbors::index_params {
    *
    *  - `true` (default) means `build` attaches the input dataset as a **non-owning view** to the
    * index. The caller is responsible for keeping the underlying dataset storage alive for as long
-   * as the index is used. A device-backed index is ready to search immediately; a host-backed index
-   * retains the dataset for operations such as serialization but is not searchable.
+   * as the index is used. A device-backed index is ready to search immediately; a host-backed
+   * index retains the dataset for operations such as serialization but is not searchable.
    *  - `false` means `build` only builds the graph and the caller is expected to attach a dataset
    * separately via `cuvs::neighbors::cagra::update_dataset` before searching.
    *
@@ -417,10 +439,11 @@ struct index_params : cuvs::neighbors::index_params {
    *
    * * IMPORTANT NOTE *
    *
-   * The reference HNSW index and the corresponding from-CAGRA generated HNSW index will NOT produce
-   * exactly the same recalls and QPS for the same parameter `ef`. The graphs are different
+   * The reference HNSW index and the corresponding from-CAGRA generated HNSW index will NOT
+   * produce exactly the same recalls and QPS for the same parameter `ef`. The graphs are different
    * internally. Depending on the selected heuristics, the CAGRA-produced graph's QPS-Recall curve
-   * may be shifted along the curve right or left. See the heuristics descriptions for more details.
+   * may be shifted along the curve right or left. See the heuristics descriptions for more
+   * details.
    *
    * Usage example:
    * @code{.cpp}
@@ -453,8 +476,8 @@ struct extend_params {
   /** The additional dataset is divided into chunks and added to the graph. This is the knob to
    * adjust the tradeoff between the recall and operation throughput. Large chunk sizes can result
    * in high throughput, but use more working memory (O(max_chunk_size*degree^2)). This can also
-   * degrade recall because no edges are added between the nodes in the same chunk. Auto select when
-   * 0. */
+   * degrade recall because no edges are added between the nodes in the same chunk. Auto select
+   * when 0. */
   uint32_t max_chunk_size = 0;
 };
 /**
@@ -724,7 +747,8 @@ struct CUVS_EXPORT index : cuvs::neighbors::index {
   }
 
   /**
-   * Replace the source indices with a new source indices taking the ownership of the passed vector.
+   * @brief Replace the source indices with a new source indices taking the ownership of the passed
+   * vector.
    */
   void update_source_indices(raft::device_vector<index_type, int64_t>&& source_indices)
   {

cpp/src/neighbors/cagra.cuh

@@ -7,6 +7,7 @@
 
 #include "detail/ann_utils.cuh"
 #include "detail/cagra/add_nodes.cuh"
+#include "detail/cagra/cagra_ace.cuh"
 #include "detail/cagra/cagra_build.cuh"
 #include "detail/cagra/cagra_merge.cuh"
 #include "detail/cagra/cagra_search.cuh"
@@ -340,6 +341,10 @@ auto build(raft::resources const& res, const index_params& params, DatasetViewT
     }
     return idx;
   } else if constexpr (cuvs::neighbors::is_dense_row_major_device_dataset_view_v<DatasetViewT>) {
+    RAFT_EXPECTS(
+      !std::holds_alternative<graph_build_params::ace_params>(params.graph_build_params),
+      "cagra::build: ACE requires a dataset in host memory; device-resident datasets are not "
+      "supported. Pass a host dataset view when using ace_params.");
     auto idx = cuvs::neighbors::cagra::detail::build_from_device_matrix<T, IdxT, DatasetViewT>(
       res, params, dataset);
     if (params.attach_dataset_on_build) {