Iterative CAGRA-Q

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

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -367,10 +367,12 @@ void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::cagra::index
   }
 
   // Parse build-algo-specific parameters and use them to decide on the algo type
-  nlohmann::json ivf_pq_build_conf  = collect_conf_with_prefix(conf, "ivf_pq_build_");
-  nlohmann::json ivf_pq_search_conf = collect_conf_with_prefix(conf, "ivf_pq_search_");
-  nlohmann::json nn_descent_conf    = collect_conf_with_prefix(conf, "nn_descent_");
-  nlohmann::json ace_conf           = collect_conf_with_prefix(conf, "ace_");
+  nlohmann::json ivf_pq_build_conf      = collect_conf_with_prefix(conf, "ivf_pq_build_");
+  nlohmann::json ivf_pq_search_conf     = collect_conf_with_prefix(conf, "ivf_pq_search_");
+  nlohmann::json nn_descent_conf        = collect_conf_with_prefix(conf, "nn_descent_");
+  nlohmann::json ace_conf               = collect_conf_with_prefix(conf, "ace_");
+  nlohmann::json build_compression_conf = collect_conf_with_prefix(conf, "build_compression_");
+  nlohmann::json build_search_conf      = collect_conf_with_prefix(conf, "build_search_");
 
   // When graph_build_algo is not specified, leave graph_build_params as monostate so the
   // CAGRA build uses AUTO selection (NN_DESCENT or IVF_PQ based on dataset/heuristics).
@@ -401,6 +403,94 @@ void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::cagra::index
       } else if constexpr (std::is_same_v<U,
                                           cuvs::neighbors::graph_build_params::nn_descent_params>) {
         parse_build_param<T, IdxT>(nn_descent_conf, arg);
+      } else if constexpr (std::is_same_v<
+                             U,
+                             cuvs::neighbors::graph_build_params::iterative_search_params>) {
+        if (!build_compression_conf.empty()) {
+          auto vpq_pams = arg.build_compression.value_or(cuvs::neighbors::vpq_params{});
+          parse_build_param(build_compression_conf, vpq_pams);
+          arg.build_compression.emplace(vpq_pams);
+        }
+        if (build_search_conf.contains("width")) {
+          arg.search_width = build_search_conf.at("width");
+        }
+        if (build_search_conf.contains("max_iterations")) {
+          arg.max_iterations = build_search_conf.at("max_iterations");
+        }
+        if (build_search_conf.contains("min_iterations")) {
+          arg.min_iterations = build_search_conf.at("min_iterations");
+        }
+        if (build_search_conf.contains("itopk")) { arg.itopk_size = build_search_conf.at("itopk"); }
+        if (build_search_conf.contains("max_queries")) {
+          arg.max_queries = build_search_conf.at("max_queries");
+        }
+        if (build_search_conf.contains("team_size")) {
+          arg.team_size = build_search_conf.at("team_size");
+        }
+        if (build_search_conf.contains("thread_block_size")) {
+          arg.thread_block_size = build_search_conf.at("thread_block_size");
+        }
+        if (build_search_conf.contains("hashmap_min_bitlen")) {
+          arg.hashmap_min_bitlen = build_search_conf.at("hashmap_min_bitlen");
+        }
+        if (build_search_conf.contains("hashmap_max_fill_rate")) {
+          arg.hashmap_max_fill_rate = build_search_conf.at("hashmap_max_fill_rate");
+        }
+        if (build_search_conf.contains("num_random_samplings")) {
+          arg.num_random_samplings = build_search_conf.at("num_random_samplings");
+        }
+        if (build_search_conf.contains("persistent")) {
+          arg.persistent = build_search_conf.at("persistent");
+        }
+        if (build_search_conf.contains("persistent_lifetime")) {
+          arg.persistent_lifetime = build_search_conf.at("persistent_lifetime");
+        }
+        if (build_search_conf.contains("persistent_device_usage")) {
+          arg.persistent_device_usage = build_search_conf.at("persistent_device_usage");
+        }
+        if (build_search_conf.contains("algo")) {
+          std::string algo = build_search_conf.at("algo");
+          if (algo == "single_cta") {
+            arg.algo = cuvs::neighbors::cagra::search_algo::SINGLE_CTA;
+          } else if (algo == "multi_cta") {
+            arg.algo = cuvs::neighbors::cagra::search_algo::MULTI_CTA;
+          } else if (algo == "multi_kernel") {
+            arg.algo = cuvs::neighbors::cagra::search_algo::MULTI_KERNEL;
+          } else if (algo == "auto") {
+            arg.algo = cuvs::neighbors::cagra::search_algo::AUTO;
+          }
+        }
+        if (build_search_conf.contains("hashmap_mode")) {
+          std::string mode = build_search_conf.at("hashmap_mode");
+          if (mode == "hash") {
+            arg.hashmap_mode = cuvs::neighbors::cagra::hash_mode::HASH;
+          } else if (mode == "small") {
+            arg.hashmap_mode = cuvs::neighbors::cagra::hash_mode::SMALL;
+          } else if (mode == "auto") {
+            arg.hashmap_mode = cuvs::neighbors::cagra::hash_mode::AUTO;
+          }
+        }
+        // Whether to shuffle the (compressed) dataset before the iterative build loop.
+        if (build_search_conf.contains("shuffle_dataset")) {
+          arg.shuffle_dataset = build_search_conf.at("shuffle_dataset").get<bool>();
+        }
+        // Precision of the codebook/query in shared memory for the VPQ search used during
+        // the iterative build. Accepts an integer code (0=F16, 1=E5M2) or a string.
+        if (build_search_conf.contains("smem_dtype")) {
+          const auto& sd = build_search_conf.at("smem_dtype");
+          if (sd.is_number_integer()) {
+            arg.smem_dtype = static_cast<cuvs::neighbors::cagra::internal_dtype>(sd.get<int>());
+          } else {
+            std::string s = sd.get<std::string>();
+            if (s == "f16" || s == "F16" || s == "fp16" || s == "half") {
+              arg.smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16;
+            } else if (s == "e5m2" || s == "E5M2" || s == "fp8") {
+              arg.smem_dtype = cuvs::neighbors::cagra::internal_dtype::E5M2;
+            } else {
+              throw std::runtime_error("invalid value for build_search smem_dtype: " + s);
+            }
+          }
+        }
       }
     },
     params.graph_build_params);