Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 20 additions & 9 deletions cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
* SPDX-License-Identifier: Apache-2.0
*/

#include "../common/ann_types.hpp"
#include "../common/conf.hpp"
#include "cuvs_ann_bench_param_parser.h"
#include "cuvs_cagra_hnswlib_wrapper.h"

Expand All @@ -27,6 +28,9 @@ auto parse_build_param(const nlohmann::json& conf) ->
hnsw_params.hierarchy = cuvs::neighbors::hnsw::HnswHierarchy::CPU;
} else if (conf.at("hierarchy") == "gpu") {
hnsw_params.hierarchy = cuvs::neighbors::hnsw::HnswHierarchy::GPU;
} else if (conf.at("hierarchy") == "gpu_layered_on_disk" ||
conf.at("hierarchy") == "gpu_layered" || conf.at("hierarchy") == "layered") {
hnsw_params.hierarchy = cuvs::neighbors::hnsw::HnswHierarchy::GPU_LAYERED_ON_DISK;
} else {
THROW("Invalid value for hierarchy: %s", conf.at("hierarchy").get<std::string>().c_str());
}
Expand All @@ -36,6 +40,11 @@ auto parse_build_param(const nlohmann::json& conf) ->
if (conf.contains("ef_construction")) {
hnsw_params.ef_construction = conf.at("ef_construction");
}
if (conf.contains("dataset_path")) {
hnsw_params.dataset_path = conf.at("dataset_path");
} else if (hnsw_params.hierarchy == cuvs::neighbors::hnsw::HnswHierarchy::GPU_LAYERED_ON_DISK) {
hnsw_params.dataset_path = configuration::singleton().get_dataset_conf().base_file;
}
if (conf.contains("num_threads")) { hnsw_params.num_threads = conf.at("num_threads"); }

// Reuse the CAGRA wrapper params parser
Expand All @@ -55,16 +64,18 @@ auto parse_build_param(const nlohmann::json& conf) ->
cuvs::neighbors::cagra::hnsw_heuristic_type::SAME_GRAPH_FOOTPRINT,
dist_type);
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")) {
// Parse ACE parameters if provided.
auto ace_conf = collect_conf_with_prefix(conf, "ace_");
if (!ace_conf.empty()) {
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 (ace_conf.contains("npartitions")) {
ace_params.npartitions = ace_conf.at("npartitions");
}
if (ace_conf.contains("build_dir")) { ace_params.build_dir = ace_conf.at("build_dir"); }
if (ace_conf.contains("ef_construction")) {
ace_params.ef_construction = ace_conf.at("ef_construction");
}
if (conf.contains("use_disk")) { ace_params.use_disk = conf.at("use_disk"); }
if (ace_conf.contains("use_disk")) { ace_params.use_disk = ace_conf.at("use_disk"); }
ps.graph_build_params = ace_params;
}
// NB: above, we only provide the defaults. Below we parse the explicit parameters as usual.
Expand Down
30 changes: 29 additions & 1 deletion cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -9,10 +9,26 @@
#include <raft/core/logger.hpp>

#include <chrono>
#include <filesystem>
#include <memory>

namespace cuvs::bench {

inline void copy_file_overwrite(const std::filesystem::path& src, const std::filesystem::path& dst)
{
std::error_code ec;
if (src == dst ||
(std::filesystem::exists(dst, ec) && std::filesystem::equivalent(src, dst, ec))) {
return;
}
if (!dst.parent_path().empty()) { std::filesystem::create_directories(dst.parent_path()); }

std::filesystem::copy_file(src, dst, std::filesystem::copy_options::overwrite_existing, ec);
const auto src_str = src.string();
const auto dst_str = dst.string();
RAFT_EXPECTS(!ec, "Failed to copy '%s' to '%s'.", src_str.c_str(), dst_str.c_str());
}

template <typename T, typename IdxT>
class cuvs_cagra_hnswlib : public algo<T>, public algo_gpu {
public:
Expand Down Expand Up @@ -130,6 +146,18 @@ void cuvs_cagra_hnswlib<T, IdxT>::set_search_param(const search_param_base& para
template <typename T, typename IdxT>
void cuvs_cagra_hnswlib<T, IdxT>::save(const std::string& file) const
{
if (build_param_.hnsw_index_params.hierarchy ==
cuvs::neighbors::hnsw::HnswHierarchy::GPU_LAYERED_ON_DISK) {
const auto src_artifact = std::filesystem::path(hnsw_index_->file_path());
RAFT_EXPECTS(!src_artifact.empty(), "Layered HNSW artifact path is not available.");
RAFT_EXPECTS(std::filesystem::exists(src_artifact),
"Layered HNSW artifact '%s' does not exist.",
src_artifact.c_str());

copy_file_overwrite(src_artifact, std::filesystem::path(file));
return;
}

if (cagra_ace_build_) {
std::string index_filename = hnsw_index_->file_path();
RAFT_EXPECTS(!index_filename.empty(), "HNSW index file path is not available.");
Expand Down
18 changes: 15 additions & 3 deletions cpp/include/cuvs/neighbors/hnsw.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
#include <cstdint>
#include <cuvs/core/export.hpp>
#include <memory>
#include <string>
#include <type_traits>
#include <variant>

Expand All @@ -41,9 +42,10 @@ namespace graph_build_params = cuvs::neighbors::graph_build_params;
* NOTE: When the value is `NONE`, the HNSW index is built as a base-layer-only index.
*/
enum class HnswHierarchy {
NONE, // base-layer-only index
CPU, // full index with CPU-built hierarchy
GPU // full index with GPU-built hierarchy
NONE, // base-layer-only index
CPU, // full index with CPU-built hierarchy
GPU, // full index with GPU-built hierarchy
GPU_LAYERED_ON_DISK // GPU-built hierarchy stored as layered on-disk topology
};

struct index_params : cuvs::neighbors::index_params {
Expand All @@ -64,6 +66,16 @@ struct index_params : cuvs::neighbors::index_params {
*/
size_t M = 32;

/** Local dataset path used by layered HNSW deserialization.
*
* When `hierarchy == HnswHierarchy::GPU_LAYERED_ON_DISK`, the index artifact stores graph
* topology only. `deserialize` loads vectors from this local dataset path to reconstruct an
* in-memory HNSW index.
* Currently supported local dataset formats are `.npy` and ANN benchmark `*.bin` files with a
* `[uint32 rows, uint32 cols]` header.
*/
std::string dataset_path;

/** Parameters to fine tune GPU graph building. By default we select the parameters based on
* dataset shape and HNSW build parameters. You can override these parameters to fine tune the
* graph building process as described in the CAGRA build docs.
Expand Down
Loading
Loading