Iterative cagra q

.gitignore

@@ -72,7 +72,9 @@ docs/source/_static/rust
 
 # clang tooling
 compile_commands.json
-.clangd/
+
+
+
 
 # serialized ann indexes
 brute_force_index
@@ -86,5 +88,8 @@ ivf_pq_index
 /datasets/
 /*.json
 
+# clangd
+*/.clangd
+
 # java
 .classpath

cpp/CMakeLists.txt

@@ -268,12 +268,12 @@ if(NOT BUILD_CPU_ONLY)
     INPUT_FILE
       "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/detail/cagra/compute_distance_vpq_inst.cu.in"
     OUTPUT_FILE_FORMAT
-      "${CMAKE_CURRENT_BINARY_DIR}/src/neighbors/detail/cagra/compute_distance_vpq_inst_data_@data_abbrev@_index_@index_abbrev@_distance_@distance_abbrev@_codebook_@codebook_abbrev@_metric_@metric@_team_@team_size@_dim_@dim@_pq_bits_@pq_bits@_pq_len_@pq_len@.cu"
+      "${CMAKE_CURRENT_BINARY_DIR}/src/neighbors/detail/cagra/compute_distance_vpq_inst_data_@data_abbrev@_index_@index_abbrev@_distance_@distance_abbrev@_codebook_@codebook_abbrev@_metric_@metric@_team_@team_size@_dim_@dim@_pq_bits_@pq_bits@_pq_len_@pq_len@_smem_@smem_abbrev@.cu"
   )
   generate_string_matrix(
     cagra_compute_distance_vpq_selector_template_params
     ITEM_FORMAT
-    "\nvpq_descriptor_spec<DistanceType::@metric@, @team_size@, @dim@, @pq_bits@, @pq_len@, @codebook_type@, @data_type@, @index_type@, @distance_type@>"
+    "\nvpq_descriptor_spec<DistanceType::@metric@, @team_size@, @dim@, @pq_bits@, @pq_len@, @codebook_type@, @data_type@, @index_type@, @distance_type@, @smem_dtype@>"
     GLUE
     ","
     MATRIX_JSON_FILE
@@ -282,7 +282,7 @@ if(NOT BUILD_CPU_ONLY)
   generate_string_matrix(
     cagra_compute_distance_vpq_template_inst
     ITEM_FORMAT
-    "extern template struct vpq_descriptor_spec<DistanceType::@metric@, @team_size@, @dim@, @pq_bits@, @pq_len@, @codebook_type@, @data_type@, @index_type@, @distance_type@>@semicolon@"
+    "extern template struct vpq_descriptor_spec<DistanceType::@metric@, @team_size@, @dim@, @pq_bits@, @pq_len@, @codebook_type@, @data_type@, @index_type@, @distance_type@, @smem_dtype@>@semicolon@"
     GLUE
     "\n"
     MATRIX_JSON_FILE
@@ -688,13 +688,13 @@ if(NOT BUILD_CPU_ONLY)
   generate_jit_lto_kernels(
     jit_lto_files
     NAME_FORMAT
-      "cagra_setup_workspace@pq_prefix@_team_size_@team_size@_dataset_block_dim_@dataset_block_dim@_@pq_bits@pq_@pq_len@subd_data_@data_abbrev@_query_@query_abbrev@"
+      "cagra_setup_workspace@pq_prefix@_team_size_@team_size@_dataset_block_dim_@dataset_block_dim@_@pq_bits@pq_@pq_len@subd_data_@data_abbrev@_query_@query_abbrev@_smem_@smem_abbrev@"
     MATRIX_JSON_FILE
       "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/detail/cagra/jit_lto_kernels/setup_workspace_matrix.json"
     KERNEL_INPUT_FILE
       "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/detail/cagra/jit_lto_kernels/setup_workspace_kernel.cu.in"
     FRAGMENT_TAG_FORMAT
-      "${cagra_ns}::fragment_tag_setup_workspace<${neighbors_ns}::tag_@data_abbrev@, ${neighbors_ns}::tag_index_@index_abbrev@, ${cagra_ns}::tag_dist_@distance_abbrev@, ${neighbors_ns}::tag_@query_abbrev@, ${cagra_ns}::tag_codebook_@codebook_abbrev@, @team_size@, @dataset_block_dim@, @pq_bits@, @pq_len@>"
+      "${cagra_ns}::fragment_tag_setup_workspace<${neighbors_ns}::tag_@data_abbrev@, ${neighbors_ns}::tag_index_@index_abbrev@, ${cagra_ns}::tag_dist_@distance_abbrev@, ${neighbors_ns}::tag_@query_abbrev@, ${cagra_ns}::tag_codebook_@codebook_abbrev@, @team_size@, @dataset_block_dim@, @pq_bits@, @pq_len@, ${cagra_ns}::tag_smem_@smem_abbrev@>"
     FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/cagra/cagra_fragments.hpp>"
                               "<cuvs/detail/jit_lto/common_fragments.hpp>"
     OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/cagra/setup_workspace"
@@ -704,13 +704,13 @@ if(NOT BUILD_CPU_ONLY)
   generate_jit_lto_kernels(
     jit_lto_files
     NAME_FORMAT
-      "cagra_compute_distance@pq_prefix@_team_size_@team_size@_dataset_block_dim_@dataset_block_dim@_@pq_bits@pq_@pq_len@subd_data_@data_abbrev@_query_@query_abbrev@"
+      "cagra_compute_distance@pq_prefix@_team_size_@team_size@_dataset_block_dim_@dataset_block_dim@_@pq_bits@pq_@pq_len@subd_data_@data_abbrev@_query_@query_abbrev@_smem_@smem_abbrev@"
     MATRIX_JSON_FILE
       "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/detail/cagra/jit_lto_kernels/compute_distance_matrix.json"
     KERNEL_INPUT_FILE
       "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/detail/cagra/jit_lto_kernels/compute_distance_kernel.cu.in"
     FRAGMENT_TAG_FORMAT
-      "${cagra_ns}::fragment_tag_compute_distance<${neighbors_ns}::tag_@data_abbrev@, ${neighbors_ns}::tag_index_@index_abbrev@, ${cagra_ns}::tag_dist_@distance_abbrev@, ${neighbors_ns}::tag_@query_abbrev@, ${cagra_ns}::tag_codebook_@codebook_abbrev@, @team_size@, @dataset_block_dim@, @pq_bits@, @pq_len@>"
+      "${cagra_ns}::fragment_tag_compute_distance<${neighbors_ns}::tag_@data_abbrev@, ${neighbors_ns}::tag_index_@index_abbrev@, ${cagra_ns}::tag_dist_@distance_abbrev@, ${neighbors_ns}::tag_@query_abbrev@, ${cagra_ns}::tag_codebook_@codebook_abbrev@, @team_size@, @dataset_block_dim@, @pq_bits@, @pq_len@, ${cagra_ns}::tag_smem_@smem_abbrev@>"
     FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/cagra/cagra_fragments.hpp>"
                               "<cuvs/detail/jit_lto/common_fragments.hpp>"
     OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/cagra/compute_distance"

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

@@ -320,10 +320,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).
@@ -354,6 +356,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);

cpp/include/cuvs/detail/jit_lto/cagra/cagra_fragments.hpp

@@ -16,6 +16,8 @@ struct tag_metric_cosine {};
 struct tag_metric_hamming {};
 struct tag_codebook_none {};
 struct tag_codebook_half {};
+struct tag_smem_f16 {};
+struct tag_smem_e5m2 {};
 struct tag_metric_l1 {};
 struct tag_norm_noop {};
 struct tag_norm_cosine {};
@@ -33,7 +35,8 @@ template <typename DataTag,
           uint32_t TeamSize,
           uint32_t DatasetBlockDim,
           uint32_t PqBits,
-          uint32_t PqLen>
+          uint32_t PqLen,
+          typename SmemTag>
 struct fragment_tag_setup_workspace {};
 
 template <typename DataTag,
@@ -44,7 +47,8 @@ template <typename DataTag,
           uint32_t TeamSize,
           uint32_t DatasetBlockDim,
           uint32_t PqBits,
-          uint32_t PqLen>
+          uint32_t PqLen,
+          typename SmemTag>
 struct fragment_tag_compute_distance {};
 
 template <typename QueryTag, typename DistanceTag, typename MetricTag>