Add ACE CAGRA Graph Reordering Option

cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h

@@ -373,6 +373,12 @@ void parse_build_param(const nlohmann::json& conf,
         ace_params.ef_construction = ace_conf.at("ef_construction");
       }
       if (ace_conf.contains("use_disk")) { ace_params.use_disk = ace_conf.at("use_disk"); }
+      if (ace_conf.contains("max_host_memory_gb")) {
+        ace_params.max_host_memory_gb = ace_conf.at("max_host_memory_gb");
+      }
+      if (ace_conf.contains("max_gpu_memory_gb")) {
+        ace_params.max_gpu_memory_gb = ace_conf.at("max_gpu_memory_gb");
+      }
       cagra_params.graph_build_params = ace_params;
     }
     ::parse_build_param<T, IdxT>(conf, cagra_params);

cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -20,6 +20,9 @@ auto parse_build_param(const nlohmann::json& conf) ->
   typename cuvs::bench::cuvs_cagra_hnswlib<T, IdxT>::build_param param;
   auto& hnsw_params  = param.hnsw_index_params;
   auto& cagra_params = param.cagra_build_params;
+  if (conf.contains("use_original_id_graph")) {
+    param.use_original_id_graph = conf.at("use_original_id_graph");
+  }
   if (conf.contains("hierarchy")) {
     if (conf.at("hierarchy") == "none") {
       hnsw_params.hierarchy = cuvs::neighbors::hnsw::HnswHierarchy::NONE;
@@ -57,14 +60,21 @@ auto parse_build_param(const nlohmann::json& conf) ->
       ps.metric = dist_type;
       // Parse ACE parameters if provided
       if (conf.contains("npartitions") || conf.contains("build_dir") ||
-          conf.contains("ef_construction") || conf.contains("use_disk")) {
+          conf.contains("ef_construction") || conf.contains("use_disk") ||
+          conf.contains("max_host_memory_gb") || conf.contains("max_gpu_memory_gb")) {
         auto ace_params = cuvs::neighbors::cagra::graph_build_params::ace_params();
         if (conf.contains("npartitions")) { ace_params.npartitions = conf.at("npartitions"); }
         if (conf.contains("build_dir")) { ace_params.build_dir = conf.at("build_dir"); }
         if (conf.contains("ef_construction")) {
           ace_params.ef_construction = conf.at("ef_construction");
         }
         if (conf.contains("use_disk")) { ace_params.use_disk = conf.at("use_disk"); }
+        if (conf.contains("max_host_memory_gb")) {
+          ace_params.max_host_memory_gb = conf.at("max_host_memory_gb");
+        }
+        if (conf.contains("max_gpu_memory_gb")) {
+          ace_params.max_gpu_memory_gb = conf.at("max_gpu_memory_gb");
+        }
         ps.graph_build_params = ace_params;
       }
       // NB: above, we only provide the defaults. Below we parse the explicit parameters as usual.

cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -22,6 +22,7 @@ class cuvs_cagra_hnswlib : public algo<T>, public algo_gpu {
     using cagra_wrapper_params = typename cuvs_cagra<T, IdxT>::build_param;
     cagra_wrapper_params cagra_build_params;
     cuvs::neighbors::hnsw::index_params hnsw_index_params;
+    bool use_original_id_graph = false;
   };
 
   struct search_param : public search_param_base {
@@ -102,9 +103,20 @@ void cuvs_cagra_hnswlib<T, IdxT>::build(const T* dataset, size_t nrow)
   cagra_wrapper.build(dataset, nrow);
   auto& cagra_index = *cagra_wrapper.get_index();
 
-  // pass the dataset directly to HNSW if it's on the host
+  const bool is_ace_disk_build = cagra_index.dataset_fd().has_value() &&
+                                 cagra_index.graph_fd().has_value() &&
+                                 cagra_index.mapping_fd().has_value();
+
   std::optional<raft::host_matrix_view<const T, int64_t>> opt_dataset_view = std::nullopt;
-  if (dataset_is_on_host) {
+  // Regular CAGRA build: pass the dataset directly to HNSW if it's on the host
+  if (dataset_is_on_host && !is_ace_disk_build) {
+    opt_dataset_view.emplace(
+      raft::make_host_matrix_view<const T, int64_t>(dataset, nrow, this->dim_));
+  }
+  // ACE disk build with original id graph: pass the dataset directly to HNSW to remap the graph to
+  // original ids
+  if (is_ace_disk_build && build_param_.use_original_id_graph) {
+    RAFT_EXPECTS(dataset_is_on_host, "Dataset must be on host for original id graph remapping.");
     opt_dataset_view.emplace(
       raft::make_host_matrix_view<const T, int64_t>(dataset, nrow, this->dim_));
   }

cpp/include/cuvs/neighbors/hnsw.hpp

@@ -452,7 +452,8 @@ std::unique_ptr<index<int8_t>> build(
  * @param[in] res raft resources
  * @param[in] params hnsw index parameters
  * @param[in] cagra_index cagra index
- * @param[in] dataset optional dataset to avoid extra memory copy when hierarchy is `CPU`
+ * @param[in] dataset optional dataset in the original row order. When provided for a disk-backed
+ *            ACE index, cuVS remaps the on-disk ACE graph back to original ids before exporting.
  *
  * Usage example:
  * @code{.cpp}
@@ -488,7 +489,8 @@ std::unique_ptr<index<float>> from_cagra(
  * @param[in] res raft resources
  * @param[in] params hnsw index parameters
  * @param[in] cagra_index cagra index
- * @param[in] dataset optional dataset to avoid extra memory copy when hierarchy is `CPU`
+ * @param[in] dataset optional dataset in the original row order. When provided for a disk-backed
+ *            ACE index, cuVS remaps the on-disk ACE graph back to original ids before exporting.
  *
  * Usage example:
  * @code{.cpp}
@@ -524,7 +526,8 @@ std::unique_ptr<index<half>> from_cagra(
  * @param[in] res raft resources
  * @param[in] params hnsw index parameters
  * @param[in] cagra_index cagra index
- * @param[in] dataset optional dataset to avoid extra memory copy when hierarchy is `CPU`
+ * @param[in] dataset optional dataset in the original row order. When provided for a disk-backed
+ *            ACE index, cuVS remaps the on-disk ACE graph back to original ids before exporting.
  *
  * Usage example:
  * @code{.cpp}
@@ -560,7 +563,8 @@ std::unique_ptr<index<uint8_t>> from_cagra(
  * @param[in] res raft resources
  * @param[in] params hnsw index parameters
  * @param[in] cagra_index cagra index
- * @param[in] dataset optional dataset to avoid extra memory copy when hierarchy is `CPU`
+ * @param[in] dataset optional dataset in the original row order. When provided for a disk-backed
+ *            ACE index, cuVS remaps the on-disk ACE graph back to original ids before exporting.
  *
  * Usage example:
  * @code{.cpp}